├── lanedet ├── engine │ ├── __init__.py │ ├── optimizer.py │ ├── scheduler.py │ ├── registry.py │ └── runner.py ├── __init__.py ├── models │ ├── necks │ │ ├── __init__.py │ │ └── fpn.py │ ├── nets │ │ ├── __init__.py │ │ └── detector.py │ ├── backbones │ │ ├── __init__.py │ │ ├── vgg.py │ │ ├── erfnet.py │ │ └── mobilenet.py │ ├── __init__.py │ ├── aggregators │ │ ├── __init__.py │ │ ├── scnn.py │ │ ├── aspp.py │ │ ├── resa.py │ │ └── transformer.py │ ├── heads │ │ ├── __init__.py │ │ ├── plain_decoder.py │ │ ├── exist_head.py │ │ ├── busd.py │ │ ├── lane_cls.py │ │ └── lane_seg.py │ ├── registry.py │ └── losses │ │ └── focal_loss.py ├── ops │ ├── __init__.py │ ├── nms.py │ └── csrc │ │ ├── nms.cpp │ │ └── nms_kernel.cu ├── utils │ ├── __init__.py │ ├── visualization_vid.py │ ├── logger.py │ ├── net_utils.py │ ├── visualization.py │ ├── registry.py │ ├── tusimple_metric.py │ ├── recorder.py │ └── culane_metric.py ├── datasets │ ├── __init__.py │ ├── process │ │ ├── __init__.py │ │ ├── process.py │ │ ├── generate_lane_cls.py │ │ ├── generate_lane_line.py │ │ └── alaug.py │ ├── registry.py │ ├── base_dataset.py │ ├── culane.py │ └── tusimple.py └── core │ └── lane.py ├── .cache ├── culane_anchors_freq.pt └── tusimple_anchors_freq.pt ├── .github ├── test-class-lvlane-ufld2.jpg └── _clips_0601_1494452613491980502_20.jpg ├── requirements.txt ├── .gitignore ├── docker ├── README.md └── Dockerfile ├── configs ├── condlane │ ├── README.md │ ├── resnet50_culane.py │ └── resnet101_culane.py ├── ufld │ ├── README.md │ ├── resnet18_culane.py │ └── resnet18_tusimple.py ├── scnn │ ├── README.md │ ├── resnet18_tusimple.py │ └── resnet50_culane.py └── resa │ ├── README.md │ ├── resa34_tusimple.py │ ├── resa34_culane.py │ ├── resa50_culane.py │ └── resa18_tusimple.py ├── main.py ├── tools ├── detect.py └── generate_seg_tusimple.py ├── setup.py └── README.md /lanedet/engine/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /lanedet/__init__.py: -------------------------------------------------------------------------------- 1 | from .ops import * 2 | -------------------------------------------------------------------------------- /lanedet/models/necks/__init__.py: -------------------------------------------------------------------------------- 1 | from .fpn import FPN 2 | -------------------------------------------------------------------------------- /lanedet/models/nets/__init__.py: -------------------------------------------------------------------------------- 1 | from .detector import Detector 2 | -------------------------------------------------------------------------------- /lanedet/ops/__init__.py: -------------------------------------------------------------------------------- 1 | from .nms import nms 2 | 3 | __all__ = ['nms'] 4 | -------------------------------------------------------------------------------- /lanedet/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .config import Config 2 | from .registry import Registry, build_from_cfg 3 | -------------------------------------------------------------------------------- /.cache/culane_anchors_freq.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zillur01/LVLane/HEAD/.cache/culane_anchors_freq.pt -------------------------------------------------------------------------------- /.cache/tusimple_anchors_freq.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zillur01/LVLane/HEAD/.cache/tusimple_anchors_freq.pt -------------------------------------------------------------------------------- /.github/test-class-lvlane-ufld2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zillur01/LVLane/HEAD/.github/test-class-lvlane-ufld2.jpg -------------------------------------------------------------------------------- /.github/_clips_0601_1494452613491980502_20.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zillur01/LVLane/HEAD/.github/_clips_0601_1494452613491980502_20.jpg -------------------------------------------------------------------------------- /lanedet/models/backbones/__init__.py: -------------------------------------------------------------------------------- 1 | from .resnet import ResNet 2 | from .vgg import VGG 3 | from .erfnet import ERFNet 4 | from .mobilenet import MobileNet 5 | -------------------------------------------------------------------------------- /lanedet/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .backbones import * 2 | from .aggregators import * 3 | from .heads import * 4 | from .nets import * 5 | from .necks import * 6 | 7 | -------------------------------------------------------------------------------- /lanedet/models/aggregators/__init__.py: -------------------------------------------------------------------------------- 1 | from .scnn import SCNN 2 | from .resa import RESA 3 | from .aspp import ASPP 4 | from .transformer import TransConvEncoderModule 5 | -------------------------------------------------------------------------------- /lanedet/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .registry import build_dataset, build_dataloader 2 | 3 | from .tusimple import TuSimple 4 | from .culane import CULane 5 | 6 | from .process import * 7 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | pandas 2 | addict 3 | scikit-learn 4 | opencv-python 5 | pytorch_warmup 6 | scikit-image 7 | tqdm 8 | p_tqdm 9 | imgaug>=0.4.0 10 | Shapely==1.7.0 11 | ujson==1.35 12 | yapf 13 | albumentations==0.4.6 14 | mmcv==1.2.5 15 | pathspec 16 | -------------------------------------------------------------------------------- /lanedet/models/heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .exist_head import ExistHead 2 | from .lane_cls import LaneCls 3 | from .busd import BUSD 4 | from .plain_decoder import PlainDecoder 5 | from .laneatt import LaneATT 6 | from .lane_seg import LaneSeg 7 | from .condlane import CondLaneHead 8 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | work_dirs/ 2 | predicts/ 3 | output/ 4 | data/ 5 | 6 | __pycache__/ 7 | */*.un~ 8 | .*.swp 9 | 10 | 11 | 12 | *.egg-info/ 13 | *.egg 14 | *.eggs 15 | 16 | output.txt 17 | .vscode/* 18 | .DS_Store 19 | tmp.* 20 | *.pt 21 | *.pth 22 | *.un~ 23 | *.so 24 | build 25 | 26 | lane 27 | debug 28 | pretrained_models 29 | -------------------------------------------------------------------------------- /lanedet/engine/optimizer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def build_optimizer(cfg, net): 5 | params = [] 6 | cfg_cp = cfg.optimizer.copy() 7 | cfg_type = cfg_cp.pop('type') 8 | 9 | if cfg_type not in dir(torch.optim): 10 | raise ValueError("{} is not defined.".format(cfg_type)) 11 | 12 | _optim = getattr(torch.optim, cfg_type) 13 | return _optim(net.parameters(), **cfg_cp) 14 | -------------------------------------------------------------------------------- /lanedet/engine/scheduler.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | 4 | def build_scheduler(cfg, optimizer): 5 | 6 | cfg_cp = cfg.scheduler.copy() 7 | cfg_type = cfg_cp.pop('type') 8 | 9 | if cfg_type not in dir(torch.optim.lr_scheduler): 10 | raise ValueError("{} is not defined.".format(cfg_type)) 11 | 12 | 13 | _scheduler = getattr(torch.optim.lr_scheduler, cfg_type) 14 | 15 | 16 | return _scheduler(optimizer, **cfg_cp) 17 | -------------------------------------------------------------------------------- /docker/README.md: -------------------------------------------------------------------------------- 1 | # Building a LaneDet container 2 | 3 | ### Install Docker and NVIDIA Container Toolkit 4 | 5 | https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html 6 | 7 | ### Build Container 8 | 9 | > 10 | 11 | ```Shell 12 | # From root of LaneDet repo 13 | cd $LANEDET_ROOT 14 | 15 | # Build: 16 | docker build -f docker/Dockerfile -t lanedet:latest . 17 | 18 | # Run: 19 | docker run --gpus all -it \ 20 | --shm-size=8gb \ 21 | --name=lanedet --ipc=host --net=host lanedet:latest 22 | ``` -------------------------------------------------------------------------------- /lanedet/utils/visualization_vid.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import os 3 | import os.path as osp 4 | 5 | def imshow_lanes(img, lanes, show=False, out_file=None): 6 | for lane in lanes: 7 | for x, y in lane: 8 | if x <= 0 or y <= 0: 9 | continue 10 | x, y = int(x), int(y) 11 | cv2.circle(img, (x, y), 4, (255, 0, 0), 2) 12 | 13 | if out_file: 14 | if not osp.exists(osp.dirname(out_file)): 15 | os.makedirs(osp.dirname(out_file)) 16 | cv2.imwrite(out_file, img) 17 | return img 18 | 19 | -------------------------------------------------------------------------------- /lanedet/utils/logger.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | def init_logger(log_file=None, log_level=logging.INFO): 4 | stream_handler = logging.StreamHandler() 5 | handlers = [stream_handler] 6 | 7 | if log_file is not None: 8 | file_handler = logging.FileHandler(log_file, 'w') 9 | handlers.append(file_handler) 10 | 11 | formatter = logging.Formatter( 12 | '%(asctime)s - %(name)s - %(levelname)s - %(message)s') 13 | for handler in handlers: 14 | handler.setFormatter(formatter) 15 | handler.setLevel(log_level) 16 | 17 | logging.basicConfig(level=log_level, handlers=handlers) 18 | -------------------------------------------------------------------------------- /lanedet/engine/registry.py: -------------------------------------------------------------------------------- 1 | from lanedet.utils import Registry, build_from_cfg 2 | 3 | TRAINER = Registry('trainer') 4 | EVALUATOR = Registry('evaluator') 5 | 6 | def build(cfg, registry, default_args=None): 7 | if isinstance(cfg, list): 8 | modules = [ 9 | build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg 10 | ] 11 | return nn.Sequential(*modules) 12 | else: 13 | return build_from_cfg(cfg, registry, default_args) 14 | 15 | def build_trainer(cfg): 16 | return build(cfg.trainer, TRAINER, default_args=dict(cfg=cfg)) 17 | 18 | def build_evaluator(cfg): 19 | return build(cfg.evaluator, EVALUATOR, default_args=dict(cfg=cfg)) 20 | -------------------------------------------------------------------------------- /lanedet/datasets/process/__init__.py: -------------------------------------------------------------------------------- 1 | from .transforms import (RandomLROffsetLABEL, RandomUDoffsetLABEL, 2 | Resize, RandomCrop, CenterCrop, RandomRotation, RandomBlur, 3 | RandomHorizontalFlip, Normalize, ToTensor) 4 | 5 | from .generate_lane_cls import GenerateLaneCls 6 | from .generate_lane_line import GenerateLaneLine 7 | from .collect_lane import CollectLane 8 | from .process import Process 9 | from .alaug import Alaug 10 | 11 | __all__ = ['Process', 'RandomLROffsetLABEL', 'RandomUDoffsetLABEL', 12 | 'Resize', 'RandomCrop', 'CenterCrop', 'RandomRotation', 'RandomBlur', 13 | 'RandomHorizontalFlip', 'Normalize', 'GenerateLaneCls', 14 | 'ToTensor', 'GenerateLaneLine'] 15 | -------------------------------------------------------------------------------- /configs/condlane/README.md: -------------------------------------------------------------------------------- 1 | # CondLaneNet: a Top-to-down Lane Detection Framework Based on Conditional Convolution 2 | 3 | ## Introduction 4 | 5 | ```latex 6 | @article{liu2021condlanenet, 7 | title={CondLaneNet: a Top-to-down Lane Detection Framework Based on Conditional Convolution}, 8 | author={Liu, Lizhe and Chen, Xiaohao and Zhu, Siyu and Tan, Ping}, 9 | journal={arXiv preprint arXiv:2105.05003}, 10 | year={2021} 11 | } 12 | ``` 13 | 14 | ## Models 15 | | Architecture| Backbone |Dataset | Metric | Config| Checkpoints | 16 | |-------------|----------|--------|--------|-------|--------------| 17 | | CondLane | ResNet101 | CULane | F1: 79.47| [config](configs/condlane/resnet101_culane.py) |[model](https://github.com/Turoad/lanedet/releases/download/1.0/condlane_r101_culane.pth.zip) | 18 | -------------------------------------------------------------------------------- /docker/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM nvcr.io/nvidia/pytorch:21.10-py3 2 | 3 | ENV DEBIAN_FRONTEND noninteractive 4 | 5 | RUN rm -rf /var/lib/apt/lists/* && rm -rf /etc/apt/sources.list.d/*\ 6 | /etc/apt/sources.list.d/cuda.list \ 7 | /etc/apt/sources.list.d/nvidia-ml.list && \ 8 | sed -i s@/archive.ubuntu.com/@/mirrors.ustc.edu.cn/@g /etc/apt/sources.list && \ 9 | apt-get update && apt-get install -y ffmpeg libsm6 libxext6 git ninja-build libglib2.0-0 libsm6 libxrender-dev libxext6 \ 10 | && apt-get clean \ 11 | && rm -rf /var/lib/apt/lists/* 12 | 13 | # Install LaneDet 14 | RUN conda clean --all 15 | RUN git clone https://github.com/turoad/lanedet.git /lanedet 16 | WORKDIR /lanedet 17 | RUN pip install --no-cache-dir -r requirements.txt -i https://mirrors.aliyun.com/pypi/simple/ && \ 18 | pip install --no-cache-dir -e . 19 | -------------------------------------------------------------------------------- /lanedet/models/heads/plain_decoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | from torch.hub import load_state_dict_from_url 5 | 6 | from ..registry import HEADS 7 | 8 | @HEADS.register_module 9 | class PlainDecoder(nn.Module): 10 | def __init__(self, cfg): 11 | super(PlainDecoder, self).__init__() 12 | self.cfg = cfg 13 | 14 | self.dropout = nn.Dropout2d(0.1) 15 | self.conv8 = nn.Conv2d(cfg.featuremap_out_channel, cfg.num_classes, 1) 16 | 17 | def forward(self, x): 18 | 19 | x = self.dropout(x) 20 | x = self.conv8(x) 21 | x = F.interpolate(x, size=[self.cfg.img_height, self.cfg.img_width], 22 | mode='bilinear', align_corners=False) 23 | 24 | output = {'seg': x} 25 | 26 | return output 27 | -------------------------------------------------------------------------------- /configs/ufld/README.md: -------------------------------------------------------------------------------- 1 | # Ultra Fast Structure-aware Deep Lane Detection 2 | 3 | ## Introduction 4 | 5 | ```latex 6 | @InProceedings{qin2020ultra, 7 | author = {Qin, Zequn and Wang, Huanyu and Li, Xi}, 8 | title = {Ultra Fast Structure-aware Deep Lane Detection}, 9 | booktitle = {The European Conference on Computer Vision (ECCV)}, 10 | year = {2020} 11 | } 12 | ``` 13 | 14 | ## Models 15 | | Architecture| Backbone |Dataset | Metric | Config| Checkpoints | 16 | |-------------|----------|--------|--------|-------|--------------| 17 | | UFLD | ResNet18 | CULane |F1: 69.47| [config](https://github.com/Turoad/lanedet/blob/main/configs/ufld/resa18_culane.py) | [model](https://github.com/Turoad/lanedet/releases/download/1.0/ufld_r18_culane.zip)| 18 | | UFLD | ResNet18 | Tusimple |acc: 95.86| [config](https://github.com/Turoad/lanedet/blob/main/configs/ufld/resa18_culane.py) | [model](https://github.com/Turoad/lanedet/releases/download/1.0/ufld_r18_tusimple.zip)| 19 | -------------------------------------------------------------------------------- /configs/scnn/README.md: -------------------------------------------------------------------------------- 1 | # Spatial As Deep: Spatial CNN for Traffic Scene Understanding 2 | 3 | ## Introduction 4 | 5 | ```latex 6 | @inproceedings{pan2018SCNN, 7 | author = {Xingang Pan, Jianping Shi, Ping Luo, Xiaogang Wang, and Xiaoou Tang}, 8 | title = {Spatial As Deep: Spatial CNN for Traffic Scene Understanding}, 9 | booktitle = {AAAI Conference on Artificial Intelligence (AAAI)}, 10 | month = {February}, 11 | year = {2018} 12 | } 13 | ``` 14 | 15 | ## Models 16 | | Architecture| Backbone |Dataset | Metric | Config| Checkpoints | 17 | |-------------|----------|--------|--------|-------|--------------| 18 | | SCNN | ResNet50 | CULane | F1: 74.89| [config](https://github.com/Turoad/lanedet/blob/main/configs/scnn/resa18_culane.py) |[model](https://github.com/Turoad/lanedet/releases/download/1.0/scnn_r50_culane.zip) | 19 | | SCNN | ResNet18 | Tusimple |acc: 96.05| [config](https://github.com/Turoad/lanedet/blob/main/configs/scnn/resa18_culane.py) | [model](https://github.com/Turoad/lanedet/releases/download/1.0/scnn_r18_tusimple.zip)| 20 | -------------------------------------------------------------------------------- /lanedet/models/heads/exist_head.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | from torch.hub import load_state_dict_from_url 5 | 6 | from ..registry import HEADS 7 | 8 | @HEADS.register_module 9 | class ExistHead(nn.Module): 10 | def __init__(self, cfg=None): 11 | super(ExistHead, self).__init__() 12 | self.cfg = cfg 13 | 14 | self.dropout = nn.Dropout2d(0.1) 15 | self.conv8 = nn.Conv2d(cfg.featuremap_out_channel, cfg.num_classes, 1) 16 | 17 | stride = cfg.featuremap_out_stride * 2 18 | self.fc9 = nn.Linear( 19 | int(cfg.num_classes * cfg.img_width / stride * cfg.img_height / stride), 128) 20 | self.fc10 = nn.Linear(128, cfg.num_classes-1) 21 | 22 | def forward(self, x): 23 | x = self.dropout(x) 24 | x = self.conv8(x) 25 | 26 | x = F.softmax(x, dim=1) 27 | x = F.avg_pool2d(x, 2, stride=2, padding=0) 28 | x = x.view(-1, x.numel() // x.shape[0]) 29 | x = self.fc9(x) 30 | x = F.relu(x) 31 | x = self.fc10(x) 32 | x = torch.sigmoid(x) 33 | 34 | output = {'exist': x} 35 | 36 | return output 37 | -------------------------------------------------------------------------------- /lanedet/models/registry.py: -------------------------------------------------------------------------------- 1 | from lanedet.utils import Registry, build_from_cfg 2 | import torch.nn as nn 3 | 4 | BACKBONES = Registry('backbones') 5 | AGGREGATORS = Registry('aggregators') 6 | HEADS = Registry('heads') 7 | NECKS = Registry('necks') 8 | NETS = Registry('nets') 9 | 10 | def build(cfg, registry, default_args=None): 11 | if isinstance(cfg, list): 12 | modules = [ 13 | build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg] 14 | return nn.Sequential(*modules) 15 | else: 16 | return build_from_cfg(cfg, registry, default_args) 17 | 18 | 19 | def build_backbones(cfg): 20 | return build(cfg.backbone, BACKBONES, default_args=dict(cfg=cfg)) 21 | 22 | def build_aggregator(cfg): 23 | return build(cfg.aggregator, AGGREGATORS, default_args=dict(cfg=cfg)) 24 | 25 | def build_heads(cfg): 26 | return build(cfg.heads, HEADS, default_args=dict(cfg=cfg)) 27 | 28 | def build_head(split_cfg, cfg): 29 | return build(split_cfg, HEADS, default_args=dict(cfg=cfg)) 30 | 31 | def build_net(cfg): 32 | return build(cfg.net, NETS, default_args=dict(cfg=cfg)) 33 | 34 | def build_necks(cfg): 35 | return build(cfg.neck, NECKS, default_args=dict(cfg=cfg)) 36 | -------------------------------------------------------------------------------- /lanedet/models/nets/detector.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch 3 | 4 | from lanedet.models.registry import NETS 5 | from ..registry import build_backbones, build_aggregator, build_heads, build_necks 6 | 7 | 8 | @NETS.register_module 9 | class Detector(nn.Module): 10 | def __init__(self, cfg): 11 | super(Detector, self).__init__() 12 | self.cfg = cfg 13 | self.backbone = build_backbones(cfg) 14 | self.aggregator = build_aggregator(cfg) if cfg.haskey('aggregator') else None 15 | self.neck = build_necks(cfg) if cfg.haskey('neck') else None 16 | self.heads = build_heads(cfg) 17 | 18 | def get_lanes(self, output): 19 | return self.heads.get_lanes(output) 20 | 21 | def forward(self, batch): 22 | output = {} 23 | fea = self.backbone(batch['img']) 24 | 25 | if self.aggregator: 26 | fea[-1] = self.aggregator(fea[-1]) 27 | 28 | if self.neck: 29 | fea = self.neck(fea) 30 | 31 | if self.training: 32 | out = self.heads(fea, batch=batch) 33 | output.update(self.heads.loss(out, batch)) 34 | else: 35 | output = self.heads(fea) 36 | 37 | return output 38 | -------------------------------------------------------------------------------- /lanedet/models/aggregators/scnn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | from torch.hub import load_state_dict_from_url 5 | 6 | from lanedet.models.registry import AGGREGATORS 7 | 8 | 9 | @AGGREGATORS.register_module 10 | class SCNN(nn.Module): 11 | def __init__(self, cfg=None): 12 | super(SCNN, self).__init__() 13 | self.conv_d = nn.Conv2d(128, 128, (1, 9), padding=(0, 4), bias=False) 14 | self.conv_u = nn.Conv2d(128, 128, (1, 9), padding=(0, 4), bias=False) 15 | self.conv_r = nn.Conv2d(128, 128, (9, 1), padding=(4, 0), bias=False) 16 | self.conv_l = nn.Conv2d(128, 128, (9, 1), padding=(4, 0), bias=False) 17 | 18 | def forward(self, x): 19 | x = x.clone() 20 | for i in range(1, x.shape[2]): 21 | x[..., i:i+1, :].add_(F.relu(self.conv_d(x[..., i-1:i, :]))) 22 | 23 | for i in range(x.shape[2] - 2, 0, -1): 24 | x[..., i:i+1, :].add_(F.relu(self.conv_u(x[..., i+1:i+2, :]))) 25 | 26 | for i in range(1, x.shape[3]): 27 | x[..., i:i+1].add_(F.relu(self.conv_r(x[..., i-1:i]))) 28 | 29 | for i in range(x.shape[3] - 2, 0, -1): 30 | x[..., i:i+1].add_(F.relu(self.conv_l(x[..., i+1:i+2]))) 31 | return x 32 | -------------------------------------------------------------------------------- /configs/resa/README.md: -------------------------------------------------------------------------------- 1 | # RESA: Recurrent Feature-Shift Aggregator for Lane Detection 2 | 3 | ## Introduction 4 | 5 | ```latex 6 | @misc{zheng2020resa, 7 | title={RESA: Recurrent Feature-Shift Aggregator for Lane Detection}, 8 | author={Tu Zheng and Hao Fang and Yi Zhang and Wenjian Tang and Zheng Yang and Haifeng Liu and Deng Cai}, 9 | year={2020}, 10 | eprint={2008.13719}, 11 | archivePrefix={arXiv}, 12 | primaryClass={cs.CV} 13 | } 14 | ``` 15 | 16 | ## Models 17 | 18 | | Architecture| Backbone |Dataset | Metric | Config| Checkpoints | 19 | |-------------|----------|--------|--------|-------|--------------| 20 | | RESA | ResNet50 | CULane |F1: 75.92| [config](https://github.com/Turoad/lanedet/blob/main/configs/resa/resa50_culane.py) | [model](https://github.com/Turoad/lanedet/releases/download/1.0/resa_r50_culane.zip)| 21 | | RESA | ResNet34 | CULane |F1: 75.85| [config](https://github.com/Turoad/lanedet/blob/main/configs/resa/resa34_culane.py) |[model](https://github.com/Turoad/lanedet/releases/download/1.0/resa_r34_culane.zip)| 22 | | RESA | ResNet34 | Tusimple |acc: 96.86| [config](https://github.com/Turoad/lanedet/blob/main/configs/resa/resa34_tusimple.py) |[model](https://github.com/Turoad/lanedet/releases/download/1.0/resa_r34_tusimple.zip)| 23 | | RESA | ResNet18 | Tusimple |acc: 96.73| [config](https://github.com/Turoad/lanedet/blob/main/configs/resa/resa18_tusimple.py) |[model](https://github.com/Turoad/lanedet/releases/download/1.0/resa_r18_tusimple.zip)| 24 | -------------------------------------------------------------------------------- /lanedet/datasets/registry.py: -------------------------------------------------------------------------------- 1 | from lanedet.utils import Registry, build_from_cfg 2 | 3 | import torch 4 | from functools import partial 5 | import numpy as np 6 | import random 7 | from mmcv.parallel import collate 8 | 9 | DATASETS = Registry('datasets') 10 | PROCESS = Registry('process') 11 | 12 | def build(cfg, registry, default_args=None): 13 | if isinstance(cfg, list): 14 | modules = [ 15 | build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg 16 | ] 17 | return nn.Sequential(*modules) 18 | else: 19 | return build_from_cfg(cfg, registry, default_args) 20 | 21 | 22 | def build_dataset(split_cfg, cfg): 23 | return build(split_cfg, DATASETS, default_args=dict(cfg=cfg)) 24 | 25 | def worker_init_fn(worker_id, seed): 26 | worker_seed = worker_id + seed 27 | np.random.seed(worker_seed) 28 | random.seed(worker_seed) 29 | 30 | def build_dataloader(split_cfg, cfg, is_train=True): 31 | if is_train: 32 | shuffle = True 33 | else: 34 | shuffle = False 35 | 36 | dataset = build_dataset(split_cfg, cfg) 37 | 38 | init_fn = partial( 39 | worker_init_fn, seed=cfg.seed) 40 | 41 | samples_per_gpu = cfg.batch_size // cfg.gpus 42 | data_loader = torch.utils.data.DataLoader( 43 | dataset, batch_size = cfg.batch_size, shuffle = shuffle, 44 | num_workers = cfg.workers, pin_memory = False, drop_last = False, 45 | collate_fn=partial(collate, samples_per_gpu=samples_per_gpu), 46 | worker_init_fn=init_fn) 47 | 48 | return data_loader 49 | -------------------------------------------------------------------------------- /lanedet/datasets/process/process.py: -------------------------------------------------------------------------------- 1 | import collections 2 | 3 | from lanedet.utils import build_from_cfg 4 | 5 | from ..registry import PROCESS 6 | 7 | class Process(object): 8 | """Compose multiple process sequentially. 9 | Args: 10 | process (Sequence[dict | callable]): Sequence of process object or 11 | config dict to be composed. 12 | """ 13 | 14 | def __init__(self, processes, cfg): 15 | assert isinstance(processes, collections.abc.Sequence) 16 | self.processes = [] 17 | 18 | for process in processes: 19 | if isinstance(process, dict): 20 | process = build_from_cfg(process, PROCESS, default_args=dict(cfg=cfg)) 21 | self.processes.append(process) 22 | elif callable(process): 23 | self.processes.append(process) 24 | else: 25 | raise TypeError('process must be callable or a dict') 26 | 27 | def __call__(self, data): 28 | """Call function to apply processes sequentially. 29 | Args: 30 | data (dict): A result dict contains the data to process. 31 | Returns: 32 | dict: Processed data. 33 | """ 34 | 35 | for t in self.processes: 36 | data = t(data) 37 | if data is None: 38 | return None 39 | return data 40 | 41 | def __repr__(self): 42 | format_string = self.__class__.__name__ + '(' 43 | for t in self.processes: 44 | format_string += '\n' 45 | format_string += f' {t}' 46 | format_string += '\n)' 47 | return format_string 48 | -------------------------------------------------------------------------------- /lanedet/core/lane.py: -------------------------------------------------------------------------------- 1 | from scipy.interpolate import InterpolatedUnivariateSpline 2 | import numpy as np 3 | 4 | 5 | class Lane: 6 | def __init__(self, points=None, invalid_value=-2., metadata=None): 7 | super(Lane, self).__init__() 8 | self.curr_iter = 0 9 | self.points = points 10 | self.invalid_value = invalid_value 11 | self.function = InterpolatedUnivariateSpline(points[:, 1], points[:, 0], k=min(3, len(points) - 1)) 12 | self.min_y = points[:, 1].min() - 0.01 13 | self.max_y = points[:, 1].max() + 0.01 14 | 15 | self.metadata = metadata or {} 16 | 17 | def __repr__(self): 18 | return '[Lane]\n' + str(self.points) + '\n[/Lane]' 19 | 20 | def __call__(self, lane_ys): 21 | lane_xs = self.function(lane_ys) 22 | 23 | lane_xs[(lane_ys < self.min_y) | (lane_ys > self.max_y)] = self.invalid_value 24 | return lane_xs 25 | 26 | def to_array(self, cfg): 27 | sample_y = cfg.sample_y 28 | img_w, img_h = cfg.ori_img_w, cfg.ori_img_h 29 | ys = np.array(sample_y) / float(img_h) 30 | xs = self(ys) 31 | valid_mask = (xs >= 0) & (xs < 1) 32 | lane_xs = xs[valid_mask] * img_w 33 | lane_ys = ys[valid_mask] * img_h 34 | lane = np.concatenate((lane_xs.reshape(-1, 1), lane_ys.reshape(-1, 1)), axis=1) 35 | return lane 36 | 37 | 38 | def __iter__(self): 39 | return self 40 | 41 | def __next__(self): 42 | if self.curr_iter < len(self.points): 43 | self.curr_iter += 1 44 | return self.points[self.curr_iter - 1] 45 | self.curr_iter = 0 46 | raise StopIteration 47 | -------------------------------------------------------------------------------- /lanedet/ops/nms.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018, Grégoire Payen de La Garanderie, Durham University 2 | # All rights reserved. 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # * Redistributions of source code must retain the above copyright notice, this 8 | # list of conditions and the following disclaimer. 9 | # 10 | # * Redistributions in binary form must reproduce the above copyright notice, 11 | # this list of conditions and the following disclaimer in the documentation 12 | # and/or other materials provided with the distribution. 13 | # 14 | # * Neither the name of the copyright holder nor the names of its 15 | # contributors may be used to endorse or promote products derived from 16 | # this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | 29 | from . import nms_impl 30 | 31 | 32 | def nms(boxes, scores, overlap, top_k): 33 | return nms_impl.nms_forward(boxes, scores, overlap, top_k) 34 | -------------------------------------------------------------------------------- /lanedet/utils/net_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os 3 | from torch import nn 4 | import numpy as np 5 | import torch.nn.functional 6 | 7 | 8 | def save_model(net, optim, scheduler, recorder, is_best=False): 9 | model_dir = os.path.join(recorder.work_dir, 'ckpt') 10 | os.system('mkdir -p {}'.format(model_dir)) 11 | epoch = recorder.epoch 12 | ckpt_name = 'best' if is_best else epoch 13 | torch.save({ 14 | 'net': net.state_dict(), 15 | 'optim': optim.state_dict(), 16 | 'scheduler': scheduler.state_dict(), 17 | 'recorder': recorder.state_dict(), 18 | 'epoch': epoch 19 | }, os.path.join(model_dir, '{}.pth'.format(ckpt_name))) 20 | 21 | # remove previous pretrained model if the number of models is too big 22 | # pths = [int(pth.split('.')[0]) for pth in os.listdir(model_dir)] 23 | # if len(pths) <= 2: 24 | # return 25 | # os.system('rm {}'.format(os.path.join(model_dir, '{}.pth'.format(min(pths))))) 26 | 27 | 28 | def load_network_specified(net, model_dir, logger=None): 29 | pretrained_net = torch.load(model_dir)['net'] 30 | net_state = net.state_dict() 31 | state = {} 32 | for k, v in pretrained_net.items(): 33 | if k not in net_state.keys() or v.size() != net_state[k].size(): 34 | if logger: 35 | logger.info('skip weights: ' + k) 36 | continue 37 | state[k] = v 38 | net.load_state_dict(state, strict=False) 39 | 40 | 41 | def load_network(net, model_dir, finetune_from=None, logger=None): 42 | if finetune_from: 43 | if logger: 44 | logger.info('Finetune model from: ' + finetune_from) 45 | load_network_specified(net, finetune_from, logger) 46 | return 47 | pretrained_model = torch.load(model_dir) 48 | net.load_state_dict(pretrained_model['net'], strict=True) 49 | -------------------------------------------------------------------------------- /lanedet/models/aggregators/aspp.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | from torch.hub import load_state_dict_from_url 5 | 6 | from lanedet.models.registry import AGGREGATORS 7 | 8 | class Atrous_module(nn.Module): 9 | def __init__(self, inplanes, planes, rate): 10 | super(Atrous_module, self).__init__() 11 | self.atrous_convolution = nn.Conv2d(inplanes, planes, kernel_size=3, 12 | stride=1, padding=rate, dilation=rate) 13 | self.batch_norm = nn.BatchNorm2d(planes) 14 | 15 | def forward(self, x): 16 | x = self.atrous_convolution(x) 17 | x = self.batch_norm(x) 18 | 19 | return x 20 | 21 | 22 | @AGGREGATORS.register_module 23 | class ASPP(nn.Module): 24 | def __init__(self, cfg): 25 | super(ASPP, self).__init__() 26 | rates = [1, 6, 12, 18] 27 | in_channel = 128 28 | self.aspp1 = Atrous_module(in_channel, 256, rate=rates[0]) 29 | self.aspp2 = Atrous_module(in_channel, 256, rate=rates[1]) 30 | self.aspp3 = Atrous_module(in_channel, 256, rate=rates[2]) 31 | self.aspp4 = Atrous_module(in_channel, 256, rate=rates[3]) 32 | self.image_pool = nn.Sequential(nn.AdaptiveMaxPool2d(1), 33 | nn.Conv2d(in_channel, 256, kernel_size=1)) 34 | self.fc1 = nn.Sequential(nn.Conv2d(1280, 128, kernel_size=1), 35 | nn.BatchNorm2d(128)) 36 | 37 | def forward(self, x): 38 | x1 = self.aspp1(x) 39 | x2 = self.aspp2(x) 40 | x3 = self.aspp3(x) 41 | x4 = self.aspp4(x) 42 | x5 = self.image_pool(x) 43 | x5 = F.upsample(x5, size=x4.size()[2:], mode='nearest') 44 | 45 | x = torch.cat((x1, x2, x3, x4, x5), dim = 1) 46 | 47 | x = self.fc1(x) 48 | 49 | return x 50 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import torch.nn.parallel 4 | import torch.backends.cudnn as cudnn 5 | import argparse 6 | import numpy as np 7 | import random 8 | from lanedet.utils.config import Config 9 | from lanedet.engine.runner import Runner 10 | from lanedet.datasets import build_dataloader 11 | 12 | 13 | def main(): 14 | args = parse_args() 15 | os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(gpu) for gpu in args.gpus) 16 | 17 | cfg = Config.fromfile(args.config) 18 | cfg.gpus = len(args.gpus) 19 | 20 | cfg.load_from = args.load_from 21 | cfg.finetune_from = args.finetune_from 22 | cfg.view = args.view 23 | cfg.seed = args.seed 24 | 25 | cfg.work_dirs = args.work_dirs + '/' + cfg.dataset.train.type 26 | 27 | cudnn.benchmark = True 28 | # cudnn.fastest = True 29 | 30 | runner = Runner(cfg) 31 | 32 | if args.validate: 33 | runner.validate() 34 | elif args.test: 35 | runner.test() 36 | else: 37 | runner.train() 38 | 39 | def parse_args(): 40 | parser = argparse.ArgumentParser(description='Train a detector') 41 | parser.add_argument('config', help='train config file path') 42 | parser.add_argument( 43 | '--work_dirs', type=str, default='work_dirs', 44 | help='work dirs') 45 | parser.add_argument( 46 | '--load_from', default=None, 47 | help='the checkpoint file to resume from') 48 | parser.add_argument( 49 | '--finetune_from', default=None, 50 | help='whether to finetune from the checkpoint') 51 | parser.add_argument( 52 | '--view', action='store_true', 53 | help='whether to view') 54 | parser.add_argument( 55 | '--validate', 56 | action='store_true', 57 | help='whether to evaluate the checkpoint during training') 58 | parser.add_argument( 59 | '--test', 60 | action='store_true', 61 | help='whether to evaluate the checkpoint during training') 62 | parser.add_argument('--gpus', nargs='+', type=int, default='0') 63 | parser.add_argument('--seed', type=int, 64 | default=0, help='random seed') 65 | args = parser.parse_args() 66 | 67 | return args 68 | 69 | 70 | if __name__ == '__main__': 71 | main() 72 | -------------------------------------------------------------------------------- /configs/scnn/resnet18_tusimple.py: -------------------------------------------------------------------------------- 1 | net = dict( 2 | type='Detector', 3 | ) 4 | 5 | backbone = dict( 6 | type='ResNetWrapper', 7 | resnet='resnet18', 8 | pretrained=True, 9 | replace_stride_with_dilation=[False, True, True], 10 | out_conv=True, 11 | ) 12 | featuremap_out_channel = 128 13 | featuremap_out_stride = 8 14 | 15 | aggregator = dict( 16 | type='SCNN', 17 | ) 18 | 19 | sample_y=range(710, 150, -10) 20 | heads = dict( 21 | type='LaneSeg', 22 | decoder=dict(type='PlainDecoder'), 23 | thr=0.6, 24 | sample_y=sample_y, 25 | ) 26 | 27 | optimizer = dict( 28 | type = 'SGD', 29 | lr = 0.025, 30 | weight_decay = 1e-4, 31 | momentum = 0.9 32 | ) 33 | 34 | epochs = 100 35 | batch_size = 8 36 | total_iter = (3616 // batch_size + 1) * epochs 37 | import math 38 | scheduler = dict( 39 | type = 'LambdaLR', 40 | lr_lambda = lambda _iter : math.pow(1 - _iter/total_iter, 0.9) 41 | ) 42 | 43 | bg_weight = 0.4 44 | 45 | img_norm = dict( 46 | mean=[103.939, 116.779, 123.68], 47 | std=[1., 1., 1.] 48 | ) 49 | 50 | img_height = 368 51 | img_width = 640 52 | cut_height = 160 53 | ori_img_h = 720 54 | ori_img_w = 1280 55 | 56 | train_process = [ 57 | dict(type='RandomRotation'), 58 | dict(type='RandomHorizontalFlip'), 59 | dict(type='Resize', size=(img_width, img_height)), 60 | dict(type='Normalize', img_norm=img_norm), 61 | dict(type='ToTensor'), 62 | ] 63 | 64 | val_process = [ 65 | dict(type='Resize', size=(img_width, img_height)), 66 | dict(type='Normalize', img_norm=img_norm), 67 | dict(type='ToTensor', keys=['img']), 68 | ] 69 | 70 | dataset_path = './data/tusimple' 71 | dataset = dict( 72 | train=dict( 73 | type='TuSimple', 74 | data_root=dataset_path, 75 | split='trainval', 76 | processes=train_process, 77 | ), 78 | val=dict( 79 | type='TuSimple', 80 | data_root=dataset_path, 81 | split='test', 82 | processes=val_process, 83 | ), 84 | test=dict( 85 | type='TuSimple', 86 | data_root=dataset_path, 87 | split='test', 88 | processes=val_process, 89 | ) 90 | ) 91 | 92 | 93 | workers = 12 94 | num_classes = 6 + 1 95 | ignore_label = 255 96 | log_interval = 100 97 | eval_ep = 1 98 | save_ep = epochs 99 | test_json_file='data/tusimple/test_label.json' 100 | lr_update_by_epoch = False -------------------------------------------------------------------------------- /configs/ufld/resnet18_culane.py: -------------------------------------------------------------------------------- 1 | net = dict( 2 | type='Detector', 3 | ) 4 | 5 | backbone = dict( 6 | type='ResNetWrapper', 7 | resnet='resnet18', 8 | pretrained=True, 9 | replace_stride_with_dilation=[False, False, False], 10 | out_conv=False, 11 | ) 12 | featuremap_out_channel = 512 13 | 14 | griding_num = 200 15 | num_classes = 4 16 | heads = dict(type='LaneCls', 17 | dim = (griding_num + 1, 18, num_classes)) 18 | 19 | optimizer = dict( 20 | type='SGD', 21 | lr=0.015, 22 | weight_decay=1e-4, 23 | momentum=0.9 24 | ) 25 | 26 | epochs = 50 27 | batch_size = 32 28 | total_iter = (88880 // batch_size + 1) * epochs 29 | import math 30 | scheduler = dict( 31 | type = 'LambdaLR', 32 | lr_lambda = lambda _iter : math.pow(1 - _iter/total_iter, 0.9) 33 | ) 34 | 35 | 36 | img_norm = dict( 37 | mean=[103.939, 116.779, 123.68], 38 | std=[1., 1., 1.] 39 | ) 40 | 41 | ori_img_h = 590 42 | ori_img_w = 1640 43 | img_h = 288 44 | img_w = 800 45 | cut_height=0 46 | sample_y = range(589, 230, -20) 47 | 48 | train_process = [ 49 | dict(type='RandomRotation', degree=(-6, 6)), 50 | dict(type='RandomUDoffsetLABEL', max_offset=100), 51 | dict(type='RandomLROffsetLABEL', max_offset=200), 52 | dict(type='GenerateLaneCls', row_anchor='culane_row_anchor', 53 | num_cols=griding_num, num_classes=num_classes), 54 | dict(type='Resize', size=(img_w, img_h)), 55 | dict(type='Normalize', img_norm=img_norm), 56 | dict(type='ToTensor', keys=['img', 'cls_label']), 57 | ] 58 | 59 | val_process = [ 60 | dict(type='Resize', size=(img_w, img_h)), 61 | dict(type='Normalize', img_norm=img_norm), 62 | dict(type='ToTensor', keys=['img']), 63 | ] 64 | 65 | dataset_type = 'CULane' 66 | dataset_path = './data/CULane' 67 | row_anchor = 'culane_row_anchor' 68 | dataset = dict( 69 | train=dict( 70 | type=dataset_type, 71 | data_root=dataset_path, 72 | split='train', 73 | processes=train_process, 74 | ), 75 | val=dict( 76 | type=dataset_type, 77 | data_root=dataset_path, 78 | split='test', 79 | processes=val_process, 80 | ), 81 | test=dict( 82 | type=dataset_type, 83 | data_root=dataset_path, 84 | split='test', 85 | processes=val_process, 86 | ) 87 | ) 88 | 89 | workers = 12 90 | ignore_label = 255 91 | log_interval = 100 92 | eval_ep = epochs // 5 93 | save_ep = epochs 94 | y_pixel_gap = 20 95 | lr_update_by_epoch = False 96 | -------------------------------------------------------------------------------- /configs/scnn/resnet50_culane.py: -------------------------------------------------------------------------------- 1 | net = dict( 2 | type='Detector', 3 | ) 4 | 5 | backbone = dict( 6 | type='ResNetWrapper', 7 | resnet='resnet50', 8 | pretrained=True, 9 | replace_stride_with_dilation=[False, True, False], 10 | out_conv=True, 11 | in_channels=[64, 128, 256, -1] 12 | ) 13 | featuremap_out_channel = 128 14 | featuremap_out_stride = 8 15 | sample_y = range(589, 230, -20) 16 | 17 | aggregator = dict( 18 | type='SCNN', 19 | ) 20 | 21 | heads = dict( 22 | type='LaneSeg', 23 | decoder=dict(type='PlainDecoder'), 24 | exist=dict(type='ExistHead'), 25 | thr=0.3, 26 | sample_y=sample_y, 27 | ) 28 | 29 | optimizer = dict( 30 | type = 'SGD', 31 | lr = 0.005, 32 | weight_decay = 1e-4, 33 | momentum = 0.9 34 | ) 35 | 36 | epochs = 12 37 | batch_size = 8 38 | total_iter = (88880 // batch_size) * epochs 39 | import math 40 | scheduler = dict( 41 | type = 'LambdaLR', 42 | lr_lambda = lambda _iter : math.pow(1 - _iter/total_iter, 0.9) 43 | ) 44 | 45 | seg_loss_weight = 1.0 46 | eval_ep = 6 47 | save_ep = epochs 48 | 49 | bg_weight = 0.4 50 | 51 | img_norm = dict( 52 | mean=[103.939, 116.779, 123.68], 53 | std=[1., 1., 1.] 54 | ) 55 | 56 | img_height = 288 57 | img_width = 800 58 | cut_height = 240 59 | ori_img_h = 590 60 | ori_img_w = 1640 61 | 62 | train_process = [ 63 | dict(type='RandomRotation', degree=(-2, 2)), 64 | dict(type='RandomHorizontalFlip'), 65 | dict(type='Resize', size=(img_width, img_height)), 66 | dict(type='Normalize', img_norm=img_norm), 67 | dict(type='ToTensor', keys=['img', 'mask', 'lane_exist']), 68 | ] 69 | 70 | val_process = [ 71 | dict(type='Resize', size=(img_width, img_height)), 72 | dict(type='Normalize', img_norm=img_norm), 73 | dict(type='ToTensor', keys=['img']), 74 | ] 75 | 76 | dataset_path = './data/CULane' 77 | dataset = dict( 78 | train=dict( 79 | type='CULane', 80 | data_root=dataset_path, 81 | split='train', 82 | processes=train_process, 83 | ), 84 | val=dict( 85 | type='CULane', 86 | data_root=dataset_path, 87 | split='test', 88 | processes=val_process, 89 | ), 90 | test=dict( 91 | type='CULane', 92 | data_root=dataset_path, 93 | split='test', 94 | processes=val_process, 95 | ) 96 | ) 97 | 98 | workers = 12 99 | num_classes = 4 + 1 100 | ignore_label = 255 101 | log_interval = 1000 102 | 103 | lr_update_by_epoch = False 104 | -------------------------------------------------------------------------------- /configs/resa/resa34_tusimple.py: -------------------------------------------------------------------------------- 1 | net = dict( 2 | type='Detector', 3 | ) 4 | 5 | backbone = dict( 6 | type='ResNetWrapper', 7 | resnet='resnet34', 8 | pretrained=True, 9 | replace_stride_with_dilation=[False, True, True], 10 | out_conv=True, 11 | ) 12 | featuremap_out_channel = 128 13 | featuremap_out_stride = 8 14 | 15 | aggregator = dict( 16 | type='RESA', 17 | direction=['d', 'u', 'r', 'l'], 18 | alpha=2.0, 19 | iter=4, 20 | conv_stride=9, 21 | ) 22 | 23 | sample_y=range(710, 350, -10) 24 | 25 | heads = dict( 26 | type='LaneSeg', 27 | decoder=dict(type='BUSD'), 28 | thr=0.6, 29 | sample_y=sample_y, 30 | ) 31 | 32 | optimizer = dict( 33 | type = 'SGD', 34 | lr = 0.025, 35 | weight_decay = 1e-4, 36 | momentum = 0.9 37 | ) 38 | 39 | 40 | epochs = 150 41 | batch_size = 8 42 | total_iter = (3616 // batch_size + 1) * epochs 43 | 44 | import math 45 | scheduler = dict( 46 | type = 'LambdaLR', 47 | lr_lambda = lambda _iter : math.pow(1 - _iter/total_iter, 0.9) 48 | ) 49 | 50 | bg_weight = 0.4 51 | 52 | img_norm = dict( 53 | mean=[103.939, 116.779, 123.68], 54 | std=[1., 1., 1.] 55 | ) 56 | 57 | img_height = 368 58 | img_width = 640 59 | cut_height = 160 60 | ori_img_h = 720 61 | ori_img_w = 1280 62 | 63 | train_process = [ 64 | dict(type='RandomRotation'), 65 | dict(type='RandomHorizontalFlip'), 66 | dict(type='Resize', size=(img_width, img_height)), 67 | dict(type='Normalize', img_norm=img_norm), 68 | dict(type='ToTensor'), 69 | ] 70 | 71 | val_process = [ 72 | dict(type='Resize', size=(img_width, img_height)), 73 | dict(type='Normalize', img_norm=img_norm), 74 | dict(type='ToTensor', keys=['img']), 75 | ] 76 | 77 | dataset_path = './data/tusimple' 78 | dataset = dict( 79 | train=dict( 80 | type='TuSimple', 81 | data_root=dataset_path, 82 | split='trainval', 83 | processes=train_process, 84 | ), 85 | val=dict( 86 | type='TuSimple', 87 | data_root=dataset_path, 88 | split='test', 89 | processes=val_process, 90 | ), 91 | test=dict( 92 | type='TuSimple', 93 | data_root=dataset_path, 94 | split='test', 95 | processes=val_process, 96 | ) 97 | ) 98 | 99 | 100 | batch_size = 8 101 | workers = 12 102 | num_classes = 6 + 1 103 | ignore_label = 255 104 | log_interval = 100 105 | eval_ep = 1 106 | save_ep = epochs 107 | test_json_file='data/tusimple/test_label.json' 108 | lr_update_by_epoch = False -------------------------------------------------------------------------------- /lanedet/ops/csrc/nms.cpp: -------------------------------------------------------------------------------- 1 | /* Copyright (c) 2018, Grégoire Payen de La Garanderie, Durham University 2 | * All rights reserved. 3 | * 4 | * Redistribution and use in source and binary forms, with or without 5 | * modification, are permitted provided that the following conditions are met: 6 | * 7 | * * Redistributions of source code must retain the above copyright notice, this 8 | * list of conditions and the following disclaimer. 9 | * 10 | * * Redistributions in binary form must reproduce the above copyright notice, 11 | * this list of conditions and the following disclaimer in the documentation 12 | * and/or other materials provided with the distribution. 13 | * 14 | * * Neither the name of the copyright holder nor the names of its 15 | * contributors may be used to endorse or promote products derived from 16 | * this software without specific prior written permission. 17 | * 18 | * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 22 | * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 23 | * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 24 | * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 25 | * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 26 | * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | */ 29 | 30 | #include 31 | #include 32 | #include 33 | 34 | std::vector nms_cuda_forward( 35 | at::Tensor boxes, 36 | at::Tensor idx, 37 | float nms_overlap_thresh, 38 | unsigned long top_k); 39 | 40 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 41 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 42 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 43 | 44 | std::vector nms_forward( 45 | at::Tensor boxes, 46 | at::Tensor scores, 47 | float thresh, 48 | unsigned long top_k) { 49 | 50 | 51 | auto idx = std::get<1>(scores.sort(0,true)); 52 | 53 | CHECK_INPUT(boxes); 54 | CHECK_INPUT(idx); 55 | 56 | return nms_cuda_forward(boxes, idx, thresh, top_k); 57 | } 58 | 59 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 60 | m.def("nms_forward", &nms_forward, "NMS"); 61 | } 62 | 63 | -------------------------------------------------------------------------------- /lanedet/utils/visualization.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import os 3 | import os.path as osp 4 | import numpy as np 5 | 6 | # Color palette for lane visualization 7 | def getcolor(code): 8 | if code == 1: 9 | return (0, 255, 0) 10 | if code == 2: 11 | return (0, 0, 255) 12 | if code == 3: 13 | return (255, 255, 0) 14 | if code == 4: 15 | return (0, 255, 255) 16 | if code == 5: 17 | return (255, 0, 255) 18 | if code == 6: 19 | return (45, 88, 200) 20 | if code == 7: 21 | return (213, 22, 224) 22 | 23 | 24 | def imshow_lanes(img, lanes, show=False, out_file=None, lane_classes = None, num_classes=2): 25 | #img = np.zeros((720, 1280, 3)) 26 | 27 | if lane_classes is not None: 28 | if num_classes == 6: 29 | cv2.putText(img,'solid-yellow',(0,40), cv2.FONT_HERSHEY_SIMPLEX, 1,getcolor(1),2,cv2.LINE_AA) 30 | cv2.putText(img,'solid-white',(0,70), cv2.FONT_HERSHEY_SIMPLEX, 1,getcolor(2),2,cv2.LINE_AA) 31 | cv2.putText(img,'dashed',(0,100), cv2.FONT_HERSHEY_SIMPLEX, 1,getcolor(3),2,cv2.LINE_AA) 32 | cv2.putText(img,'Botts\'-dots',(0,170), cv2.FONT_HERSHEY_SIMPLEX, 1,getcolor(4),2,cv2.LINE_AA) 33 | cv2.putText(img,'double-solid-yellow',(0,200), cv2.FONT_HERSHEY_SIMPLEX, 1,getcolor(5),2,cv2.LINE_AA) 34 | cv2.putText(img,'unknown',(0,230), cv2.FONT_HERSHEY_SIMPLEX, 1,getcolor(6),2,cv2.LINE_AA) 35 | else: 36 | cv2.putText(img,'solid',(0,40), cv2.FONT_HERSHEY_SIMPLEX, 1,getcolor(1),2,cv2.LINE_AA) 37 | cv2.putText(img,'dashed',(0,70), cv2.FONT_HERSHEY_SIMPLEX, 1,getcolor(2),2,cv2.LINE_AA) 38 | #df = {0:0, 1:1, 2:1, 3:2, 4:2, 5:1, 6:1} 39 | #lane_classes = list(map(df.get,lane_classes)) 40 | 41 | 42 | for i, lane in enumerate(lanes): 43 | for x, y in lane: 44 | if x <= 0 or y <= 0: 45 | continue 46 | x, y = int(x), int(y) 47 | if lane_classes is not None: 48 | color = getcolor(lane_classes[i]) 49 | else: 50 | color = (255, 0, 0) 51 | cv2.circle(img, (x, y), 4, color, 2) 52 | 53 | ''' 54 | for i, lane in enumerate(lanes): 55 | for j in range(len(lane)-1): 56 | if lane_classes is not None: 57 | color = getcolor(lane_classes[i]) 58 | else: 59 | color = (255, 0, 0) 60 | cv2.line(img, lane[j], lane[j+1], color, 5) 61 | ''' 62 | 63 | if show: 64 | cv2.imshow('view', img) 65 | cv2.waitKey(0) 66 | 67 | if out_file: 68 | if not osp.exists(osp.dirname(out_file)): 69 | os.makedirs(osp.dirname(out_file)) 70 | cv2.imwrite(out_file, img) 71 | 72 | -------------------------------------------------------------------------------- /configs/resa/resa34_culane.py: -------------------------------------------------------------------------------- 1 | net = dict( 2 | type='Detector', 3 | ) 4 | 5 | backbone = dict( 6 | type='ResNetWrapper', 7 | resnet='resnet34', 8 | pretrained=True, 9 | replace_stride_with_dilation=[False, True, True], 10 | out_conv=True, 11 | in_channels=[64, 128, 256, -1] 12 | ) 13 | featuremap_out_channel = 128 14 | featuremap_out_stride = 8 15 | sample_y = range(589, 230, -20) 16 | 17 | aggregator = dict( 18 | type='RESA', 19 | direction=['d', 'u', 'r', 'l'], 20 | alpha=2.0, 21 | iter=4, 22 | conv_stride=9, 23 | ) 24 | 25 | heads = dict( 26 | type='LaneSeg', 27 | decoder=dict(type='PlainDecoder'), 28 | exist=dict(type='ExistHead'), 29 | thr=0.3, 30 | sample_y=sample_y, 31 | ) 32 | 33 | trainer = dict( 34 | type='RESA' 35 | ) 36 | 37 | evaluator = dict( 38 | type='CULane', 39 | ) 40 | 41 | optimizer = dict( 42 | type = 'SGD', 43 | lr = 0.030, 44 | weight_decay = 1e-4, 45 | momentum = 0.9 46 | ) 47 | 48 | epochs = 12 49 | batch_size = 8 50 | total_iter = (88880 // batch_size) * epochs 51 | import math 52 | scheduler = dict( 53 | type = 'LambdaLR', 54 | lr_lambda = lambda _iter : math.pow(1 - _iter/total_iter, 0.9) 55 | ) 56 | 57 | seg_loss_weight = 1.0 58 | eval_ep = 6 59 | save_ep = epochs 60 | 61 | bg_weight = 0.4 62 | 63 | img_norm = dict( 64 | mean=[103.939, 116.779, 123.68], 65 | std=[1., 1., 1.] 66 | ) 67 | 68 | img_height = 288 69 | img_width = 800 70 | cut_height = 240 71 | ori_img_h = 590 72 | ori_img_w = 1640 73 | 74 | train_process = [ 75 | dict(type='RandomRotation', degree=(-2, 2)), 76 | dict(type='RandomHorizontalFlip'), 77 | dict(type='Resize', size=(img_width, img_height)), 78 | dict(type='Normalize', img_norm=img_norm), 79 | dict(type='ToTensor', keys=['img', 'mask', 'lane_exist']), 80 | ] 81 | 82 | val_process = [ 83 | dict(type='Resize', size=(img_width, img_height)), 84 | dict(type='Normalize', img_norm=img_norm), 85 | dict(type='ToTensor', keys=['img']), 86 | ] 87 | 88 | dataset_path = './data/CULane' 89 | dataset = dict( 90 | train=dict( 91 | type='CULane', 92 | data_root=dataset_path, 93 | split='train', 94 | processes=train_process, 95 | ), 96 | val=dict( 97 | type='CULane', 98 | data_root=dataset_path, 99 | split='test', 100 | processes=val_process, 101 | ), 102 | test=dict( 103 | type='CULane', 104 | data_root=dataset_path, 105 | split='test', 106 | processes=val_process, 107 | ) 108 | ) 109 | 110 | 111 | workers = 12 112 | num_classes = 4 + 1 113 | ignore_label = 255 114 | log_interval = 1000 115 | 116 | lr_update_by_epoch = False 117 | -------------------------------------------------------------------------------- /configs/resa/resa50_culane.py: -------------------------------------------------------------------------------- 1 | net = dict( 2 | type='Detector', 3 | ) 4 | 5 | backbone = dict( 6 | type='ResNetWrapper', 7 | resnet='resnet50', 8 | pretrained=True, 9 | replace_stride_with_dilation=[False, True, True], 10 | out_conv=True, 11 | in_channels=[64, 128, 256, -1] 12 | ) 13 | featuremap_out_channel = 128 14 | featuremap_out_stride = 8 15 | sample_y = range(589, 230, -20) 16 | 17 | aggregator = dict( 18 | type='RESA', 19 | direction=['d', 'u', 'r', 'l'], 20 | alpha=2.0, 21 | iter=4, 22 | conv_stride=9, 23 | ) 24 | 25 | heads = dict( 26 | type='LaneSeg', 27 | decoder=dict(type='PlainDecoder'), 28 | exist=dict(type='ExistHead'), 29 | thr=0.3, 30 | sample_y=sample_y, 31 | ) 32 | 33 | trainer = dict( 34 | type='RESA' 35 | ) 36 | 37 | evaluator = dict( 38 | type='CULane', 39 | ) 40 | 41 | optimizer = dict( 42 | type = 'SGD', 43 | lr = 0.030, 44 | weight_decay = 1e-4, 45 | momentum = 0.9 46 | ) 47 | 48 | epochs = 12 49 | batch_size = 8 50 | total_iter = (88880 // batch_size) * epochs 51 | import math 52 | scheduler = dict( 53 | type = 'LambdaLR', 54 | lr_lambda = lambda _iter : math.pow(1 - _iter/total_iter, 0.9) 55 | ) 56 | 57 | seg_loss_weight = 1.0 58 | eval_ep = 6 59 | save_ep = epochs 60 | 61 | bg_weight = 0.4 62 | 63 | img_norm = dict( 64 | mean=[103.939, 116.779, 123.68], 65 | std=[1., 1., 1.] 66 | ) 67 | 68 | img_height = 288 69 | img_width = 800 70 | cut_height = 240 71 | ori_img_h = 590 72 | ori_img_w = 1640 73 | 74 | train_process = [ 75 | dict(type='RandomRotation', degree=(-2, 2)), 76 | dict(type='RandomHorizontalFlip'), 77 | dict(type='Resize', size=(img_width, img_height)), 78 | dict(type='Normalize', img_norm=img_norm), 79 | dict(type='ToTensor', keys=['img', 'mask', 'lane_exist']), 80 | ] 81 | 82 | val_process = [ 83 | dict(type='Resize', size=(img_width, img_height)), 84 | dict(type='Normalize', img_norm=img_norm), 85 | dict(type='ToTensor', keys=['img']), 86 | ] 87 | 88 | dataset_path = './data/CULane' 89 | dataset = dict( 90 | train=dict( 91 | type='CULane', 92 | data_root=dataset_path, 93 | split='train', 94 | processes=train_process, 95 | ), 96 | val=dict( 97 | type='CULane', 98 | data_root=dataset_path, 99 | split='test', 100 | processes=val_process, 101 | ), 102 | test=dict( 103 | type='CULane', 104 | data_root=dataset_path, 105 | split='test', 106 | processes=val_process, 107 | ) 108 | ) 109 | 110 | 111 | workers = 12 112 | num_classes = 4 + 1 113 | ignore_label = 255 114 | log_interval = 1000 115 | 116 | lr_update_by_epoch = False 117 | -------------------------------------------------------------------------------- /lanedet/utils/registry.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | 3 | import six 4 | 5 | # borrow from mmdetection 6 | 7 | def is_str(x): 8 | """Whether the input is an string instance.""" 9 | return isinstance(x, six.string_types) 10 | 11 | class Registry(object): 12 | 13 | def __init__(self, name): 14 | self._name = name 15 | self._module_dict = dict() 16 | 17 | def __repr__(self): 18 | format_str = self.__class__.__name__ + '(name={}, items={})'.format( 19 | self._name, list(self._module_dict.keys())) 20 | return format_str 21 | 22 | @property 23 | def name(self): 24 | return self._name 25 | 26 | @property 27 | def module_dict(self): 28 | return self._module_dict 29 | 30 | def get(self, key): 31 | return self._module_dict.get(key, None) 32 | 33 | def _register_module(self, module_class): 34 | """Register a module. 35 | 36 | Args: 37 | module (:obj:`nn.Module`): Module to be registered. 38 | """ 39 | if not inspect.isclass(module_class): 40 | raise TypeError('module must be a class, but got {}'.format( 41 | type(module_class))) 42 | module_name = module_class.__name__ 43 | if module_name in self._module_dict: 44 | raise KeyError('{} is already registered in {}'.format( 45 | module_name, self.name)) 46 | self._module_dict[module_name] = module_class 47 | 48 | def register_module(self, cls): 49 | self._register_module(cls) 50 | return cls 51 | 52 | 53 | def build_from_cfg(cfg, registry, default_args=None): 54 | """Build a module from config dict. 55 | 56 | Args: 57 | cfg (dict): Config dict. It should at least contain the key "type". 58 | registry (:obj:`Registry`): The registry to search the type from. 59 | default_args (dict, optional): Default initialization arguments. 60 | 61 | Returns: 62 | obj: The constructed object. 63 | """ 64 | assert isinstance(cfg, dict) and 'type' in cfg 65 | assert isinstance(default_args, dict) or default_args is None 66 | 67 | args = cfg.copy() 68 | obj_type = args.pop('type') 69 | 70 | if is_str(obj_type): 71 | obj_cls = registry.get(obj_type) 72 | #print(obj_cls) 73 | if obj_cls is None: 74 | raise KeyError('{} is not in the {} registry'.format( 75 | obj_type, registry.name)) 76 | elif inspect.isclass(obj_type): 77 | obj_cls = obj_type 78 | else: 79 | raise TypeError('type must be a str or valid type, but got {}'.format( 80 | type(obj_type))) 81 | 82 | if default_args is not None: 83 | for name, value in default_args.items(): 84 | args.setdefault(name, value) 85 | return obj_cls(**args) -------------------------------------------------------------------------------- /configs/resa/resa18_tusimple.py: -------------------------------------------------------------------------------- 1 | net = dict( 2 | type='Detector', 3 | ) 4 | 5 | backbone = dict( 6 | type='ResNetWrapper', 7 | resnet='resnet18', 8 | pretrained=True, 9 | replace_stride_with_dilation=[False, True, True], 10 | out_conv=True, 11 | ) 12 | featuremap_out_channel = 128 13 | featuremap_out_stride = 8 14 | num_classes = 3 15 | num_lanes = 6 + 1 16 | classification = True 17 | autocast = True 18 | 19 | aggregator = dict( 20 | type='RESA', 21 | direction=['d', 'u', 'r', 'l'], 22 | alpha=2.0, 23 | iter=4, 24 | conv_stride=9, 25 | ) 26 | 27 | sample_y=range(710, 150, -10) 28 | heads = dict( 29 | type='LaneSeg', 30 | decoder=dict(type='BUSD'), 31 | thr=0.6, 32 | sample_y=sample_y, 33 | cat_dim = (num_classes, num_lanes - 1) 34 | ) 35 | 36 | optimizer = dict( 37 | type = 'SGD', 38 | lr = 0.025, 39 | weight_decay = 1e-4, 40 | momentum = 0.9 41 | ) 42 | 43 | 44 | epochs = 25 45 | batch_size = 16 46 | total_training_samples = 3626 47 | total_iter = (total_training_samples // batch_size + 1) * epochs 48 | import math 49 | scheduler = dict( 50 | type = 'LambdaLR', 51 | lr_lambda = lambda _iter : math.pow(1 - _iter/total_iter, 0.9) 52 | ) 53 | 54 | bg_weight = 0.4 55 | 56 | img_norm = dict( 57 | mean=[103.939, 116.779, 123.68], 58 | std=[1., 1., 1.] 59 | ) 60 | 61 | img_height = 368 62 | img_width = 640 63 | cut_height = 160 64 | ori_img_h = 720 65 | ori_img_w = 1280 66 | 67 | train_process = [ 68 | dict(type='RandomRotation'), 69 | dict(type='RandomHorizontalFlip'), 70 | dict(type='Resize', size=(img_width, img_height)), 71 | dict(type='Normalize', img_norm=img_norm), 72 | dict(type='ToTensor'), 73 | ] 74 | 75 | val_process = [ 76 | dict(type='Resize', size=(img_width, img_height)), 77 | dict(type='Normalize', img_norm=img_norm), 78 | dict(type='ToTensor'), 79 | ] 80 | 81 | infer_process = [ 82 | dict(type='Resize', size=(img_width, img_height)), 83 | dict(type='Normalize', img_norm=img_norm), 84 | dict(type='ToTensor', keys=['img']), 85 | ] 86 | 87 | dataset_path = './data/tusimple' 88 | dataset = dict( 89 | train=dict( 90 | type='TuSimple', 91 | data_root=dataset_path, 92 | split='trainval', 93 | processes=train_process, 94 | ), 95 | val=dict( 96 | type='TuSimple', 97 | data_root=dataset_path, 98 | split='val', 99 | processes=val_process, 100 | ), 101 | test=dict( 102 | type='TuSimple', 103 | data_root=dataset_path, 104 | split='test', 105 | processes=val_process, 106 | ) 107 | ) 108 | 109 | 110 | workers = 8 111 | ignore_label = 255 112 | log_interval = 200 113 | eval_ep = 1 114 | save_ep = epochs 115 | test_json_file='data/tusimple/test_label.json' 116 | lr_update_by_epoch = False -------------------------------------------------------------------------------- /lanedet/datasets/base_dataset.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import os 3 | import numpy as np 4 | import cv2 5 | import torch 6 | from torch.utils.data import Dataset 7 | import torchvision 8 | import logging 9 | from .registry import DATASETS 10 | from .process import Process 11 | from lanedet.utils.visualization import imshow_lanes 12 | from mmcv.parallel import DataContainer as DC 13 | import inspect 14 | from torch.nn.functional import one_hot 15 | 16 | @DATASETS.register_module 17 | class BaseDataset(Dataset): 18 | def __init__(self, data_root, split, processes=None, 19 | cfg=None): 20 | self.cfg = cfg 21 | self.logger = logging.getLogger(__name__) 22 | self.data_root = data_root 23 | self.training = 'train' in split 24 | self.validation = 'val' in split 25 | self.processes = Process(processes, cfg) 26 | 27 | def view(self, predictions, img_metas): 28 | img_metas = [item for img_meta in img_metas.data for item in img_meta] 29 | for lanes, img_meta in zip(predictions, img_metas): 30 | img_name = img_meta['img_name'] 31 | img = cv2.imread(osp.join(self.data_root, img_name)) 32 | out_file = osp.join(self.cfg.work_dir, 'visualization', 33 | img_name.replace('/', '_')) 34 | lanes = [lane.to_array(self.cfg) for lane in lanes] 35 | imshow_lanes(img, lanes, out_file=out_file) 36 | 37 | def __len__(self): 38 | 'Denotes the total number of samples' 39 | return len(self.data_infos) 40 | 41 | def __getitem__(self, idx): 42 | 'Generates one sample of data' 43 | data_info = self.data_infos[idx] 44 | if not osp.isfile(data_info['img_path']): 45 | raise FileNotFoundError('cannot find file: {}'.format(data_info['img_path'])) 46 | img = cv2.imread(data_info['img_path']) 47 | 48 | img = img[self.cfg.cut_height:, :, :] 49 | sample = data_info.copy() 50 | sample.update({'img': img}) 51 | 52 | if self.training or self.validation: 53 | label = cv2.imread(sample['mask_path'], cv2.IMREAD_UNCHANGED) 54 | if len(label.shape) > 2: 55 | label = label[:, :, 0] 56 | label = label.squeeze() 57 | label = label[self.cfg.cut_height:, :] 58 | sample.update({'mask': label}) 59 | 60 | sample = self.processes(sample) 61 | meta = {'full_img_path': data_info['img_path'], 62 | 'img_name': data_info['img_name']} 63 | meta = DC(meta, cpu_only=True) 64 | sample.update({'meta': meta}) #generate one dict with img, img_path, lane pixels, seg_img 65 | 66 | category = data_info['categories'] 67 | #category = [0 if np.all(sample['cls_label'][:,i].numpy() == 100) else category[i] for i in range(6)] 68 | sample['category'] = torch.LongTensor(category) 69 | #print(sample.keys()) 70 | 71 | return sample -------------------------------------------------------------------------------- /configs/ufld/resnet18_tusimple.py: -------------------------------------------------------------------------------- 1 | net = dict( 2 | type='Detector', 3 | ) 4 | 5 | backbone = dict( 6 | type='ResNetWrapper', 7 | resnet='resnet18', 8 | pretrained=True, 9 | replace_stride_with_dilation=[False, False, False], 10 | out_conv=False, 11 | ) 12 | featuremap_out_channel = 512 13 | 14 | griding_num = 100 15 | num_lanes = 6 16 | classification = True 17 | num_classes = 3 18 | autocast = True 19 | 20 | heads = dict(type='LaneCls', 21 | dim = (griding_num + 1, 56, num_lanes), 22 | cat_dim =(num_classes, num_lanes)) 23 | 24 | trainer = dict( 25 | type='LaneCls' 26 | ) 27 | 28 | evaluator = dict( 29 | type='Tusimple', 30 | ) 31 | 32 | optimizer = dict( 33 | type = 'SGD', 34 | lr = 0.025, 35 | weight_decay = 1e-4, 36 | momentum = 0.9 37 | ) 38 | #optimizer = dict(type='Adam', lr= 0.025, weight_decay = 0.0001) # 3e-4 for batchsize 8 39 | 40 | epochs = 2 41 | batch_size = 4 42 | total_training_samples = 3626 43 | total_iter = (total_training_samples // batch_size + 1) * epochs 44 | 45 | import math 46 | 47 | scheduler = dict( 48 | type = 'LambdaLR', 49 | lr_lambda = lambda _iter : math.pow(1 - _iter/total_iter, 0.9) 50 | ) 51 | img_norm = dict( 52 | mean=[103.939, 116.779, 123.68], 53 | std=[1., 1., 1.] 54 | ) 55 | 56 | ori_img_h = 720 57 | ori_img_w = 1280 58 | img_h = 288 59 | img_w = 800 60 | cut_height= 0 61 | sample_y = range(710, 150, -10) 62 | 63 | dataset_type = 'TuSimple' 64 | dataset_path = './data/tusimple' 65 | row_anchor = 'tusimple_row_anchor' 66 | 67 | train_process = [ 68 | dict(type='RandomRotation', degree=(-6, 6)), 69 | dict(type='RandomUDoffsetLABEL', max_offset=100), 70 | dict(type='RandomLROffsetLABEL', max_offset=200), 71 | dict(type='GenerateLaneCls', row_anchor=row_anchor, 72 | num_cols=griding_num, num_lanes=num_lanes), 73 | dict(type='Resize', size=(img_w, img_h)), 74 | dict(type='Normalize', img_norm=img_norm), 75 | dict(type='ToTensor', keys=['img', 'cls_label']), 76 | ] 77 | 78 | val_process = [ 79 | dict(type='GenerateLaneCls', row_anchor=row_anchor, 80 | num_cols=griding_num, num_lanes=num_lanes), 81 | dict(type='Resize', size=(img_w, img_h)), 82 | dict(type='Normalize', img_norm=img_norm), 83 | dict(type='ToTensor', keys=['img', 'cls_label']), 84 | ] 85 | 86 | infer_process = [ 87 | dict(type='Resize', size=(img_w, img_h)), 88 | dict(type='Normalize', img_norm=img_norm), 89 | dict(type='ToTensor', keys=['img']), 90 | ] 91 | 92 | dataset = dict( 93 | train=dict( 94 | type=dataset_type, 95 | data_root=dataset_path, 96 | split='trainval', 97 | processes=train_process, 98 | ), 99 | val=dict( 100 | type=dataset_type, 101 | data_root=dataset_path, 102 | split='val', 103 | processes=val_process, 104 | ), 105 | test=dict( 106 | type=dataset_type, 107 | data_root=dataset_path, 108 | split='val', 109 | processes=val_process, 110 | ) 111 | ) 112 | 113 | 114 | workers = 8 115 | ignore_label = 255 116 | log_interval = 200 117 | eval_ep = 1 118 | save_ep = epochs 119 | row_anchor='tusimple_row_anchor' 120 | test_json_file='data/tusimple/test_label.json' 121 | lr_update_by_epoch = False 122 | -------------------------------------------------------------------------------- /lanedet/models/aggregators/resa.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | from torch.hub import load_state_dict_from_url 5 | 6 | from lanedet.models.registry import AGGREGATORS 7 | from .aspp import ASPP 8 | 9 | @AGGREGATORS.register_module 10 | class RESA(nn.Module): 11 | def __init__(self, 12 | direction, 13 | alpha, 14 | iter, 15 | conv_stride, 16 | cfg): 17 | super(RESA, self).__init__() 18 | self.cfg = cfg 19 | self.iter = iter 20 | chan = cfg.featuremap_out_channel 21 | fea_stride = cfg.featuremap_out_stride 22 | self.height = cfg.img_height // fea_stride 23 | self.width = cfg.img_width // fea_stride 24 | self.alpha = alpha 25 | 26 | for i in range(self.iter): 27 | conv_vert1 = nn.Conv2d( 28 | chan, chan, (1, conv_stride), 29 | padding=(0, conv_stride//2), groups=1, bias=False) 30 | conv_vert2 = nn.Conv2d( 31 | chan, chan, (1, conv_stride), 32 | padding=(0, conv_stride//2), groups=1, bias=False) 33 | 34 | setattr(self, 'conv_d'+str(i), conv_vert1) 35 | setattr(self, 'conv_u'+str(i), conv_vert2) 36 | 37 | conv_hori1 = nn.Conv2d( 38 | chan, chan, (conv_stride, 1), 39 | padding=(conv_stride//2, 0), groups=1, bias=False) 40 | conv_hori2 = nn.Conv2d( 41 | chan, chan, (conv_stride, 1), 42 | padding=(conv_stride//2, 0), groups=1, bias=False) 43 | 44 | setattr(self, 'conv_r'+str(i), conv_hori1) 45 | setattr(self, 'conv_l'+str(i), conv_hori2) 46 | 47 | idx_d = (torch.arange(self.height) + self.height // 48 | 2**(self.iter - i)) % self.height 49 | setattr(self, 'idx_d'+str(i), idx_d) 50 | 51 | idx_u = (torch.arange(self.height) - self.height // 52 | 2**(self.iter - i)) % self.height 53 | setattr(self, 'idx_u'+str(i), idx_u) 54 | 55 | idx_r = (torch.arange(self.width) + self.width // 56 | 2**(self.iter - i)) % self.width 57 | setattr(self, 'idx_r'+str(i), idx_r) 58 | 59 | idx_l = (torch.arange(self.width) - self.width // 60 | 2**(self.iter - i)) % self.width 61 | setattr(self, 'idx_l'+str(i), idx_l) 62 | 63 | def update(self, x): 64 | height, width = x.size(2), x.size(3) 65 | for i in range(self.iter): 66 | idx_d = (torch.arange(height) + height // 67 | 2**(self.iter - i)) % height 68 | setattr(self, 'idx_d'+str(i), idx_d) 69 | 70 | idx_u = (torch.arange(height) - height // 71 | 2**(self.iter - i)) % height 72 | setattr(self, 'idx_u'+str(i), idx_u) 73 | 74 | idx_r = (torch.arange(width) + width // 75 | 2**(self.iter - i)) % width 76 | setattr(self, 'idx_r'+str(i), idx_r) 77 | 78 | idx_l = (torch.arange(width) - width // 79 | 2**(self.iter - i)) % width 80 | setattr(self, 'idx_l'+str(i), idx_l) 81 | 82 | def forward(self, x): 83 | x = x.clone() 84 | self.update(x) 85 | 86 | for direction in self.cfg.aggregator.direction: 87 | for i in range(self.iter): 88 | conv = getattr(self, 'conv_' + direction + str(i)) 89 | idx = getattr(self, 'idx_' + direction + str(i)) 90 | if direction in ['d', 'u']: 91 | x.add_(self.alpha * F.relu(conv(x[..., idx, :]))) 92 | else: 93 | x.add_(self.alpha * F.relu(conv(x[..., idx]))) 94 | 95 | return x 96 | -------------------------------------------------------------------------------- /lanedet/datasets/culane.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path as osp 3 | import numpy as np 4 | from .base_dataset import BaseDataset 5 | from .registry import DATASETS 6 | import lanedet.utils.culane_metric as culane_metric 7 | import cv2 8 | from tqdm import tqdm 9 | import logging 10 | 11 | LIST_FILE = { 12 | 'train': 'list/train_gt.txt', 13 | 'val': 'list/test.txt', 14 | 'test': 'list/test.txt', 15 | } 16 | 17 | @DATASETS.register_module 18 | class CULane(BaseDataset): 19 | def __init__(self, data_root, split, processes=None, cfg=None): 20 | super().__init__(data_root, split, processes=processes, cfg=cfg) 21 | self.list_path = osp.join(data_root, LIST_FILE[split]) 22 | self.load_annotations() 23 | 24 | def load_annotations(self): 25 | self.logger.info('Loading CULane annotations...') 26 | self.data_infos = [] 27 | with open(self.list_path) as list_file: 28 | for line in list_file: 29 | infos = self.load_annotation(line.split()) 30 | self.data_infos.append(infos) 31 | 32 | def load_annotation(self, line): 33 | infos = {} 34 | img_line = line[0] 35 | img_line = img_line[1 if img_line[0] == '/' else 0::] 36 | img_path = os.path.join(self.data_root, img_line) 37 | infos['img_name'] = img_line 38 | infos['img_path'] = img_path 39 | if len(line) > 1: 40 | mask_line = line[1] 41 | mask_line = mask_line[1 if mask_line[0] == '/' else 0::] 42 | mask_path = os.path.join(self.data_root, mask_line) 43 | infos['mask_path'] = mask_path 44 | 45 | if len(line) > 2: 46 | exist_list = [int(l) for l in line[2:]] 47 | infos['lane_exist'] = np.array(exist_list) 48 | 49 | anno_path = img_path[:-3] + 'lines.txt' # remove sufix jpg and add lines.txt 50 | with open(anno_path, 'r') as anno_file: 51 | data = [list(map(float, line.split())) for line in anno_file.readlines()] 52 | lanes = [[(lane[i], lane[i + 1]) for i in range(0, len(lane), 2) if lane[i] >= 0 and lane[i + 1] >= 0] 53 | for lane in data] 54 | lanes = [list(set(lane)) for lane in lanes] # remove duplicated points 55 | lanes = [lane for lane in lanes if len(lane) > 3] # remove lanes with less than 2 points 56 | 57 | lanes = [sorted(lane, key=lambda x: x[1]) for lane in lanes] # sort by y 58 | infos['lanes'] = lanes 59 | 60 | return infos 61 | 62 | def get_prediction_string(self, pred): 63 | ys = np.array(list(self.cfg.sample_y))[::-1] / self.cfg.ori_img_h 64 | out = [] 65 | for lane in pred: 66 | xs = lane(ys) 67 | valid_mask = (xs >= 0) & (xs < 1) 68 | xs = xs * self.cfg.ori_img_w 69 | lane_xs = xs[valid_mask] 70 | lane_ys = ys[valid_mask] * self.cfg.ori_img_h 71 | lane_xs, lane_ys = lane_xs[::-1], lane_ys[::-1] 72 | lane_str = ' '.join(['{:.5f} {:.5f}'.format(x, y) for x, y in zip(lane_xs, lane_ys)]) 73 | if lane_str != '': 74 | out.append(lane_str) 75 | 76 | return '\n'.join(out) 77 | 78 | def evaluate(self, predictions, output_basedir): 79 | print('Generating prediction output...') 80 | for idx, pred in enumerate(tqdm(predictions)): 81 | output_dir = os.path.join(output_basedir, os.path.dirname(self.data_infos[idx]['img_name'])) 82 | output_filename = os.path.basename(self.data_infos[idx]['img_name'])[:-3] + 'lines.txt' 83 | os.makedirs(output_dir, exist_ok=True) 84 | output = self.get_prediction_string(pred) 85 | with open(os.path.join(output_dir, output_filename), 'w') as out_file: 86 | out_file.write(output) 87 | result = culane_metric.eval_predictions(output_basedir, self.data_root, self.list_path, official=True) 88 | self.logger.info(result) 89 | return result['F1'] 90 | -------------------------------------------------------------------------------- /tools/detect.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import cv2 4 | import os 5 | import os.path as osp 6 | import glob 7 | import argparse 8 | from lanedet.datasets.process import Process 9 | from lanedet.models.registry import build_net 10 | from lanedet.utils.config import Config 11 | from lanedet.utils.visualization import imshow_lanes 12 | from lanedet.utils.net_utils import load_network 13 | from pathlib import Path 14 | from tqdm import tqdm 15 | import torch.nn.functional as F 16 | 17 | class Detect(object): 18 | def __init__(self, cfg): 19 | self.cfg = cfg 20 | self.processes = Process(cfg.infer_process, cfg) 21 | self.net = build_net(self.cfg) 22 | #print(self.net) 23 | self.net = torch.nn.parallel.DataParallel( 24 | self.net, device_ids = range(1)).cuda() 25 | total_params = sum(param.numel() for param in self.net.parameters()) 26 | print(f"total number of params: {total_params}") 27 | self.net.eval() 28 | load_network(self.net, self.cfg.load_from) 29 | 30 | def preprocess(self, img_path): 31 | ori_img = cv2.imread(img_path) 32 | img = ori_img[self.cfg.cut_height:, :, :].astype(np.float32) 33 | data = {'img': img, 'lanes': []} 34 | data = self.processes(data) 35 | data['img'] = data['img'].unsqueeze(0) # add one dimension; make it (1,3,h,w) shape 36 | data.update({'img_path':img_path, 'ori_img':ori_img}) # add image path and original image in the dict 37 | return data 38 | 39 | def inference(self, data): 40 | with torch.no_grad(): 41 | data = self.net(data) 42 | lane_detection, lane_indx = list(self.net.module.get_lanes(data).values()) 43 | if self.cfg.classification: 44 | lane_classes = self.get_lane_class(data, lane_indx) 45 | #print(lane_classes) 46 | return lane_detection[0], lane_classes 47 | return lane_detection 48 | 49 | def get_lane_class(self, predictions, lane_indx): 50 | score = F.softmax(predictions['category'], dim=1) 51 | y_pred = score.argmax(dim=1).squeeze() 52 | return y_pred[lane_indx].detach().cpu().numpy() 53 | 54 | def show(self, data, lane_classes=None): 55 | out_file = self.cfg.savedir 56 | if out_file: 57 | out_file = osp.join(out_file, osp.basename(data['img_path'])) 58 | lanes = [lane.to_array(self.cfg).astype(int) for lane in data['lanes']] 59 | #print(lanes) 60 | imshow_lanes(data['ori_img'], lanes, show=self.cfg.show, out_file=out_file, lane_classes=lane_classes) 61 | 62 | def run(self, data): 63 | data = self.preprocess(data) 64 | lane_classes = None 65 | if self.cfg.classification: 66 | data['lanes'], lane_classes = self.inference(data) 67 | else: 68 | data['lanes'] = self.inference(data) 69 | if self.cfg.show or self.cfg.savedir: 70 | self.show(data, lane_classes) 71 | #return data 72 | 73 | def get_img_paths(path): 74 | p = str(Path(path).absolute()) # os-agnostic absolute path 75 | if '*' in p: 76 | paths = sorted(glob.glob(p, recursive=True)) # glob 77 | elif os.path.isdir(p): 78 | paths = sorted(glob.glob(os.path.join(p, '*.*'))) # dir 79 | elif os.path.isfile(p): 80 | paths = [p] # files 81 | else: 82 | raise Exception(f'ERROR: {p} does not exist') 83 | return paths 84 | 85 | def process(args): 86 | cfg = Config.fromfile(args.config) 87 | cfg.show = args.show 88 | cfg.savedir = args.savedir 89 | cfg.load_from = args.load_from 90 | detect = Detect(cfg) 91 | paths = get_img_paths(args.img) 92 | for p in tqdm(paths): 93 | detect.run(p) 94 | 95 | if __name__ == '__main__': 96 | parser = argparse.ArgumentParser() 97 | parser.add_argument('config', help='The path of config file') 98 | parser.add_argument('--img', help='The path of the img (img file or img_folder), for example: data/*.png') 99 | parser.add_argument('--show', action='store_true', default=False, help='Whether to show the image') 100 | parser.add_argument('--savedir', type=str, default='./', help='The root of save directory') 101 | parser.add_argument('--load_from', type=str, default='best.pth', help='The path of model') 102 | args = parser.parse_args() 103 | process(args) 104 | -------------------------------------------------------------------------------- /lanedet/utils/tusimple_metric.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from sklearn.linear_model import LinearRegression 3 | import json as json 4 | 5 | 6 | class LaneEval(object): 7 | lr = LinearRegression() 8 | pixel_thresh = 20 9 | pt_thresh = 0.85 10 | 11 | @staticmethod 12 | def get_angle(xs, y_samples): 13 | xs, ys = xs[xs >= 0], y_samples[xs >= 0] 14 | if len(xs) > 1: 15 | LaneEval.lr.fit(ys[:, None], xs) 16 | k = LaneEval.lr.coef_[0] 17 | theta = np.arctan(k) 18 | else: 19 | theta = 0 20 | return theta 21 | 22 | @staticmethod 23 | def line_accuracy(pred, gt, thresh): 24 | pred = np.array([p if p >= 0 else -100 for p in pred]) 25 | gt = np.array([g if g >= 0 else -100 for g in gt]) 26 | return np.sum(np.where(np.abs(pred - gt) < thresh, 1., 0.)) / len(gt) 27 | 28 | @staticmethod 29 | def bench(pred, gt, y_samples, running_time): 30 | if any(len(p) != len(y_samples) for p in pred): 31 | raise Exception('Format of lanes error.') 32 | if running_time > 200 or len(gt) + 2 < len(pred): 33 | return 0., 0., 1. 34 | angles = [LaneEval.get_angle( 35 | np.array(x_gts), np.array(y_samples)) for x_gts in gt] 36 | threshs = [LaneEval.pixel_thresh / np.cos(angle) for angle in angles] 37 | line_accs = [] 38 | fp, fn = 0., 0. 39 | matched = 0. 40 | for x_gts, thresh in zip(gt, threshs): 41 | accs = [LaneEval.line_accuracy( 42 | np.array(x_preds), np.array(x_gts), thresh) for x_preds in pred] 43 | max_acc = np.max(accs) if len(accs) > 0 else 0. 44 | if max_acc < LaneEval.pt_thresh: 45 | fn += 1 46 | else: 47 | matched += 1 48 | line_accs.append(max_acc) 49 | fp = len(pred) - matched 50 | if len(gt) > 4 and fn > 0: 51 | fn -= 1 52 | s = sum(line_accs) 53 | if len(gt) > 4: 54 | s -= min(line_accs) 55 | return s / max(min(4.0, len(gt)), 1.), fp / len(pred) if len(pred) > 0 else 0., fn / max(min(len(gt), 4.), 1.) 56 | 57 | @staticmethod 58 | def bench_one_submit(pred_file, gt_file): 59 | try: 60 | json_pred = [json.loads(line) 61 | for line in open(pred_file).readlines()] 62 | except BaseException as e: 63 | raise Exception('Fail to load json file of the prediction.') 64 | json_gt = [json.loads(line) for line in open(gt_file).readlines()] 65 | if len(json_gt) != len(json_pred): 66 | raise Exception( 67 | 'We do not get the predictions of all the test tasks') 68 | gts = {l['raw_file']: l for l in json_gt} 69 | accuracy, fp, fn = 0., 0., 0. 70 | for pred in json_pred: 71 | if 'raw_file' not in pred or 'lanes' not in pred or 'run_time' not in pred: 72 | raise Exception( 73 | 'raw_file or lanes or run_time not in some predictions.') 74 | raw_file = pred['raw_file'] 75 | pred_lanes = pred['lanes'] 76 | run_time = pred['run_time'] 77 | if raw_file not in gts: 78 | raise Exception( 79 | 'Some raw_file from your predictions do not exist in the test tasks.') 80 | gt = gts[raw_file] 81 | gt_lanes = gt['lanes'] 82 | y_samples = gt['h_samples'] 83 | try: 84 | a, p, n = LaneEval.bench( 85 | pred_lanes, gt_lanes, y_samples, run_time) 86 | except BaseException as e: 87 | raise Exception('Format of lanes error.') 88 | accuracy += a 89 | fp += p 90 | fn += n 91 | num = len(gts) 92 | # the first return parameter is the default ranking parameter 93 | return json.dumps([ 94 | {'name': 'Accuracy', 'value': accuracy / num, 'order': 'desc'}, 95 | {'name': 'FP', 'value': fp / num, 'order': 'asc'}, 96 | {'name': 'FN', 'value': fn / num, 'order': 'asc'} 97 | ]), accuracy / num 98 | 99 | 100 | if __name__ == '__main__': 101 | import sys 102 | try: 103 | if len(sys.argv) != 3: 104 | raise Exception('Invalid input arguments') 105 | print(LaneEval.bench_one_submit(sys.argv[1], sys.argv[2])) 106 | except Exception as e: 107 | print(e.message) 108 | sys.exit(e.message) 109 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import os 3 | import re 4 | from setuptools import find_packages, setup 5 | from torch.utils.cpp_extension import CUDAExtension, BuildExtension 6 | 7 | 8 | def parse_requirements(fname='requirements.txt', with_version=True): 9 | """Parse the package dependencies listed in a requirements file but strips 10 | specific versioning information. 11 | Args: 12 | fname (str): path to requirements file 13 | with_version (bool, default=False): if True include version specs 14 | Returns: 15 | List[str]: list of requirements items 16 | CommandLine: 17 | python -c "import setup; print(setup.parse_requirements())" 18 | """ 19 | import sys 20 | from os.path import exists 21 | require_fpath = fname 22 | 23 | def parse_line(line): 24 | """Parse information from a line in a requirements text file.""" 25 | if line.startswith('-r '): 26 | # Allow specifying requirements in other files 27 | target = line.split(' ')[1] 28 | for info in parse_require_file(target): 29 | yield info 30 | else: 31 | info = {'line': line} 32 | if line.startswith('-e '): 33 | info['package'] = line.split('#egg=')[1] 34 | else: 35 | # Remove versioning from the package 36 | pat = '(' + '|'.join(['>=', '==', '>']) + ')' 37 | parts = re.split(pat, line, maxsplit=1) 38 | parts = [p.strip() for p in parts] 39 | 40 | info['package'] = parts[0] 41 | if len(parts) > 1: 42 | op, rest = parts[1:] 43 | if ';' in rest: 44 | # Handle platform specific dependencies 45 | # http://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-platform-specific-dependencies 46 | version, platform_deps = map(str.strip, 47 | rest.split(';')) 48 | info['platform_deps'] = platform_deps 49 | else: 50 | version = rest # NOQA 51 | info['version'] = (op, version) 52 | yield info 53 | 54 | def parse_require_file(fpath): 55 | with open(fpath, 'r') as f: 56 | for line in f.readlines(): 57 | line = line.strip() 58 | if line and not line.startswith('#'): 59 | for info in parse_line(line): 60 | yield info 61 | 62 | def gen_packages_items(): 63 | if exists(require_fpath): 64 | for info in parse_require_file(require_fpath): 65 | parts = [info['package']] 66 | if with_version and 'version' in info: 67 | parts.extend(info['version']) 68 | if not sys.version.startswith('3.4'): 69 | # apparently package_deps are broken in 3.4 70 | platform_deps = info.get('platform_deps') 71 | if platform_deps is not None: 72 | parts.append(';' + platform_deps) 73 | item = ''.join(parts) 74 | yield item 75 | 76 | packages = list(gen_packages_items()) 77 | return packages 78 | 79 | 80 | install_requires = parse_requirements() 81 | 82 | 83 | def get_extensions(): 84 | extensions = [] 85 | 86 | op_files = glob.glob('./lanedet/ops/csrc/*.c*') 87 | extension = CUDAExtension 88 | ext_name = 'lanedet.ops.nms_impl' 89 | 90 | ext_ops = extension( 91 | name=ext_name, 92 | sources=op_files, 93 | ) 94 | 95 | extensions.append(ext_ops) 96 | 97 | return extensions 98 | 99 | 100 | setup( 101 | name='lanedet', 102 | version="1.0", 103 | description='Lane Detection open toolbox via Pytorch', 104 | keywords='computer vision & lane detection', 105 | classifiers=['License :: OSI Approved :: MIT License', 106 | 'Programming Language :: Python :: 3', 107 | 'Intended Audience :: Developers', 108 | 'Operating System :: OS Independent'], 109 | packages=find_packages(), 110 | include_package_data=True, 111 | url='https://github.com//turoad/lanedet', 112 | author='Tu Zheng', 113 | author_email='zhengtuzju@gmail.com', 114 | setup_requires=['pytest-runner'], 115 | tests_require=['pytest'], 116 | install_requires=install_requires, 117 | ext_modules=get_extensions(), 118 | cmdclass={'build_ext': BuildExtension}, 119 | zip_safe=False) 120 | -------------------------------------------------------------------------------- /lanedet/models/backbones/vgg.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | from torch.hub import load_state_dict_from_url 5 | 6 | from lanedet.models.registry import BACKBONES 7 | 8 | model_urls = { 9 | 'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth', 10 | 'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth', 11 | 'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth', 12 | 'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth', 13 | 'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth', 14 | 'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth', 15 | 'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth', 16 | 'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth', 17 | } 18 | 19 | 20 | 21 | @BACKBONES.register_module 22 | class VGG(nn.Module): 23 | def __init__(self, cfg): 24 | super(VGG, self).__init__() 25 | 26 | self.conv1_1 = nn.Conv2d(3, 64, 3, padding=1, bias=False) 27 | self.bn1_1 = nn.BatchNorm2d(64) 28 | self.conv1_2 = nn.Conv2d(64, 64, 3, padding=1, bias=False) 29 | self.bn1_2 = nn.BatchNorm2d(64) 30 | 31 | self.conv2_1 = nn.Conv2d(64, 128, 3, padding=1, bias=False) 32 | self.bn2_1 = nn.BatchNorm2d(128) 33 | self.conv2_2 = nn.Conv2d(128, 128, 3, padding=1, bias=False) 34 | self.bn2_2 = nn.BatchNorm2d(128) 35 | 36 | self.conv3_1 = nn.Conv2d(128, 256, 3, padding=1, bias=False) 37 | self.bn3_1 = nn.BatchNorm2d(256) 38 | self.conv3_2 = nn.Conv2d(256, 256, 3, padding=1, bias=False) 39 | self.bn3_2 = nn.BatchNorm2d(256) 40 | self.conv3_3 = nn.Conv2d(256, 256, 3, padding=1, bias=False) 41 | self.bn3_3 = nn.BatchNorm2d(256) 42 | 43 | self.conv4_1 = nn.Conv2d(256, 512, 3, padding=1, bias=False) 44 | self.bn4_1 = nn.BatchNorm2d(512) 45 | self.conv4_2 = nn.Conv2d(512, 512, 3, padding=1, bias=False) 46 | self.bn4_2 = nn.BatchNorm2d(512) 47 | self.conv4_3 = nn.Conv2d(512, 512, 3, padding=1, bias=False) 48 | self.bn4_3 = nn.BatchNorm2d(512) 49 | 50 | self.conv5_1 = nn.Conv2d( 51 | 512, 512, 3, padding=2, dilation=2, bias=False) 52 | self.bn5_1 = nn.BatchNorm2d(512) 53 | self.conv5_2 = nn.Conv2d( 54 | 512, 512, 3, padding=2, dilation=2, bias=False) 55 | self.bn5_2 = nn.BatchNorm2d(512) 56 | self.conv5_3 = nn.Conv2d( 57 | 512, 512, 3, padding=2, dilation=2, bias=False) 58 | self.bn5_3 = nn.BatchNorm2d(512) 59 | 60 | self.conv6 = nn.Conv2d(512, 1024, 3, padding=4, dilation=4, bias=False) 61 | self.bn6 = nn.BatchNorm2d(1024) 62 | self.conv7 = nn.Conv2d(1024, 128, 1, bias=False) 63 | self.bn7 = nn.BatchNorm2d(128) 64 | self._initialize_weights() 65 | 66 | 67 | def _initialize_weights(self) -> None: 68 | for m in self.modules(): 69 | if isinstance(m, nn.Conv2d): 70 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 71 | if m.bias is not None: 72 | nn.init.constant_(m.bias, 0) 73 | elif isinstance(m, nn.BatchNorm2d): 74 | nn.init.constant_(m.weight, 1) 75 | nn.init.constant_(m.bias, 0) 76 | elif isinstance(m, nn.Linear): 77 | nn.init.normal_(m.weight, 0, 0.01) 78 | nn.init.constant_(m.bias, 0) 79 | 80 | 81 | 82 | def forward(self, x): 83 | """ 84 | return 1/8 feature 85 | """ 86 | x = F.relu(self.bn1_1(self.conv1_1(x))) 87 | x = F.relu(self.bn1_2(self.conv1_2(x))) 88 | x = F.max_pool2d(x, 2, stride=2, padding=0) 89 | 90 | x = F.relu(self.bn2_1(self.conv2_1(x))) 91 | x = F.relu(self.bn2_2(self.conv2_2(x))) 92 | x = F.max_pool2d(x, 2, stride=2, padding=0) 93 | 94 | x = F.relu(self.bn3_1(self.conv3_1(x))) 95 | x = F.relu(self.bn3_2(self.conv3_2(x))) 96 | x = F.relu(self.bn3_3(self.conv3_3(x))) 97 | x = F.max_pool2d(x, 2, stride=2, padding=0) 98 | 99 | x = F.relu(self.bn4_1(self.conv4_1(x))) 100 | x = F.relu(self.bn4_2(self.conv4_2(x))) 101 | x = F.relu(self.bn4_3(self.conv4_3(x))) 102 | 103 | x = F.relu(self.bn5_1(self.conv5_1(x))) 104 | x = F.relu(self.bn5_2(self.conv5_2(x))) 105 | x = F.relu(self.bn5_3(self.conv5_3(x))) 106 | 107 | x = F.relu(self.bn6(self.conv6(x))) 108 | x = F.relu(self.bn7(self.conv7(x))) 109 | 110 | return [x] 111 | -------------------------------------------------------------------------------- /lanedet/utils/recorder.py: -------------------------------------------------------------------------------- 1 | from collections import deque, defaultdict 2 | import torch 3 | import os 4 | import datetime 5 | from .logger import init_logger 6 | import logging 7 | import pathspec 8 | 9 | 10 | class SmoothedValue(object): 11 | """Track a series of values and provide access to smoothed values over a 12 | window or the global series average. 13 | """ 14 | 15 | def __init__(self, window_size=20): 16 | self.deque = deque(maxlen=window_size) 17 | self.total = 0.0 18 | self.count = 0 19 | 20 | def update(self, value): 21 | self.deque.append(value) 22 | self.count += 1 23 | self.total += value 24 | 25 | @property 26 | def median(self): 27 | d = torch.tensor(list(self.deque)) 28 | return d.median().item() 29 | 30 | @property 31 | def avg(self): 32 | d = torch.tensor(list(self.deque)) 33 | return d.mean().item() 34 | 35 | @property 36 | def global_avg(self): 37 | return self.total / self.count 38 | 39 | 40 | class Recorder(object): 41 | def __init__(self, cfg): 42 | self.cfg = cfg 43 | self.work_dir = self.get_work_dir() 44 | cfg.work_dir = self.work_dir 45 | self.log_path = os.path.join(self.work_dir, 'log.txt') 46 | 47 | init_logger(self.log_path) 48 | self.logger = logging.getLogger(__name__) 49 | #self.logger.info('Config: \n' + cfg.text) 50 | 51 | self.save_cfg(cfg) 52 | self.cp_projects(self.work_dir) 53 | 54 | # scalars 55 | self.epoch = 0 56 | self.step = 0 57 | self.loss_stats = defaultdict(SmoothedValue) 58 | self.batch_time = SmoothedValue() 59 | self.data_time = SmoothedValue() 60 | self.max_iter = self.cfg.total_iter 61 | self.lr = 0. 62 | 63 | def save_cfg(self, cfg): 64 | cfg_path = os.path.join(self.work_dir, 'config.py') 65 | with open(cfg_path, 'w') as cfg_file: 66 | cfg_file.write(cfg.text) 67 | 68 | def cp_projects(self, to_path): 69 | with open('./.gitignore','r') as fp: 70 | ign = fp.read() 71 | ign += '\n.git' 72 | spec = pathspec.PathSpec.from_lines(pathspec.patterns.GitWildMatchPattern, ign.splitlines()) 73 | all_files = {os.path.join(root,name) for root,dirs,files in os.walk('./') for name in files} 74 | matches = spec.match_files(all_files) 75 | matches = set(matches) 76 | to_cp_files = all_files - matches 77 | for f in to_cp_files: 78 | dirs = os.path.join(to_path,'code',os.path.split(f[2:])[0]) 79 | if not os.path.exists(dirs): 80 | os.makedirs(dirs) 81 | os.system('cp %s %s'%(f,os.path.join(to_path,'code',f[2:]))) 82 | 83 | def get_work_dir(self): 84 | now = datetime.datetime.now().strftime('%Y%m%d_%H%M%S') 85 | hyper_param_str = '_lr_%1.0e_b_%d' % (self.cfg.optimizer.lr, self.cfg.batch_size) 86 | work_dir = os.path.join(self.cfg.work_dirs, now + hyper_param_str) 87 | if not os.path.exists(work_dir): 88 | os.makedirs(work_dir) 89 | return work_dir 90 | 91 | def update_loss_stats(self, loss_dict): 92 | for k, v in loss_dict.items(): 93 | if not isinstance(v, torch.Tensor): continue 94 | self.loss_stats[k].update(v.detach().cpu()) 95 | 96 | def record(self, prefix, step=-1, loss_stats=None, image_stats=None): 97 | self.logger.info(self) 98 | # self.write(str(self)) 99 | 100 | def write(self, content): 101 | with open(self.log_path, 'a+') as f: 102 | f.write(content) 103 | f.write('\n') 104 | 105 | def state_dict(self): 106 | scalar_dict = {} 107 | scalar_dict['step'] = self.step 108 | return scalar_dict 109 | 110 | def load_state_dict(self, scalar_dict): 111 | self.step = scalar_dict['step'] 112 | 113 | def __str__(self): 114 | loss_state = [] 115 | for k, v in self.loss_stats.items(): 116 | loss_state.append('{}: {:.4f}'.format(k, v.avg)) 117 | loss_state = ' '.join(loss_state) 118 | 119 | recording_state = ' '.join(['epoch: {}', 'step: {}', 'lr: {:.4f}', '{}', 'data: {:.4f}', 'batch: {:.4f}', 'eta: {}']) 120 | eta_seconds = self.batch_time.global_avg * (self.max_iter - self.step) 121 | eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) 122 | return recording_state.format(self.epoch, self.step, self.lr, loss_state, self.data_time.avg, self.batch_time.avg, eta_string) 123 | 124 | 125 | def build_recorder(cfg): 126 | return Recorder(cfg) 127 | 128 | -------------------------------------------------------------------------------- /lanedet/models/heads/busd.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | from torch.hub import load_state_dict_from_url 5 | 6 | from ..registry import HEADS 7 | 8 | def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): 9 | """3x3 convolution with padding""" 10 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 11 | padding=dilation, groups=groups, bias=False, dilation=dilation) 12 | 13 | 14 | def conv1x1(in_planes, out_planes, stride=1): 15 | """1x1 convolution""" 16 | return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) 17 | 18 | 19 | class non_bottleneck_1d(nn.Module): 20 | def __init__(self, chann, dropprob, dilated): 21 | super().__init__() 22 | 23 | self.conv3x1_1 = nn.Conv2d( 24 | chann, chann, (3, 1), stride=1, padding=(1, 0), bias=True) 25 | 26 | self.conv1x3_1 = nn.Conv2d( 27 | chann, chann, (1, 3), stride=1, padding=(0, 1), bias=True) 28 | 29 | self.bn1 = nn.BatchNorm2d(chann, eps=1e-03) 30 | 31 | self.conv3x1_2 = nn.Conv2d(chann, chann, (3, 1), stride=1, padding=(1 * dilated, 0), bias=True, 32 | dilation=(dilated, 1)) 33 | 34 | self.conv1x3_2 = nn.Conv2d(chann, chann, (1, 3), stride=1, padding=(0, 1 * dilated), bias=True, 35 | dilation=(1, dilated)) 36 | 37 | self.bn2 = nn.BatchNorm2d(chann, eps=1e-03) 38 | 39 | self.dropout = nn.Dropout2d(dropprob) 40 | 41 | def forward(self, input): 42 | output = self.conv3x1_1(input) 43 | output = F.relu(output) 44 | output = self.conv1x3_1(output) 45 | output = self.bn1(output) 46 | output = F.relu(output) 47 | 48 | output = self.conv3x1_2(output) 49 | output = F.relu(output) 50 | output = self.conv1x3_2(output) 51 | output = self.bn2(output) 52 | 53 | if (self.dropout.p != 0): 54 | output = self.dropout(output) 55 | 56 | # +input = identity (residual connection) 57 | return F.relu(output + input) 58 | 59 | 60 | class UpsamplerBlock(nn.Module): 61 | def __init__(self, ninput, noutput, up_width, up_height): 62 | super().__init__() 63 | 64 | self.conv = nn.ConvTranspose2d( 65 | ninput, noutput, 3, stride=2, padding=1, output_padding=1, bias=True) 66 | 67 | self.bn = nn.BatchNorm2d(noutput, eps=1e-3, track_running_stats=True) 68 | 69 | self.follows = nn.ModuleList() 70 | self.follows.append(non_bottleneck_1d(noutput, 0, 1)) 71 | self.follows.append(non_bottleneck_1d(noutput, 0, 1)) 72 | 73 | # interpolate 74 | self.up_width = up_width 75 | self.up_height = up_height 76 | self.interpolate_conv = conv1x1(ninput, noutput) 77 | self.interpolate_bn = nn.BatchNorm2d( 78 | noutput, eps=1e-3, track_running_stats=True) 79 | 80 | def forward(self, input): 81 | output = self.conv(input) 82 | output = self.bn(output) 83 | out = F.relu(output) 84 | for follow in self.follows: 85 | out = follow(out) 86 | 87 | interpolate_output = self.interpolate_conv(input) 88 | interpolate_output = self.interpolate_bn(interpolate_output) 89 | interpolate_output = F.relu(interpolate_output) 90 | 91 | interpolate = F.interpolate(interpolate_output, size=[self.up_height, self.up_width], 92 | mode='bilinear', align_corners=False) 93 | 94 | return out + interpolate 95 | 96 | @HEADS.register_module 97 | class BUSD(nn.Module): 98 | def __init__(self, cfg): 99 | super().__init__() 100 | img_height = cfg.img_height 101 | img_width = cfg.img_width 102 | num_lanes = cfg.num_lanes 103 | 104 | self.layers = nn.ModuleList() 105 | 106 | self.layers.append(UpsamplerBlock(ninput=128, noutput=64, 107 | up_height=int(img_height)//4, up_width=int(img_width)//4)) 108 | self.layers.append(UpsamplerBlock(ninput=64, noutput=32, 109 | up_height=int(img_height)//2, up_width=int(img_width)//2)) 110 | self.layers.append(UpsamplerBlock(ninput=32, noutput=16, 111 | up_height=int(img_height)//1, up_width=int(img_width)//1)) 112 | 113 | self.output_conv = conv1x1(16, num_lanes) 114 | 115 | def forward(self, input): 116 | output = input 117 | 118 | for layer in self.layers: 119 | output = layer(output) 120 | 121 | output = self.output_conv(output) 122 | output = {'seg': output} 123 | 124 | return output 125 | -------------------------------------------------------------------------------- /lanedet/datasets/process/generate_lane_cls.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import numpy as np 3 | import cv2 4 | import os 5 | import json 6 | import torchvision 7 | import inspect 8 | from ..registry import PROCESS 9 | 10 | tusimple_row_anchor = [ 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 11 | 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164, 12 | 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216, 13 | 220, 224, 228, 232, 236, 240, 244, 248, 252, 256, 260, 264, 268, 14 | 272, 276, 280, 284] 15 | 16 | culane_row_anchor = [121, 131, 141, 150, 160, 170, 180, 189, 199, 209, 219, 228, 238, 248, 258, 267, 277, 287] 17 | 18 | def find_start_pos(row_sample,start_line): 19 | l,r = 0,len(row_sample)-1 20 | while True: 21 | mid = int((l+r)/2) 22 | if r - l == 1: 23 | return r 24 | if row_sample[mid] < start_line: 25 | l = mid 26 | if row_sample[mid] > start_line: 27 | r = mid 28 | if row_sample[mid] == start_line: 29 | return mid 30 | 31 | def _grid_pts(pts, num_cols, w): 32 | # pts : numlane,n,2 33 | num_lane, n, n2 = pts.shape 34 | col_sample = np.linspace(0, w - 1, num_cols) 35 | 36 | assert n2 == 2 37 | to_pts = np.zeros((n, num_lane)) 38 | tot_len = col_sample[1] - col_sample[0] 39 | 40 | for i in range(num_lane): 41 | pti = pts[i, :, 1] 42 | to_pts[:, i] = np.asarray( 43 | [int(pt // tot_len) if pt != -1 else num_cols for pt in pti]) 44 | return to_pts.astype(int) 45 | 46 | @PROCESS.register_module 47 | class GenerateLaneCls(object): 48 | def __init__(self, row_anchor, num_cols, num_lanes, cfg): 49 | self.row_anchor = eval(row_anchor) 50 | self.num_cols = num_cols #100 51 | self.num_lanes = num_lanes #6 52 | 53 | def __call__(self, sample): 54 | label = sample['mask'] # seg_mask 55 | h, w = label.shape # 720x1280 56 | if h != 288: 57 | scale_f = lambda x : int((x * 1.0/288) * h) 58 | sample_tmp = list(map(scale_f, self.row_anchor)) # list [160, ..... 710] 59 | 60 | all_idx = np.zeros((self.num_lanes, len(sample_tmp),2)) # 6x56x2 61 | 62 | for i,r in enumerate(sample_tmp): 63 | label_r = np.asarray(label)[int(round(r))] # 1280 pixels in each row anchor: shape = 1280x1 64 | # pixels are actually lane numbers like 1 to 6 65 | for lane_idx in range(1, self.num_lanes+1): # 1 to 6 66 | pos = np.where(label_r == lane_idx)[0] # x pixels of the lane location 67 | if len(pos) == 0: 68 | all_idx[lane_idx - 1, i, 0] = r # if no lane, just put y values like 160, 170 69 | all_idx[lane_idx - 1, i, 1] = -1 # in x values, put -1 70 | continue 71 | pos = np.mean(pos) 72 | all_idx[lane_idx - 1, i, 0] = r 73 | all_idx[lane_idx - 1, i, 1] = pos # in x values, put mean of x pixels of the lane 74 | 75 | all_idx_cp = all_idx.copy() 76 | for i in range(self.num_lanes): 77 | if np.all(all_idx_cp[i,:,1] == -1): # if all x values are -1, ignore that lane 78 | continue 79 | 80 | valid = all_idx_cp[i,:,1] != -1 81 | valid_idx = all_idx_cp[i,valid,:] # index of valid lanes (y,x) 82 | if valid_idx[-1,0] == all_idx_cp[0,-1,0]: # if last y value is 710, ignore that too 83 | continue 84 | if len(valid_idx) < 6: 85 | continue 86 | 87 | valid_idx_half = valid_idx[len(valid_idx) // 2:,:] 88 | p = np.polyfit(valid_idx_half[:,0], valid_idx_half[:,1],deg = 1) # create a line using upper half of the lane 89 | start_line = valid_idx_half[-1,0] 90 | pos = find_start_pos(all_idx_cp[i,:,0],start_line) + 1 91 | 92 | fitted = np.polyval(p,all_idx_cp[i,pos:,0]) # get x values from the 1D poly using y values 93 | fitted = np.array([-1 if x< 0 or x > w-1 else x for x in fitted]) # if x value is out of bound, make it -1 94 | 95 | assert np.all(all_idx_cp[i,pos:,1] == -1) 96 | all_idx_cp[i,pos:,1] = fitted # make all x values after pos equal to fitted x values 97 | if -1 in all_idx[:, :, 0]: 98 | pdb.set_trace() 99 | sample['cls_label'] = _grid_pts(all_idx_cp, self.num_cols, w) # return int 6x56x2 array as final class label 100 | 101 | return sample 102 | -------------------------------------------------------------------------------- /lanedet/models/heads/lane_cls.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | from torch.hub import load_state_dict_from_url 5 | import numpy as np 6 | from lanedet.core.lane import Lane 7 | import scipy 8 | from lanedet.models.losses.focal_loss import SoftmaxFocalLoss 9 | import inspect 10 | from ..registry import HEADS 11 | 12 | @HEADS.register_module 13 | class LaneCls(nn.Module): 14 | def __init__(self, dim, cat_dim, cfg=None): 15 | super(LaneCls, self).__init__() 16 | self.cfg = cfg 17 | chan = cfg.featuremap_out_channel 18 | self.pool = torch.nn.Conv2d(chan, 8, 1) 19 | self.cat_dim = cat_dim 20 | self.dim = dim 21 | self.total_dim = np.prod(dim) 22 | 23 | self.det = torch.nn.Sequential( 24 | torch.nn.Linear(1800, 2048), 25 | torch.nn.ReLU(), 26 | torch.nn.Linear(2048, self.total_dim), 27 | ) 28 | 29 | if self.cfg.classification: 30 | self.category = torch.nn.Sequential( 31 | torch.nn.Linear(1800, 256), 32 | torch.nn.BatchNorm1d(256), 33 | torch.nn.ReLU(), 34 | torch.nn.Linear(256, 100), 35 | torch.nn.ReLU(), 36 | torch.nn.Linear(100, np.prod(self.cat_dim)) 37 | ) 38 | 39 | def postprocess(self, out, localization_type='rel', flip_updown=True): 40 | predictions = [] 41 | griding_num = self.cfg.griding_num 42 | for j in range(out.shape[0]): 43 | out_j = out[j].data.cpu().numpy() 44 | if flip_updown: 45 | out_j = out_j[:, ::-1, :] 46 | if localization_type == 'abs': 47 | out_j = np.argmax(out_j, axis=0) 48 | out_j[out_j == griding_num] = -1 49 | out_j = out_j + 1 50 | elif localization_type == 'rel': 51 | prob = scipy.special.softmax(out_j[:-1, :, :], axis=0) 52 | idx = np.arange(griding_num) + 1 53 | idx = idx.reshape(-1, 1, 1) 54 | loc = np.sum(prob * idx, axis=0) 55 | out_j = np.argmax(out_j, axis=0) 56 | loc[out_j == griding_num] = 0 57 | out_j = loc 58 | else: 59 | raise NotImplementedError 60 | predictions.append(out_j) 61 | return predictions 62 | 63 | def loss(self, output, batch): 64 | criterion = SoftmaxFocalLoss(2) 65 | total_loss = 0 66 | loss_stats = {} 67 | det_loss = criterion(output['det'], batch['cls_label']) 68 | classification_loss_weight = 0.7 69 | #print(batch['cls_label'].shape, batch['category'].shape) 70 | 71 | if self.cfg.classification: 72 | loss_fn = torch.nn.CrossEntropyLoss() 73 | cat_loss = loss_fn(output['category'], batch['category']) 74 | 75 | loss_stats.update({'det_loss': det_loss, 'cls_loss': cat_loss}) 76 | total_loss = det_loss + cat_loss*classification_loss_weight 77 | else: 78 | loss_stats.update({'det_loss': det_loss}) 79 | total_loss = det_loss 80 | 81 | ret = {'loss': total_loss , 'loss_stats': loss_stats} 82 | 83 | return ret 84 | 85 | def get_lanes(self, pred): 86 | predictions = self.postprocess(pred['det']) 87 | ret = {} 88 | lane_output = [] 89 | lane_indexes = [] 90 | griding_num = self.cfg.griding_num 91 | sample_y = list(self.cfg.sample_y) 92 | for out in predictions: 93 | lane_indx = [] 94 | lanes = [] 95 | for i in range(out.shape[1]): 96 | if sum(out[:, i] != 0) <= 2: continue 97 | out_i = out[:, i] 98 | lane_indx.append(i) 99 | coord = [] 100 | for k in range(out.shape[0]): 101 | if out[k, i] <= 0: continue 102 | x = ((out_i[k]-0.5) * self.cfg.ori_img_w / (griding_num - 1)) 103 | y = sample_y[k] 104 | coord.append([x, y]) 105 | coord = np.array(coord) 106 | coord = np.flip(coord, axis=0) 107 | coord[:, 0] /= self.cfg.ori_img_w 108 | coord[:, 1] /= self.cfg.ori_img_h 109 | lanes.append(Lane(coord)) 110 | lane_indexes.append(lane_indx) 111 | lane_output.append(lanes) 112 | ret.update({'lane_output': lane_output, 'lane indexes': lane_indexes}) 113 | return ret 114 | 115 | def forward(self, x, **kwargs): 116 | x = x[-1] 117 | #print(x.shape) 118 | x = self.pool(x).view(-1, 1800) # shape will be batch size x 1800 though 119 | det = self.det(x).view(-1, *self.dim) 120 | if self.cfg.classification: 121 | category = self.category(x).view(-1, *self.cat_dim) 122 | #print(category.shape) 123 | output = {'det': det, 'category': category} 124 | else: 125 | output = {'det': det} 126 | return output 127 | -------------------------------------------------------------------------------- /tools/generate_seg_tusimple.py: -------------------------------------------------------------------------------- 1 | import json 2 | import numpy as np 3 | import cv2 4 | import os 5 | import argparse 6 | 7 | TRAIN_SET = ['label_data_0313.json', 'label_data_0601.json'] 8 | VAL_SET = ['label_data_0531.json'] 9 | 10 | # create mask labels and class labels 11 | 12 | def gen_label_from_json(data_root, data_savedir, file_name): 13 | H, W = 720, 1280 14 | SEG_WIDTH = 30 15 | save_dir = data_savedir 16 | annot_file = file_name + '.json' 17 | class_file = file_name + '_classes.txt' 18 | 19 | os.makedirs(os.path.join(data_root, data_savedir, "list"), exist_ok=True) 20 | 21 | with open(os.path.join(data_root, data_savedir, 'list', annot_file), 'w') as outfile: 22 | json_path = os.path.join(data_root, annot_file) 23 | with open(json_path) as f, open(data_root + class_file) as cat_file: 24 | json_lines = f.readlines() 25 | class_lines = cat_file.readlines() 26 | line_index = 0 27 | while line_index < len(json_lines): 28 | line = json_lines[line_index] 29 | label = json.loads(line) 30 | class_line = class_lines[line_index] 31 | 32 | class_line = class_line.strip() 33 | class_list = class_line.split(' ') 34 | 35 | # ---------- clean and sort lanes ------------- 36 | lanes = [] 37 | _lanes = [] 38 | slope = [] # identify 0th, 1st, 2nd, 3rd, 4th, 5th lane through slope 39 | for i in range(len(label['lanes'])): 40 | l = [(x, y) for x, y in zip(label['lanes'][i], label['h_samples']) if x >= 0] 41 | if (len(l)>1): 42 | _lanes.append(l) 43 | slope.append(np.arctan2(l[-1][1]-l[0][1], l[0][0]-l[-1][0]) / np.pi * 180) 44 | _lanes = [_lanes[i] for i in np.argsort(slope)]# arrange lanes based on slope 45 | data = [(slp, cls) for slp, cls in zip(slope, class_list)] 46 | data.sort(key = lambda x: x[0]) # arrange (slope, class_list) based on slope 47 | slope = [slope[i] for i in np.argsort(slope)] # arrange slope low to high 48 | #print(data) 49 | #print(_lanes) 50 | 51 | idx = [None for i in range(6)] 52 | for i in range(len(slope)): 53 | if slope[i] <= 90: 54 | idx[2] = i 55 | idx[1] = i-1 if i > 0 else None 56 | idx[0] = i-2 if i > 1 else None 57 | else: 58 | idx[3] = i 59 | idx[4] = i+1 if i+1 < len(slope) else None 60 | idx[5] = i+2 if i+2 < len(slope) else None 61 | break 62 | for i in range(6): 63 | lanes.append([] if idx[i] is None else _lanes[idx[i]]) # keep max 3 on left and 3 on right 64 | 65 | # --------------------------------------------- 66 | data = [data[i] for i in idx if i is not None] 67 | 68 | img_path = label['raw_file'] 69 | seg_img = np.zeros((H, W, 3)) 70 | list_str = [] # str to be written to list.txt 71 | for i in range(len(lanes)): 72 | coords = lanes[i] 73 | if len(coords) < 4: 74 | list_str.append(0) 75 | continue 76 | for j in range(len(coords)-1): 77 | cv2.line(seg_img, coords[j], coords[j+1], (i+1, i+1, i+1), SEG_WIDTH//2) 78 | list_str.append(1) # from left 3 to right 3, put 1 if there is a lane 79 | 80 | seg_path = img_path.split("/") 81 | seg_path, img_name = os.path.join(data_root, data_savedir, seg_path[1]), seg_path[2] 82 | os.makedirs(seg_path, exist_ok=True) 83 | seg_path = os.path.join(seg_path, img_name[:-3]+"png") 84 | cv2.imwrite(seg_path, seg_img) 85 | 86 | cls = [c[1] for c in data] 87 | non_zero_ind = [i for i, e in enumerate(list_str) if e != 0] 88 | for i,val in enumerate(non_zero_ind): 89 | list_str[val]= int(cls[i]) 90 | 91 | line_index += 1 92 | label['categories']= list_str 93 | json_object = json.dumps(label) 94 | outfile.write(json_object) 95 | outfile.write('\n') 96 | 97 | def generate_label(args): 98 | save_dir = os.path.join(args.root, args.savedir) 99 | os.makedirs(save_dir, exist_ok=True) 100 | 101 | print("generating masks...") 102 | gen_label_from_json(args.root, save_dir, file_name = args.filename) 103 | 104 | 105 | if __name__ == '__main__': 106 | parser = argparse.ArgumentParser() 107 | parser.add_argument('--root', required=True, help='The root of the Tusimple dataset') 108 | parser.add_argument('--savedir', type=str, default='seg_label', help='The root of the Tusimple dataset') 109 | parser.add_argument('--filename', type=str, default='LVLane_test_sunny', help= 'Name of the json file') 110 | args = parser.parse_args() 111 | 112 | generate_label(args) 113 | -------------------------------------------------------------------------------- /configs/condlane/resnet50_culane.py: -------------------------------------------------------------------------------- 1 | net = dict( 2 | type='Detector', 3 | ) 4 | 5 | backbone = dict( 6 | type='ResNetWrapper', 7 | resnet='resnet50', 8 | pretrained=True, 9 | replace_stride_with_dilation=[False, False, False], 10 | out_conv=False, 11 | in_channels=[64, 128, 256, 512] 12 | ) 13 | 14 | sample_y = range(590, 270, -8) 15 | 16 | batch_size = 8 17 | aggregator = dict( 18 | type='TransConvEncoderModule', 19 | in_dim=2048, 20 | attn_in_dims=[2048, 256], 21 | attn_out_dims=[256, 256], 22 | strides=[1, 1], 23 | ratios=[4, 4], 24 | pos_shape=(batch_size, 10, 25), 25 | ) 26 | 27 | neck=dict( 28 | type='FPN', 29 | in_channels=[256, 512, 1024, 256], 30 | out_channels=64, 31 | num_outs=4, 32 | #trans_idx=-1, 33 | ) 34 | 35 | loss_weights=dict( 36 | hm_weight=1, 37 | kps_weight=0.4, 38 | row_weight=1., 39 | range_weight=1., 40 | ) 41 | 42 | num_lane_classes=1 43 | heads=dict( 44 | type='CondLaneHead', 45 | heads=dict(hm=num_lane_classes), 46 | in_channels=(64, ), 47 | num_classes=num_lane_classes, 48 | head_channels=64, 49 | head_layers=1, 50 | disable_coords=False, 51 | branch_in_channels=64, 52 | branch_channels=64, 53 | branch_out_channels=64, 54 | reg_branch_channels=64, 55 | branch_num_conv=1, 56 | hm_idx=2, 57 | mask_idx=0, 58 | compute_locations_pre=True, 59 | location_configs=dict(size=(batch_size, 1, 80, 200), device='cuda:0') 60 | ) 61 | 62 | optimizer = dict(type='AdamW', lr=3e-4, betas=(0.9, 0.999), eps=1e-8) 63 | 64 | epochs = 16 65 | total_iter = (88880 // batch_size) * epochs 66 | import math 67 | scheduler = dict( 68 | type = 'MultiStepLR', 69 | milestones=[8, 14], 70 | gamma=0.1 71 | ) 72 | 73 | seg_loss_weight = 1.0 74 | eval_ep = 1 75 | save_ep = 1 76 | 77 | img_norm = dict( 78 | mean=[75.3, 76.6, 77.6], 79 | std=[50.5, 53.8, 54.3] 80 | ) 81 | 82 | img_height = 320 83 | img_width = 800 84 | cut_height = 0 85 | ori_img_h = 590 86 | ori_img_w = 1640 87 | 88 | mask_down_scale = 4 89 | hm_down_scale = 16 90 | num_lane_classes = 1 91 | line_width = 3 92 | radius = 6 93 | nms_thr = 4 94 | img_scale = (800, 320) 95 | crop_bbox = [0, 270, 1640, 590] 96 | mask_size = (1, 80, 200) 97 | 98 | train_process = [ 99 | dict(type='Alaug', 100 | transforms=[dict(type='Compose', params=dict(bboxes=False, keypoints=True, masks=False)), 101 | dict( 102 | type='Crop', 103 | x_min=crop_bbox[0], 104 | x_max=crop_bbox[2], 105 | y_min=crop_bbox[1], 106 | y_max=crop_bbox[3], 107 | p=1), 108 | dict(type='Resize', height=img_scale[1], width=img_scale[0], p=1), 109 | dict( 110 | type='OneOf', 111 | transforms=[ 112 | dict( 113 | type='RGBShift', 114 | r_shift_limit=10, 115 | g_shift_limit=10, 116 | b_shift_limit=10, 117 | p=1.0), 118 | dict( 119 | type='HueSaturationValue', 120 | hue_shift_limit=(-10, 10), 121 | sat_shift_limit=(-15, 15), 122 | val_shift_limit=(-10, 10), 123 | p=1.0), 124 | ], 125 | p=0.7), 126 | dict(type='JpegCompression', quality_lower=85, quality_upper=95, p=0.2), 127 | dict( 128 | type='OneOf', 129 | transforms=[ 130 | dict(type='Blur', blur_limit=3, p=1.0), 131 | dict(type='MedianBlur', blur_limit=3, p=1.0) 132 | ], 133 | p=0.2), 134 | dict(type='RandomBrightness', limit=0.2, p=0.6), 135 | dict( 136 | type='ShiftScaleRotate', 137 | shift_limit=0.1, 138 | scale_limit=(-0.2, 0.2), 139 | rotate_limit=10, 140 | border_mode=0, 141 | p=0.6), 142 | dict( 143 | type='RandomResizedCrop', 144 | height=img_scale[1], 145 | width=img_scale[0], 146 | scale=(0.8, 1.2), 147 | ratio=(1.7, 2.7), 148 | p=0.6), 149 | dict(type='Resize', height=img_scale[1], width=img_scale[0], p=1),] 150 | ), 151 | dict(type='CollectLane', 152 | down_scale=mask_down_scale, 153 | hm_down_scale=hm_down_scale, 154 | max_mask_sample=5, 155 | line_width=line_width, 156 | radius=radius, 157 | keys=['img', 'gt_hm'], 158 | meta_keys=[ 159 | 'gt_masks', 'mask_shape', 'hm_shape', 160 | 'down_scale', 'hm_down_scale', 'gt_points' 161 | ] 162 | ), 163 | #dict(type='Resize', size=(img_width, img_height)), 164 | dict(type='Normalize', img_norm=img_norm), 165 | dict(type='ToTensor', keys=['img', 'gt_hm'], collect_keys=['img_metas']), 166 | ] 167 | 168 | 169 | val_process = [ 170 | dict(type='Alaug', 171 | transforms=[dict(type='Compose', params=dict(bboxes=False, keypoints=True, masks=False)), 172 | dict(type='Crop', 173 | x_min=crop_bbox[0], 174 | x_max=crop_bbox[2], 175 | y_min=crop_bbox[1], 176 | y_max=crop_bbox[3], 177 | p=1), 178 | dict(type='Resize', height=img_scale[1], width=img_scale[0], p=1)] 179 | ), 180 | #dict(type='Resize', size=(img_width, img_height)), 181 | dict(type='Normalize', img_norm=img_norm), 182 | dict(type='ToTensor', keys=['img']), 183 | ] 184 | 185 | dataset_path = './data/CULane' 186 | dataset = dict( 187 | train=dict( 188 | type='CULane', 189 | data_root=dataset_path, 190 | split='train', 191 | processes=train_process, 192 | ), 193 | val=dict( 194 | type='CULane', 195 | data_root=dataset_path, 196 | split='test', 197 | processes=val_process, 198 | ), 199 | test=dict( 200 | type='CULane', 201 | data_root=dataset_path, 202 | split='test', 203 | processes=val_process, 204 | ) 205 | ) 206 | 207 | 208 | workers = 12 209 | log_interval = 1000 210 | lr_update_by_epoch=True 211 | -------------------------------------------------------------------------------- /configs/condlane/resnet101_culane.py: -------------------------------------------------------------------------------- 1 | net = dict( 2 | type='Detector', 3 | ) 4 | 5 | backbone = dict( 6 | type='ResNetWrapper', 7 | resnet='resnet101', 8 | pretrained=True, 9 | replace_stride_with_dilation=[False, False, False], 10 | out_conv=False, 11 | in_channels=[64, 128, 256, 512] 12 | ) 13 | 14 | sample_y = range(590, 270, -8) 15 | 16 | batch_size = 8 17 | aggregator = dict( 18 | type='TransConvEncoderModule', 19 | in_dim=2048, 20 | attn_in_dims=[2048, 256], 21 | attn_out_dims=[256, 256], 22 | strides=[1, 1], 23 | ratios=[4, 4], 24 | pos_shape=(batch_size, 10, 25), 25 | ) 26 | 27 | neck=dict( 28 | type='FPN', 29 | in_channels=[256, 512, 1024, 256], 30 | out_channels=64, 31 | num_outs=4, 32 | #trans_idx=-1, 33 | ) 34 | 35 | loss_weights=dict( 36 | hm_weight=1, 37 | kps_weight=0.4, 38 | row_weight=1., 39 | range_weight=1., 40 | ) 41 | 42 | num_lane_classes=1 43 | heads=dict( 44 | type='CondLaneHead', 45 | heads=dict(hm=num_lane_classes), 46 | in_channels=(64, ), 47 | num_classes=num_lane_classes, 48 | head_channels=64, 49 | head_layers=1, 50 | disable_coords=False, 51 | branch_in_channels=64, 52 | branch_channels=64, 53 | branch_out_channels=64, 54 | reg_branch_channels=64, 55 | branch_num_conv=1, 56 | hm_idx=2, 57 | mask_idx=0, 58 | compute_locations_pre=True, 59 | location_configs=dict(size=(batch_size, 1, 80, 200), device='cuda:0') 60 | ) 61 | 62 | optimizer = dict(type='AdamW', lr=3e-4, betas=(0.9, 0.999), eps=1e-8) 63 | 64 | epochs = 16 65 | total_iter = (88880 // batch_size) * epochs 66 | import math 67 | scheduler = dict( 68 | type = 'MultiStepLR', 69 | milestones=[8, 14], 70 | gamma=0.1 71 | ) 72 | 73 | seg_loss_weight = 1.0 74 | eval_ep = 1 75 | save_ep = 1 76 | 77 | img_norm = dict( 78 | mean=[75.3, 76.6, 77.6], 79 | std=[50.5, 53.8, 54.3] 80 | ) 81 | 82 | img_height = 320 83 | img_width = 800 84 | cut_height = 0 85 | ori_img_h = 590 86 | ori_img_w = 1640 87 | 88 | mask_down_scale = 4 89 | hm_down_scale = 16 90 | num_lane_classes = 1 91 | line_width = 3 92 | radius = 6 93 | nms_thr = 4 94 | img_scale = (800, 320) 95 | crop_bbox = [0, 270, 1640, 590] 96 | mask_size = (1, 80, 200) 97 | 98 | train_process = [ 99 | dict(type='Alaug', 100 | transforms=[dict(type='Compose', params=dict(bboxes=False, keypoints=True, masks=False)), 101 | dict( 102 | type='Crop', 103 | x_min=crop_bbox[0], 104 | x_max=crop_bbox[2], 105 | y_min=crop_bbox[1], 106 | y_max=crop_bbox[3], 107 | p=1), 108 | dict(type='Resize', height=img_scale[1], width=img_scale[0], p=1), 109 | dict( 110 | type='OneOf', 111 | transforms=[ 112 | dict( 113 | type='RGBShift', 114 | r_shift_limit=10, 115 | g_shift_limit=10, 116 | b_shift_limit=10, 117 | p=1.0), 118 | dict( 119 | type='HueSaturationValue', 120 | hue_shift_limit=(-10, 10), 121 | sat_shift_limit=(-15, 15), 122 | val_shift_limit=(-10, 10), 123 | p=1.0), 124 | ], 125 | p=0.7), 126 | dict(type='JpegCompression', quality_lower=85, quality_upper=95, p=0.2), 127 | dict( 128 | type='OneOf', 129 | transforms=[ 130 | dict(type='Blur', blur_limit=3, p=1.0), 131 | dict(type='MedianBlur', blur_limit=3, p=1.0) 132 | ], 133 | p=0.2), 134 | dict(type='RandomBrightness', limit=0.2, p=0.6), 135 | dict( 136 | type='ShiftScaleRotate', 137 | shift_limit=0.1, 138 | scale_limit=(-0.2, 0.2), 139 | rotate_limit=10, 140 | border_mode=0, 141 | p=0.6), 142 | dict( 143 | type='RandomResizedCrop', 144 | height=img_scale[1], 145 | width=img_scale[0], 146 | scale=(0.8, 1.2), 147 | ratio=(1.7, 2.7), 148 | p=0.6), 149 | dict(type='Resize', height=img_scale[1], width=img_scale[0], p=1),] 150 | 151 | ), 152 | dict(type='CollectLane', 153 | down_scale=mask_down_scale, 154 | hm_down_scale=hm_down_scale, 155 | max_mask_sample=5, 156 | line_width=line_width, 157 | radius=radius, 158 | keys=['img', 'gt_hm'], 159 | meta_keys=[ 160 | 'gt_masks', 'mask_shape', 'hm_shape', 161 | 'down_scale', 'hm_down_scale', 'gt_points' 162 | ] 163 | ), 164 | #dict(type='Resize', size=(img_width, img_height)), 165 | dict(type='Normalize', img_norm=img_norm), 166 | dict(type='ToTensor', keys=['img', 'gt_hm'], collect_keys=['img_metas']), 167 | ] 168 | 169 | 170 | val_process = [ 171 | dict(type='Alaug', 172 | transforms=[dict(type='Compose', params=dict(bboxes=False, keypoints=True, masks=False)), 173 | dict(type='Crop', 174 | x_min=crop_bbox[0], 175 | x_max=crop_bbox[2], 176 | y_min=crop_bbox[1], 177 | y_max=crop_bbox[3], 178 | p=1), 179 | dict(type='Resize', height=img_scale[1], width=img_scale[0], p=1)] 180 | ), 181 | #dict(type='Resize', size=(img_width, img_height)), 182 | dict(type='Normalize', img_norm=img_norm), 183 | dict(type='ToTensor', keys=['img']), 184 | ] 185 | 186 | dataset_path = './data/CULane' 187 | dataset = dict( 188 | train=dict( 189 | type='CULane', 190 | data_root=dataset_path, 191 | split='train', 192 | processes=train_process, 193 | ), 194 | val=dict( 195 | type='CULane', 196 | data_root=dataset_path, 197 | split='test', 198 | processes=val_process, 199 | ), 200 | test=dict( 201 | type='CULane', 202 | data_root=dataset_path, 203 | split='test', 204 | processes=val_process, 205 | ) 206 | ) 207 | 208 | 209 | workers = 12 210 | log_interval = 1000 211 | lr_update_by_epoch=True 212 | -------------------------------------------------------------------------------- /lanedet/models/heads/lane_seg.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | from lanedet.core.lane import Lane 5 | import cv2 6 | import numpy as np 7 | 8 | from ..registry import HEADS, build_head 9 | 10 | @HEADS.register_module 11 | class LaneSeg(nn.Module): 12 | def __init__(self, decoder, exist=None, thr=0.6, 13 | sample_y=None, cat_dim = None, cfg=None, in_channels=6, out_channels=6): 14 | super(LaneSeg, self).__init__() 15 | self.cfg = cfg 16 | self.thr = thr 17 | self.sample_y = sample_y 18 | self.cat_dim = cat_dim 19 | 20 | self.decoder = build_head(decoder, cfg) 21 | self.exist = build_head(exist, cfg) if exist else None 22 | if self.cfg.classification: 23 | 24 | self.maxpool = torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 25 | self.conv1 = torch.nn.Conv2d( 26 | in_channels, 27 | out_channels, 28 | kernel_size=3, 29 | stride=1, 30 | padding=1, 31 | bias=False 32 | ) 33 | #self.bn1 = torch.nn.BatchNorm2d(out_channels) 34 | self.relu = torch.nn.ReLU(inplace=True) 35 | 36 | self.category = torch.nn.Sequential( 37 | torch.nn.Dropout(p=0.3), 38 | torch.nn.Linear(353280, 256), 39 | torch.nn.BatchNorm1d(256), 40 | torch.nn.ReLU(), 41 | torch.nn.Linear(256, 100), 42 | torch.nn.ReLU(), 43 | torch.nn.Linear(100, np.prod(self.cat_dim)) 44 | ) 45 | 46 | def get_lanes(self, output): 47 | segs = output['seg'] 48 | segs = F.softmax(segs, dim=1) 49 | segs = segs.detach().cpu().numpy() 50 | #print(segs.shape) 51 | #print("------------") 52 | if 'exist' in output: 53 | exists = output['exist'] 54 | exists = exists.detach().cpu().numpy() 55 | exists = exists > 0.5 56 | else: 57 | exists = [None for _ in segs] 58 | ret= {} 59 | 60 | lane_output = [] 61 | lane_indexes = [] 62 | for seg, exist in zip(segs, exists): 63 | #print(seg.shape) 64 | lanes, lane_indx = self.probmap2lane(seg, exist) 65 | lane_output.append(lanes) 66 | lane_indexes.append(lane_indx) 67 | ret.update({'lane_output': lane_output, 'lane indexes': lane_indexes}) 68 | return ret 69 | 70 | def probmap2lane(self, probmaps, exists=None): 71 | lanes = [] 72 | probmaps = probmaps[1:, ...] 73 | #print(probmaps.shape) 74 | if exists is None: 75 | exists = [True for _ in probmaps] 76 | 77 | lane_indx = [] 78 | for i, (probmap, exist) in enumerate(zip(probmaps, exists)): 79 | if exist == 0: 80 | continue 81 | probmap = cv2.blur(probmap, (9, 9), borderType=cv2.BORDER_REPLICATE) 82 | cut_height = self.cfg.cut_height 83 | ori_h = self.cfg.ori_img_h - cut_height 84 | coord = [] 85 | for y in self.sample_y: 86 | proj_y = round((y - cut_height) * self.cfg.img_height/ori_h) 87 | line = probmap[proj_y] 88 | if np.max(line) < self.thr: 89 | continue 90 | value = np.argmax(line) 91 | x = value*self.cfg.ori_img_w/self.cfg.img_width#-1. 92 | if x > 0: 93 | coord.append([x, y]) 94 | if len(coord) < 5: 95 | continue 96 | 97 | coord = np.array(coord) 98 | coord = np.flip(coord, axis=0) 99 | coord[:, 0] /= self.cfg.ori_img_w 100 | coord[:, 1] /= self.cfg.ori_img_h 101 | lanes.append(Lane(coord)) 102 | lane_indx.append(i) 103 | 104 | return lanes, lane_indx 105 | 106 | def loss(self, output, batch): 107 | weights = torch.ones(self.cfg.num_lanes) 108 | weights[0] = self.cfg.bg_weight 109 | weights = weights.cuda() 110 | criterion = torch.nn.NLLLoss(ignore_index=self.cfg.ignore_label, 111 | weight=weights).cuda() 112 | criterion_exist = torch.nn.BCEWithLogitsLoss().cuda() 113 | loss = 0. 114 | loss_stats = {} 115 | seg_loss = criterion(F.log_softmax( 116 | output['seg'], dim=1), batch['mask'].long()) 117 | loss += seg_loss 118 | loss_stats.update({'seg_loss': seg_loss}) 119 | 120 | if self.cfg.classification: 121 | loss_fn = torch.nn.CrossEntropyLoss() 122 | cat_loss = loss_fn(output['category'], batch['category']) 123 | loss += cat_loss*0.6 124 | loss_stats.update({'cls_loss': cat_loss}) 125 | 126 | if 'exist' in output: 127 | exist_loss = 0.1 * \ 128 | criterion_exist(output['exist'], batch['lane_exist'].float()) 129 | loss += exist_loss 130 | loss_stats.update({'exist_loss': exist_loss}) 131 | 132 | ret = {'loss': loss, 'loss_stats': loss_stats} 133 | return ret 134 | 135 | 136 | def forward(self, x, **kwargs): 137 | output = {} 138 | x = x[-1] 139 | output.update(self.decoder(x)) 140 | if self.exist: 141 | output.update(self.exist(x)) 142 | 143 | if self.cfg.classification: 144 | x= output['seg'][:,1:, ...] 145 | #print(x.shape) 146 | x = self.maxpool(x) 147 | #print(x.shape) 148 | x = self.conv1(x) 149 | #print(x.shape) 150 | #x = self.bn1(x) 151 | x = self.relu(x).view(-1, 353280) 152 | #print(x.shape) 153 | category = self.category(x).view(-1, *self.cat_dim) 154 | output.update({'category': category}) 155 | 156 | return output 157 | -------------------------------------------------------------------------------- /lanedet/models/backbones/erfnet.py: -------------------------------------------------------------------------------- 1 | # ERFNET full network definition for Pytorch 2 | # Sept 2017 3 | # Eduardo Romera 4 | ####################### 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.init as init 9 | import torch.nn.functional as F 10 | from lanedet.models.registry import BACKBONES 11 | 12 | 13 | class DownsamplerBlock(nn.Module): 14 | def __init__(self, ninput, noutput): 15 | super().__init__() 16 | 17 | self.conv = nn.Conv2d(ninput, noutput - ninput, (3, 3), stride=2, padding=1, bias=True) 18 | self.pool = nn.MaxPool2d(2, stride=2) 19 | self.bn = nn.BatchNorm2d(noutput, eps=1e-3) 20 | 21 | def forward(self, input): 22 | output = torch.cat([self.conv(input), self.pool(input)], 1) 23 | output = self.bn(output) 24 | return F.relu(output) 25 | 26 | 27 | class non_bottleneck_1d(nn.Module): 28 | def __init__(self, chann, dropprob, dilated, use_dcn=False): 29 | super().__init__() 30 | 31 | if not use_dcn: 32 | self.conv3x1_1 = nn.Conv2d(chann, chann, (3, 1), stride=1, padding=(1, 0), bias=True) 33 | 34 | self.conv1x3_1 = nn.Conv2d(chann, chann, (1, 3), stride=1, padding=(0, 1), bias=True) 35 | else: 36 | self.conv3x1_1 = DCN(chann, chann, (3, 1), stride=1, padding=(1, 0)) 37 | 38 | self.conv1x3_1 = DCN(chann, chann, (1, 3), stride=1, padding=(0, 1)) 39 | 40 | self.bn1 = nn.BatchNorm2d(chann, eps=1e-03) 41 | 42 | self.conv3x1_2 = nn.Conv2d(chann, chann, (3, 1), stride=1, padding=(1 * dilated, 0), bias=True, 43 | dilation=(dilated, 1)) 44 | 45 | self.conv1x3_2 = nn.Conv2d(chann, chann, (1, 3), stride=1, padding=(0, 1 * dilated), bias=True, 46 | dilation=(1, dilated)) 47 | 48 | self.bn2 = nn.BatchNorm2d(chann, eps=1e-03) 49 | 50 | self.dropout = nn.Dropout2d(dropprob) 51 | 52 | def forward(self, input): 53 | output = self.conv3x1_1(input) 54 | output = F.relu(output) 55 | output = self.conv1x3_1(output) 56 | output = self.bn1(output) 57 | output = F.relu(output) 58 | 59 | output = self.conv3x1_2(output) 60 | output = F.relu(output) 61 | output = self.conv1x3_2(output) 62 | output = self.bn2(output) 63 | 64 | if (self.dropout.p != 0): 65 | output = self.dropout(output) 66 | 67 | return F.relu(output + input) # +input = identity (residual connection) 68 | 69 | 70 | class Encoder(nn.Module): 71 | def __init__(self, num_classes): 72 | super().__init__() 73 | self.initial_block = DownsamplerBlock(3, 16) 74 | 75 | self.layers = nn.ModuleList() 76 | 77 | self.layers.append(DownsamplerBlock(16, 64)) 78 | 79 | for x in range(0, 5): # 5 times 80 | self.layers.append(non_bottleneck_1d(64, 0.1, 1)) 81 | 82 | self.layers.append(DownsamplerBlock(64, 128)) 83 | 84 | for x in range(0, 2): # 2 times 85 | self.layers.append(non_bottleneck_1d(128, 0.1, 2)) 86 | self.layers.append(non_bottleneck_1d(128, 0.1, 4)) 87 | self.layers.append(non_bottleneck_1d(128, 0.1, 8)) 88 | self.layers.append(non_bottleneck_1d(128, 0.1, 16)) 89 | 90 | # only for encoder mode: 91 | self.output_conv = nn.Conv2d(128, num_classes, 1, stride=1, padding=0, bias=True) 92 | 93 | def forward(self, input, predict=False): 94 | output = self.initial_block(input) 95 | 96 | for layer in self.layers: 97 | output = layer(output) 98 | 99 | if predict: 100 | output = self.output_conv(output) 101 | 102 | return output 103 | 104 | 105 | class UpsamplerBlock(nn.Module): 106 | def __init__(self, ninput, noutput): 107 | super().__init__() 108 | self.conv = nn.ConvTranspose2d(ninput, noutput, 3, stride=2, padding=1, output_padding=1, bias=True) 109 | self.bn = nn.BatchNorm2d(noutput, eps=1e-3, track_running_stats=True) 110 | 111 | def forward(self, input): 112 | output = self.conv(input) 113 | output = self.bn(output) 114 | return F.relu(output) 115 | 116 | 117 | class Lane_exist(nn.Module): 118 | def __init__(self, cfg, num_output): 119 | super().__init__() 120 | 121 | self.layers = nn.ModuleList() 122 | 123 | self.layers.append(nn.Conv2d(128, 32, (3, 3), stride=1, padding=(4, 4), bias=False, dilation=(4, 4))) 124 | self.layers.append(nn.BatchNorm2d(32, eps=1e-03)) 125 | 126 | self.layers_final = nn.ModuleList() 127 | 128 | self.layers_final.append(nn.Dropout2d(0.1)) 129 | self.layers_final.append(nn.Conv2d(32, 5, (1, 1), stride=1, padding=(0, 0), bias=True)) 130 | 131 | self.maxpool = nn.MaxPool2d(2, stride=2) 132 | self.linear_dim = int(cfg.img_width / 16 * cfg.img_height / 16 * 5) 133 | self.linear1 = nn.Linear(self.linear_dim, 128) 134 | self.linear2 = nn.Linear(128, num_output) 135 | 136 | def forward(self, input): 137 | output = input 138 | 139 | for layer in self.layers: 140 | output = layer(output) 141 | 142 | output = F.relu(output) 143 | 144 | for layer in self.layers_final: 145 | output = layer(output) 146 | 147 | output = F.softmax(output, dim=1) 148 | output = self.maxpool(output) 149 | output = output.view(-1, self.linear_dim) 150 | output = self.linear1(output) 151 | output = F.relu(output) 152 | output = self.linear2(output) 153 | output = F.sigmoid(output) 154 | 155 | return output 156 | 157 | 158 | @BACKBONES.register_module 159 | class ERFNet(nn.Module): 160 | def __init__(self, cfg): # use encoder to pass pretrained encoder 161 | super().__init__() 162 | 163 | self.encoder = Encoder(cfg.num_classes) 164 | 165 | def forward(self, input): 166 | output = self.encoder(input) # predict=False by default 167 | return [output] 168 | -------------------------------------------------------------------------------- /lanedet/datasets/process/generate_lane_line.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import numpy as np 3 | import cv2 4 | import os 5 | import json 6 | import imgaug.augmenters as iaa 7 | from imgaug.augmenters import Resize 8 | from imgaug.augmentables.lines import LineString, LineStringsOnImage 9 | from scipy.interpolate import InterpolatedUnivariateSpline 10 | 11 | from ..registry import PROCESS 12 | 13 | @PROCESS.register_module 14 | class GenerateLaneLine(object): 15 | def __init__(self, transforms=None, wh=(640, 360), cfg=None): 16 | self.transforms = transforms 17 | self.img_w, self.img_h = cfg.img_w, cfg.img_h 18 | self.num_points = cfg.num_points 19 | self.n_offsets = cfg.num_points 20 | self.n_strips = cfg.num_points - 1 21 | self.strip_size = self.img_h / self.n_strips 22 | self.max_lanes = cfg.max_lanes 23 | self.offsets_ys = np.arange(self.img_h, -1, -self.strip_size) 24 | transformations = iaa.Sequential([Resize({'height': self.img_h, 'width': self.img_w})]) 25 | if transforms is not None: 26 | transforms = [getattr(iaa, aug['name'])(**aug['parameters']) 27 | for aug in transforms] # add augmentation 28 | else: 29 | transforms = [] 30 | self.transform = iaa.Sequential([iaa.Sometimes(then_list=transforms, p=1.0), transformations]) 31 | 32 | def lane_to_linestrings(self, lanes): 33 | lines = [] 34 | for lane in lanes: 35 | lines.append(LineString(lane)) 36 | 37 | return lines 38 | 39 | def sample_lane(self, points, sample_ys): 40 | # this function expects the points to be sorted 41 | points = np.array(points) 42 | if not np.all(points[1:, 1] < points[:-1, 1]): 43 | raise Exception('Annotaion points have to be sorted') 44 | x, y = points[:, 0], points[:, 1] 45 | 46 | # interpolate points inside domain 47 | assert len(points) > 1 48 | interp = InterpolatedUnivariateSpline(y[::-1], x[::-1], k=min(3, len(points) - 1)) 49 | domain_min_y = y.min() 50 | domain_max_y = y.max() 51 | sample_ys_inside_domain = sample_ys[(sample_ys >= domain_min_y) & (sample_ys <= domain_max_y)] 52 | assert len(sample_ys_inside_domain) > 0 53 | interp_xs = interp(sample_ys_inside_domain) 54 | 55 | # extrapolate lane to the bottom of the image with a straight line using the 2 points closest to the bottom 56 | two_closest_points = points[:2] 57 | extrap = np.polyfit(two_closest_points[:, 1], two_closest_points[:, 0], deg=1) 58 | extrap_ys = sample_ys[sample_ys > domain_max_y] 59 | extrap_xs = np.polyval(extrap, extrap_ys) 60 | all_xs = np.hstack((extrap_xs, interp_xs)) 61 | 62 | # separate between inside and outside points 63 | inside_mask = (all_xs >= 0) & (all_xs < self.img_w) 64 | xs_inside_image = all_xs[inside_mask] 65 | xs_outside_image = all_xs[~inside_mask] 66 | 67 | return xs_outside_image, xs_inside_image 68 | 69 | def filter_lane(self, lane): 70 | assert lane[-1][1] <= lane[0][1] 71 | filtered_lane = [] 72 | used = set() 73 | for p in lane: 74 | if p[1] not in used: 75 | filtered_lane.append(p) 76 | used.add(p[1]) 77 | 78 | return filtered_lane 79 | 80 | def transform_annotation(self, anno, img_wh=None): 81 | img_w, img_h = self.img_w, self.img_h 82 | 83 | old_lanes = anno['lanes'] 84 | 85 | # removing lanes with less than 2 points 86 | old_lanes = filter(lambda x: len(x) > 1, old_lanes) 87 | # sort lane points by Y (bottom to top of the image) 88 | old_lanes = [sorted(lane, key=lambda x: -x[1]) for lane in old_lanes] 89 | # remove points with same Y (keep first occurrence) 90 | old_lanes = [self.filter_lane(lane) for lane in old_lanes] 91 | # normalize the annotation coordinates 92 | old_lanes = [[[x * self.img_w / float(img_w), y * self.img_h / float(img_h)] for x, y in lane] 93 | for lane in old_lanes] 94 | # create tranformed annotations 95 | lanes = np.ones((self.max_lanes, 2 + 1 + 1 + 1 + self.n_offsets), 96 | dtype=np.float32) * -1e5 # 2 scores, 1 start_y, 1 start_x, 1 length, S+1 coordinates 97 | # lanes are invalid by default 98 | lanes[:, 0] = 1 99 | lanes[:, 1] = 0 100 | for lane_idx, lane in enumerate(old_lanes): 101 | try: 102 | xs_outside_image, xs_inside_image = self.sample_lane(lane, self.offsets_ys) 103 | except AssertionError: 104 | continue 105 | if len(xs_inside_image) == 0: 106 | continue 107 | all_xs = np.hstack((xs_outside_image, xs_inside_image)) 108 | lanes[lane_idx, 0] = 0 109 | lanes[lane_idx, 1] = 1 110 | lanes[lane_idx, 2] = len(xs_outside_image) / self.n_strips 111 | lanes[lane_idx, 3] = xs_inside_image[0] 112 | lanes[lane_idx, 4] = len(xs_inside_image) 113 | lanes[lane_idx, 5:5 + len(all_xs)] = all_xs 114 | 115 | new_anno = {'label': lanes, 'old_anno': anno} 116 | return new_anno 117 | 118 | def linestrings_to_lanes(self, lines): 119 | lanes = [] 120 | for line in lines: 121 | lanes.append(line.coords) 122 | 123 | return lanes 124 | 125 | def __call__(self, sample): 126 | img_org = sample['img'] 127 | line_strings_org = self.lane_to_linestrings(sample['lanes']) 128 | line_strings_org = LineStringsOnImage(line_strings_org, shape=img_org.shape) 129 | 130 | for i in range(30): 131 | img, line_strings = self.transform(image=img_org.copy(), line_strings=line_strings_org) 132 | line_strings.clip_out_of_image_() 133 | new_anno = {'lanes': self.linestrings_to_lanes(line_strings)} 134 | try: 135 | label = self.transform_annotation(new_anno, img_wh=(self.img_w, self.img_h))['label'] 136 | break 137 | except: 138 | if (i + 1) == 30: 139 | self.logger.critical('Transform annotation failed 30 times :(') 140 | exit() 141 | 142 | sample['img'] = (img / 255.).astype(np.float32) 143 | sample['lane_line'] = label 144 | 145 | return sample 146 | -------------------------------------------------------------------------------- /lanedet/models/aggregators/transformer.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from lanedet.models.registry import AGGREGATORS 6 | 7 | from mmcv.cnn import ConvModule 8 | 9 | class PositionEmbeddingSine(nn.Module): 10 | """ 11 | This is a more standard version of the position embedding, very similar to the one 12 | used by the Attention is all you need paper, generalized to work on images. 13 | """ 14 | 15 | def __init__(self, 16 | num_pos_feats=64, 17 | temperature=10000, 18 | normalize=False, 19 | scale=None): 20 | super().__init__() 21 | self.num_pos_feats = num_pos_feats 22 | self.temperature = temperature 23 | self.normalize = normalize 24 | if scale is not None and normalize is False: 25 | raise ValueError("normalize should be True if scale is passed") 26 | if scale is None: 27 | scale = 2 * math.pi 28 | self.scale = scale 29 | 30 | def forward(self, mask): 31 | assert mask is not None 32 | not_mask = ~mask 33 | y_embed = not_mask.cumsum(1, dtype=torch.float32) 34 | x_embed = not_mask.cumsum(2, dtype=torch.float32) 35 | if self.normalize: 36 | eps = 1e-6 37 | y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale 38 | x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale 39 | 40 | dim_t = torch.arange( 41 | self.num_pos_feats, dtype=torch.float32, device=mask.device) 42 | dim_t = self.temperature**(2 * (dim_t // 2) / self.num_pos_feats) 43 | 44 | pos_x = x_embed[:, :, :, None] / dim_t 45 | pos_y = y_embed[:, :, :, None] / dim_t 46 | pos_x = torch.stack( 47 | (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), 48 | dim=4).flatten(3) 49 | pos_y = torch.stack( 50 | (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), 51 | dim=4).flatten(3) 52 | pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) 53 | return pos 54 | 55 | def build_position_encoding(hidden_dim, shape): 56 | mask = torch.zeros(shape, dtype=torch.bool) 57 | pos_module = PositionEmbeddingSine(hidden_dim // 2) 58 | pos_embs = pos_module(mask) 59 | return pos_embs 60 | 61 | 62 | class AttentionLayer(nn.Module): 63 | """ Position attention module""" 64 | 65 | def __init__(self, in_dim, out_dim, ratio=4, stride=1): 66 | super(AttentionLayer, self).__init__() 67 | self.chanel_in = in_dim 68 | norm_cfg = dict(type='BN', requires_grad=True) 69 | act_cfg = dict(type='ReLU') 70 | self.pre_conv = ConvModule( 71 | in_dim, 72 | out_dim, 73 | kernel_size=3, 74 | stride=stride, 75 | padding=1, 76 | norm_cfg=norm_cfg, 77 | act_cfg=act_cfg, 78 | inplace=False) 79 | self.query_conv = nn.Conv2d( 80 | in_channels=out_dim, out_channels=out_dim // ratio, kernel_size=1) 81 | self.key_conv = nn.Conv2d( 82 | in_channels=out_dim, out_channels=out_dim // ratio, kernel_size=1) 83 | self.value_conv = nn.Conv2d( 84 | in_channels=out_dim, out_channels=out_dim, kernel_size=1) 85 | self.final_conv = ConvModule( 86 | out_dim, 87 | out_dim, 88 | kernel_size=3, 89 | padding=1, 90 | norm_cfg=norm_cfg, 91 | act_cfg=act_cfg) 92 | self.softmax = nn.Softmax(dim=-1) 93 | self.gamma = nn.Parameter(torch.zeros(1)) 94 | 95 | def forward(self, x, pos=None): 96 | """ 97 | inputs : 98 | x : inpput feature maps( B X C X H X W) 99 | returns : 100 | out : attention value + input feature 101 | attention: B X (HxW) X (HxW) 102 | """ 103 | x = self.pre_conv(x) 104 | m_batchsize, _, height, width = x.size() 105 | if pos is not None: 106 | x += pos 107 | proj_query = self.query_conv(x).view(m_batchsize, -1, 108 | width * height).permute(0, 2, 1) 109 | proj_key = self.key_conv(x).view(m_batchsize, -1, width * height) 110 | 111 | energy = torch.bmm(proj_query, proj_key) 112 | attention = self.softmax(energy) 113 | attention = attention.permute(0, 2, 1) 114 | proj_value = self.value_conv(x).view(m_batchsize, -1, width * height) 115 | out = torch.bmm(proj_value, attention) 116 | out = out.view(m_batchsize, -1, height, width) 117 | proj_value = proj_value.view(m_batchsize, -1, height, width) 118 | out_feat = self.gamma * out + x 119 | out_feat = self.final_conv(out_feat) 120 | return out_feat 121 | 122 | @AGGREGATORS.register_module 123 | class TransConvEncoderModule(nn.Module): 124 | def __init__(self, in_dim, attn_in_dims, attn_out_dims, strides, ratios, downscale=True, pos_shape=None, cfg=None): 125 | super(TransConvEncoderModule, self).__init__() 126 | if downscale: 127 | stride = 2 128 | else: 129 | stride = 1 130 | # self.first_conv = ConvModule(in_dim, 2*in_dim, kernel_size=3, stride=stride, padding=1) 131 | # self.final_conv = ConvModule(attn_out_dims[-1], attn_out_dims[-1], kernel_size=3, stride=1, padding=1) 132 | attn_layers = [] 133 | for dim1, dim2, stride, ratio in zip(attn_in_dims, attn_out_dims, strides, ratios): 134 | attn_layers.append(AttentionLayer(dim1, dim2, ratio, stride)) 135 | if pos_shape is not None: 136 | self.attn_layers = nn.ModuleList(attn_layers) 137 | else: 138 | self.attn_layers = nn.Sequential(*attn_layers) 139 | self.pos_shape = pos_shape 140 | self.pos_embeds = [] 141 | if pos_shape is not None: 142 | for dim in attn_out_dims: 143 | pos_embed = build_position_encoding(dim, pos_shape).cuda() 144 | self.pos_embeds.append(pos_embed) 145 | 146 | def forward(self, src): 147 | # src = self.first_conv(src) 148 | if self.pos_shape is None: 149 | src = self.attn_layers(src) 150 | else: 151 | for layer, pos in zip(self.attn_layers, self.pos_embeds): 152 | src = layer(src, pos.to(src.device)) 153 | # src = self.final_conv(src) 154 | return src 155 | -------------------------------------------------------------------------------- /lanedet/datasets/tusimple.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import numpy as np 3 | import cv2 4 | import os 5 | import json 6 | import torch 7 | import torchvision 8 | from .base_dataset import BaseDataset 9 | from lanedet.utils.tusimple_metric import LaneEval 10 | from .registry import DATASETS 11 | import logging 12 | import random 13 | import torch.nn.functional as F 14 | import pandas as pd 15 | import seaborn as sns 16 | import matplotlib.pyplot as plt 17 | from sklearn.metrics import confusion_matrix 18 | 19 | SPLIT_FILES = { 20 | 21 | 'trainval': ['label_data_0313.json', 'label_data_0601.json', 'label_data_0531.json'], 22 | 'val': ['test_label.json'], 23 | 'test': ['test_label.json'] 24 | } 25 | 26 | 27 | @DATASETS.register_module 28 | class TuSimple(BaseDataset): 29 | def __init__(self, data_root, split, processes=None, cfg=None): 30 | super().__init__(data_root, split, processes, cfg) 31 | self.anno_files = SPLIT_FILES[split] 32 | self.load_annotations() 33 | self.h_samples = list(range(160, 720, 10)) 34 | 35 | def load_annotations(self): 36 | self.logger.info('Loading TuSimple annotations...') 37 | self.data_infos = [] 38 | max_lanes = 0 39 | #df = {0:0, 1:1, 2:2, 3:3, 4:3, 5:4, 6:5, 7:6} # for 6 class 40 | df = {0:0, 1:1, 2:1, 3:2, 4:2, 5:2, 6:1, 7:1} # for 2 class 41 | #df = {0:0, 1:1, 2:1, 3:2, 4:2, 5:1} # for caltech 2 class 42 | for anno_file in self.anno_files: 43 | anno_file = osp.join(self.data_root, anno_file) 44 | with open(anno_file, 'r') as anno_obj: 45 | lines = anno_obj.readlines() 46 | for line in lines: 47 | data = json.loads(line) 48 | y_samples = data['h_samples'] 49 | gt_lanes = data['lanes'] 50 | category = data['categories'] 51 | category = list(map(df.get,category)) 52 | 53 | mask_path = data['raw_file'].replace('clips', 'seg_label')[:-3] + 'png' 54 | lanes = [[(x, y) for (x, y) in zip(lane, y_samples) if x >= 0] for lane in gt_lanes] 55 | lanes = [lane for lane in lanes if len(lane) > 0] 56 | max_lanes = max(max_lanes, len(lanes)) 57 | self.data_infos.append({ 58 | 'img_path': osp.join(self.data_root, data['raw_file']), #append all the samples in all the json files 59 | 'img_name': data['raw_file'], 60 | 'mask_path': osp.join(self.data_root, mask_path), 61 | 'lanes': lanes, 62 | 'categories':category 63 | }) 64 | if self.training: 65 | random.shuffle(self.data_infos) 66 | self.max_lanes = max_lanes 67 | 68 | def pred2lanes(self, pred): 69 | ys = np.array(self.h_samples) / self.cfg.ori_img_h 70 | lanes = [] 71 | for lane in pred: 72 | xs = lane(ys) 73 | invalid_mask = xs < 0 74 | lane = (xs * self.cfg.ori_img_w).astype(int) 75 | lane[invalid_mask] = -2 76 | lanes.append(lane.tolist()) 77 | 78 | return lanes 79 | 80 | def pred2tusimpleformat(self, idx, pred, runtime): 81 | runtime *= 1000. # s to ms 82 | img_name = self.data_infos[idx]['img_name'] 83 | lanes = self.pred2lanes(pred) 84 | output = {'raw_file': img_name, 'lanes': lanes, 'run_time': runtime} 85 | return json.dumps(output) 86 | 87 | def save_tusimple_predictions(self, predictions, filename, runtimes=None): 88 | if runtimes is None: 89 | runtimes = np.ones(len(predictions)) * 1.e-3 90 | lines = [] 91 | for idx, (prediction, runtime) in enumerate(zip(predictions, runtimes)): 92 | line = self.pred2tusimpleformat(idx, prediction, runtime) 93 | lines.append(line) 94 | with open(filename, 'w') as output_file: 95 | output_file.write('\n'.join(lines)) 96 | 97 | def evaluate_detection(self, predictions, output_basedir, runtimes=None): 98 | pred_filename = os.path.join(output_basedir, 'tusimple_predictions.json') 99 | self.save_tusimple_predictions(predictions, pred_filename, runtimes) 100 | result, acc = LaneEval.bench_one_submit(pred_filename, self.cfg.test_json_file) 101 | self.logger.info(result) 102 | return acc 103 | 104 | # Calculate accuracy (a classification metric) 105 | def accuracy_fn(self, y_true, y_pred): 106 | """Calculates accuracy between truth labels and predictions. 107 | Args: 108 | y_true (torch.Tensor): Truth labels for predictions. 109 | y_pred (torch.Tensor): Predictions to be compared to predictions. 110 | Returns: 111 | [torch.float]: Accuracy value between y_true and y_pred, e.g. 78.45 112 | """ 113 | correct = torch.eq(y_true, y_pred).sum().item() 114 | acc = (correct / torch.numel(y_pred)) 115 | return acc 116 | 117 | def evaluate_classification(self, predictions, ground_truth): 118 | score = F.softmax(predictions, dim=1) 119 | y_pred = score.argmax(dim=1) 120 | return self.accuracy_fn(ground_truth, y_pred) 121 | 122 | def plot_confusion_matrix(self, y_true, y_pred): 123 | 124 | cf_matrix = confusion_matrix(y_true, y_pred) 125 | #class_names = ('background','solid-yellow', 'solid-white', 'dashed','botts\'-dots', 'double-solid-yellow','unknown') 126 | class_names = ('background', 'solid', 'dashed') 127 | # Create pandas dataframe 128 | dataframe = pd.DataFrame(cf_matrix, index=class_names, columns=class_names) 129 | total_number_of_instances = dataframe.sum(1)[1:].sum() 130 | 131 | 132 | #df = {0:0, 1:1, 2:1, 3:2, 4:2, 5:1, 6:1} 133 | #y_true = list(map(df.get,y_true)) 134 | #y_pred = list(map(df.get,y_pred)) 135 | #cf_matrix_2 = confusion_matrix(y_true, y_pred) 136 | #true_positives_2 = np.diag(cf_matrix_2)[1:].sum() 137 | #accuracy_2 = true_positives_2 / total_number_of_instances 138 | #print(f"Accuracy for 2 classes: {accuracy_2}") 139 | 140 | true_positives = np.diag(cf_matrix)[1:].sum() 141 | accuracy = true_positives / total_number_of_instances 142 | print(f"Accuracy for 2 classes: {accuracy}") 143 | 144 | # compute metrices from confusion matrix 145 | FP = cf_matrix.sum(axis=0) - np.diag(cf_matrix) 146 | FN = cf_matrix.sum(axis=1) - np.diag(cf_matrix) 147 | TP = np.diag(cf_matrix) 148 | TN = cf_matrix.sum() - (FP + FN + TP) 149 | 150 | # Overall accuracy 151 | ACC = (TP+TN)/(TP+FP+FN+TN) 152 | 153 | # plot the confusion matrix 154 | plt.figure(figsize=(8, 6)) 155 | 156 | # Create heatmap 157 | sns.heatmap(dataframe, annot=True, cbar=None,cmap="YlGnBu",fmt="d") 158 | 159 | plt.title("Confusion Matrix"), plt.tight_layout() 160 | 161 | plt.ylabel("True Class"), 162 | plt.xlabel("Predicted Class") 163 | plt.show() 164 | -------------------------------------------------------------------------------- /lanedet/models/necks/fpn.py: -------------------------------------------------------------------------------- 1 | import warnings 2 | 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | from mmcv.cnn import ConvModule 7 | 8 | from ..registry import NECKS 9 | 10 | 11 | @NECKS.register_module 12 | class FPN(nn.Module): 13 | def __init__(self, 14 | in_channels, 15 | out_channels, 16 | num_outs, 17 | start_level=0, 18 | end_level=-1, 19 | add_extra_convs=False, 20 | extra_convs_on_inputs=True, 21 | relu_before_extra_convs=False, 22 | no_norm_on_lateral=False, 23 | conv_cfg=None, 24 | norm_cfg=None, 25 | act_cfg=None, 26 | upsample_cfg=dict(mode='nearest'), 27 | init_cfg=dict( 28 | type='Xavier', layer='Conv2d', distribution='uniform'), 29 | cfg=None): 30 | super(FPN, self).__init__() 31 | assert isinstance(in_channels, list) 32 | self.in_channels = in_channels 33 | self.out_channels = out_channels 34 | self.num_ins = len(in_channels) 35 | self.num_outs = num_outs 36 | self.relu_before_extra_convs = relu_before_extra_convs 37 | self.no_norm_on_lateral = no_norm_on_lateral 38 | self.upsample_cfg = upsample_cfg.copy() 39 | 40 | if end_level == -1: 41 | self.backbone_end_level = self.num_ins 42 | assert num_outs >= self.num_ins - start_level 43 | else: 44 | # if end_level < inputs, no extra level is allowed 45 | self.backbone_end_level = end_level 46 | assert end_level <= len(in_channels) 47 | assert num_outs == end_level - start_level 48 | self.start_level = start_level 49 | self.end_level = end_level 50 | self.add_extra_convs = add_extra_convs 51 | assert isinstance(add_extra_convs, (str, bool)) 52 | if isinstance(add_extra_convs, str): 53 | # Extra_convs_source choices: 'on_input', 'on_lateral', 'on_output' 54 | assert add_extra_convs in ('on_input', 'on_lateral', 'on_output') 55 | elif add_extra_convs: # True 56 | if extra_convs_on_inputs: 57 | # TODO: deprecate `extra_convs_on_inputs` 58 | warnings.simplefilter('once') 59 | warnings.warn( 60 | '"extra_convs_on_inputs" will be deprecated in v2.9.0,' 61 | 'Please use "add_extra_convs"', DeprecationWarning) 62 | self.add_extra_convs = 'on_input' 63 | else: 64 | self.add_extra_convs = 'on_output' 65 | 66 | self.lateral_convs = nn.ModuleList() 67 | self.fpn_convs = nn.ModuleList() 68 | 69 | for i in range(self.start_level, self.backbone_end_level): 70 | l_conv = ConvModule( 71 | in_channels[i], 72 | out_channels, 73 | 1, 74 | conv_cfg=conv_cfg, 75 | norm_cfg=norm_cfg if not self.no_norm_on_lateral else None, 76 | act_cfg=act_cfg, 77 | inplace=False) 78 | fpn_conv = ConvModule( 79 | out_channels, 80 | out_channels, 81 | 3, 82 | padding=1, 83 | conv_cfg=conv_cfg, 84 | norm_cfg=norm_cfg, 85 | act_cfg=act_cfg, 86 | inplace=False) 87 | 88 | self.lateral_convs.append(l_conv) 89 | self.fpn_convs.append(fpn_conv) 90 | 91 | # add extra conv layers (e.g., RetinaNet) 92 | extra_levels = num_outs - self.backbone_end_level + self.start_level 93 | if self.add_extra_convs and extra_levels >= 1: 94 | for i in range(extra_levels): 95 | if i == 0 and self.add_extra_convs == 'on_input': 96 | in_channels = self.in_channels[self.backbone_end_level - 1] 97 | else: 98 | in_channels = out_channels 99 | extra_fpn_conv = ConvModule( 100 | in_channels, 101 | out_channels, 102 | 3, 103 | stride=2, 104 | padding=1, 105 | conv_cfg=conv_cfg, 106 | norm_cfg=norm_cfg, 107 | act_cfg=act_cfg, 108 | inplace=False) 109 | self.fpn_convs.append(extra_fpn_conv) 110 | 111 | def forward(self, inputs): 112 | """Forward function.""" 113 | assert len(inputs) >= len(self.in_channels) 114 | 115 | if len(inputs) > len(self.in_channels): 116 | for _ in range(len(inputs) - len(self.in_channels)): 117 | del inputs[0] 118 | 119 | # build laterals 120 | laterals = [ 121 | lateral_conv(inputs[i + self.start_level]) 122 | for i, lateral_conv in enumerate(self.lateral_convs) 123 | ] 124 | 125 | # build top-down path 126 | used_backbone_levels = len(laterals) 127 | for i in range(used_backbone_levels - 1, 0, -1): 128 | # In some cases, fixing `scale factor` (e.g. 2) is preferred, but 129 | # it cannot co-exist with `size` in `F.interpolate`. 130 | if 'scale_factor' in self.upsample_cfg: 131 | laterals[i - 1] += F.interpolate(laterals[i], 132 | **self.upsample_cfg) 133 | else: 134 | prev_shape = laterals[i - 1].shape[2:] 135 | laterals[i - 1] += F.interpolate( 136 | laterals[i], size=prev_shape, **self.upsample_cfg) 137 | 138 | # build outputs 139 | # part 1: from original levels 140 | outs = [ 141 | self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels) 142 | ] 143 | # part 2: add extra levels 144 | if self.num_outs > len(outs): 145 | # use max pool to get more levels on top of outputs 146 | # (e.g., Faster R-CNN, Mask R-CNN) 147 | if not self.add_extra_convs: 148 | for i in range(self.num_outs - used_backbone_levels): 149 | outs.append(F.max_pool2d(outs[-1], 1, stride=2)) 150 | # add conv layers on top of original feature maps (RetinaNet) 151 | else: 152 | if self.add_extra_convs == 'on_input': 153 | extra_source = inputs[self.backbone_end_level - 1] 154 | elif self.add_extra_convs == 'on_lateral': 155 | extra_source = laterals[-1] 156 | elif self.add_extra_convs == 'on_output': 157 | extra_source = outs[-1] 158 | else: 159 | raise NotImplementedError 160 | outs.append(self.fpn_convs[used_backbone_levels](extra_source)) 161 | for i in range(used_backbone_levels + 1, self.num_outs): 162 | if self.relu_before_extra_convs: 163 | outs.append(self.fpn_convs[i](F.relu(outs[-1]))) 164 | else: 165 | outs.append(self.fpn_convs[i](outs[-1])) 166 | return tuple(outs) 167 | -------------------------------------------------------------------------------- /lanedet/utils/culane_metric.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | from functools import partial 4 | 5 | import cv2 6 | import numpy as np 7 | from tqdm import tqdm 8 | from p_tqdm import t_map, p_map 9 | from scipy.interpolate import splprep, splev 10 | from scipy.optimize import linear_sum_assignment 11 | from shapely.geometry import LineString, Polygon 12 | 13 | 14 | def draw_lane(lane, img=None, img_shape=None, width=30): 15 | if img is None: 16 | img = np.zeros(img_shape, dtype=np.uint8) 17 | lane = lane.astype(np.int32) 18 | for p1, p2 in zip(lane[:-1], lane[1:]): 19 | cv2.line(img, tuple(p1), tuple(p2), color=(255, 255, 255), thickness=width) 20 | return img 21 | 22 | 23 | def discrete_cross_iou(xs, ys, width=30, img_shape=(590, 1640, 3)): 24 | xs = [draw_lane(lane, img_shape=img_shape, width=width) > 0 for lane in xs] 25 | ys = [draw_lane(lane, img_shape=img_shape, width=width) > 0 for lane in ys] 26 | 27 | ious = np.zeros((len(xs), len(ys))) 28 | for i, x in enumerate(xs): 29 | for j, y in enumerate(ys): 30 | ious[i, j] = (x & y).sum() / (x | y).sum() 31 | return ious 32 | 33 | 34 | def continuous_cross_iou(xs, ys, width=30, img_shape=(590, 1640, 3)): 35 | h, w, _ = img_shape 36 | image = Polygon([(0, 0), (0, h - 1), (w - 1, h - 1), (w - 1, 0)]) 37 | xs = [LineString(lane).buffer(distance=width / 2., cap_style=1, join_style=2).intersection(image) for lane in xs] 38 | ys = [LineString(lane).buffer(distance=width / 2., cap_style=1, join_style=2).intersection(image) for lane in ys] 39 | 40 | ious = np.zeros((len(xs), len(ys))) 41 | for i, x in enumerate(xs): 42 | for j, y in enumerate(ys): 43 | ious[i, j] = x.intersection(y).area / x.union(y).area 44 | 45 | return ious 46 | 47 | 48 | def interp(points, n=50): 49 | x = [x for x, _ in points] 50 | y = [y for _, y in points] 51 | tck, u = splprep([x, y], s=0, t=n, k=min(3, len(points) - 1)) 52 | 53 | u = np.linspace(0., 1., num=(len(u) - 1) * n + 1) 54 | return np.array(splev(u, tck)).T 55 | 56 | 57 | def culane_metric(pred, anno, width=30, iou_threshold=0.5, official=True, img_shape=(590, 1640, 3)): 58 | if len(pred) == 0: 59 | return 0, 0, len(anno), np.zeros(len(pred)), np.zeros(len(pred), dtype=bool) 60 | if len(anno) == 0: 61 | return 0, len(pred), 0, np.zeros(len(pred)), np.zeros(len(pred), dtype=bool) 62 | interp_pred = np.array([interp(pred_lane, n=5) for pred_lane in pred], dtype=object) # (4, 50, 2) 63 | interp_anno = np.array([interp(anno_lane, n=5) for anno_lane in anno], dtype=object) # (4, 50, 2) 64 | 65 | if official: 66 | ious = discrete_cross_iou(interp_pred, interp_anno, width=width, img_shape=img_shape) 67 | else: 68 | ious = continuous_cross_iou(interp_pred, interp_anno, width=width, img_shape=img_shape) 69 | 70 | row_ind, col_ind = linear_sum_assignment(1 - ious) 71 | tp = int((ious[row_ind, col_ind] > iou_threshold).sum()) 72 | fp = len(pred) - tp 73 | fn = len(anno) - tp 74 | pred_ious = np.zeros(len(pred)) 75 | pred_ious[row_ind] = ious[row_ind, col_ind] 76 | return tp, fp, fn, pred_ious, pred_ious > iou_threshold 77 | 78 | 79 | def load_culane_img_data(path): 80 | with open(path, 'r') as data_file: 81 | img_data = data_file.readlines() 82 | img_data = [line.split() for line in img_data] 83 | img_data = [list(map(float, lane)) for lane in img_data] 84 | img_data = [[(lane[i], lane[i + 1]) for i in range(0, len(lane), 2)] for lane in img_data] 85 | img_data = [lane for lane in img_data if len(lane) >= 2] 86 | 87 | return img_data 88 | 89 | 90 | def load_culane_data(data_dir, file_list_path): 91 | with open(file_list_path, 'r') as file_list: 92 | filepaths = [ 93 | os.path.join(data_dir, line[1 if line[0] == '/' else 0:].rstrip().replace('.jpg', '.lines.txt')) 94 | for line in file_list.readlines() 95 | ] 96 | 97 | data = [] 98 | for path in tqdm(filepaths): 99 | img_data = load_culane_img_data(path) 100 | data.append(img_data) 101 | 102 | return data 103 | 104 | 105 | def eval_predictions(pred_dir, anno_dir, list_path, width=30, official=True, sequential=False): 106 | print('List: {}'.format(list_path)) 107 | print('Loading prediction data...') 108 | predictions = load_culane_data(pred_dir, list_path) 109 | print('Loading annotation data...') 110 | annotations = load_culane_data(anno_dir, list_path) 111 | print('Calculating metric {}...'.format('sequentially' if sequential else 'in parallel')) 112 | img_shape = (590, 1640, 3) 113 | if sequential: 114 | results = t_map(partial(culane_metric, width=width, official=official, img_shape=img_shape), predictions, 115 | annotations) 116 | else: 117 | results = p_map(partial(culane_metric, width=width, official=official, img_shape=img_shape), predictions, 118 | annotations) 119 | total_tp = sum(tp for tp, _, _, _, _ in results) 120 | total_fp = sum(fp for _, fp, _, _, _ in results) 121 | total_fn = sum(fn for _, _, fn, _, _ in results) 122 | if total_tp == 0: 123 | precision = 0 124 | recall = 0 125 | f1 = 0 126 | else: 127 | precision = float(total_tp) / (total_tp + total_fp) 128 | recall = float(total_tp) / (total_tp + total_fn) 129 | f1 = 2 * precision * recall / (precision + recall) 130 | 131 | return {'TP': total_tp, 'FP': total_fp, 'FN': total_fn, 'Precision': precision, 'Recall': recall, 'F1': f1} 132 | 133 | 134 | def main(): 135 | args = parse_args() 136 | for list_path in args.list: 137 | results = eval_predictions(args.pred_dir, 138 | args.anno_dir, 139 | list_path, 140 | width=args.width, 141 | official=args.official, 142 | sequential=args.sequential) 143 | 144 | header = '=' * 20 + ' Results ({})'.format(os.path.basename(list_path)) + '=' * 20 145 | print(header) 146 | for metric, value in results.items(): 147 | if isinstance(value, float): 148 | print('{}: {:.4f}'.format(metric, value)) 149 | else: 150 | print('{}: {}'.format(metric, value)) 151 | print('=' * len(header)) 152 | 153 | 154 | def parse_args(): 155 | parser = argparse.ArgumentParser(description="Measure CULane's metric") 156 | parser.add_argument("--pred_dir", help="Path to directory containing the predicted lanes", required=True) 157 | parser.add_argument("--anno_dir", help="Path to directory containing the annotated lanes", required=True) 158 | parser.add_argument("--width", type=int, default=30, help="Width of the lane") 159 | parser.add_argument("--list", nargs='+', help="Path to txt file containing the list of files", required=True) 160 | parser.add_argument("--sequential", action='store_true', help="Run sequentially instead of in parallel") 161 | parser.add_argument("--official", action='store_true', help="Use official way to calculate the metric") 162 | 163 | return parser.parse_args() 164 | 165 | 166 | if __name__ == '__main__': 167 | main() 168 | -------------------------------------------------------------------------------- /lanedet/models/losses/focal_loss.py: -------------------------------------------------------------------------------- 1 | # pylint: disable-all 2 | from typing import Optional 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | 8 | # Source: https://github.com/kornia/kornia/blob/f4f70fefb63287f72bc80cd96df9c061b1cb60dd/kornia/losses/focal.py 9 | 10 | class SoftmaxFocalLoss(nn.Module): 11 | def __init__(self, gamma, ignore_lb=255, *args, **kwargs): 12 | super(SoftmaxFocalLoss, self).__init__() 13 | self.gamma = gamma 14 | self.nll = nn.NLLLoss(ignore_index=ignore_lb) 15 | 16 | def forward(self, logits, labels): 17 | #print(f"logit shape: {logits.shape}") 18 | #print(f"label shape: {labels.shape}") 19 | scores = F.softmax(logits, dim=1) 20 | factor = torch.pow(1.-scores, self.gamma) 21 | log_score = F.log_softmax(logits, dim=1) 22 | log_score = factor * log_score 23 | #print(f"log score shape: {log_score.shape}") 24 | #print("-------------") 25 | loss = self.nll(log_score, labels) 26 | return loss 27 | 28 | def one_hot(labels: torch.Tensor, 29 | num_classes: int, 30 | device: Optional[torch.device] = None, 31 | dtype: Optional[torch.dtype] = None, 32 | eps: Optional[float] = 1e-6) -> torch.Tensor: 33 | r"""Converts an integer label x-D tensor to a one-hot (x+1)-D tensor. 34 | 35 | Args: 36 | labels (torch.Tensor) : tensor with labels of shape :math:`(N, *)`, 37 | where N is batch size. Each value is an integer 38 | representing correct classification. 39 | num_classes (int): number of classes in labels. 40 | device (Optional[torch.device]): the desired device of returned tensor. 41 | Default: if None, uses the current device for the default tensor type 42 | (see torch.set_default_tensor_type()). device will be the CPU for CPU 43 | tensor types and the current CUDA device for CUDA tensor types. 44 | dtype (Optional[torch.dtype]): the desired data type of returned 45 | tensor. Default: if None, infers data type from values. 46 | 47 | Returns: 48 | torch.Tensor: the labels in one hot tensor of shape :math:`(N, C, *)`, 49 | 50 | Examples:: 51 | >>> labels = torch.LongTensor([[[0, 1], [2, 0]]]) 52 | >>> kornia.losses.one_hot(labels, num_classes=3) 53 | tensor([[[[1., 0.], 54 | [0., 1.]], 55 | [[0., 1.], 56 | [0., 0.]], 57 | [[0., 0.], 58 | [1., 0.]]]] 59 | """ 60 | if not torch.is_tensor(labels): 61 | raise TypeError("Input labels type is not a torch.Tensor. Got {}".format(type(labels))) 62 | if not labels.dtype == torch.int64: 63 | raise ValueError("labels must be of the same dtype torch.int64. Got: {}".format(labels.dtype)) 64 | if num_classes < 1: 65 | raise ValueError("The number of classes must be bigger than one." " Got: {}".format(num_classes)) 66 | shape = labels.shape 67 | one_hot = torch.zeros(shape[0], num_classes, *shape[1:], device=device, dtype=dtype) 68 | return one_hot.scatter_(1, labels.unsqueeze(1), 1.0) + eps 69 | 70 | 71 | def focal_loss(input: torch.Tensor, 72 | target: torch.Tensor, 73 | alpha: float, 74 | gamma: float = 2.0, 75 | reduction: str = 'none', 76 | eps: float = 1e-8) -> torch.Tensor: 77 | r"""Function that computes Focal loss. 78 | 79 | See :class:`~kornia.losses.FocalLoss` for details. 80 | """ 81 | if not torch.is_tensor(input): 82 | raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(input))) 83 | 84 | if not len(input.shape) >= 2: 85 | raise ValueError("Invalid input shape, we expect BxCx*. Got: {}".format(input.shape)) 86 | 87 | if input.size(0) != target.size(0): 88 | raise ValueError('Expected input batch_size ({}) to match target batch_size ({}).'.format( 89 | input.size(0), target.size(0))) 90 | 91 | n = input.size(0) 92 | out_size = (n, ) + input.size()[2:] 93 | if target.size()[1:] != input.size()[2:]: 94 | raise ValueError('Expected target size {}, got {}'.format(out_size, target.size())) 95 | 96 | if not input.device == target.device: 97 | raise ValueError("input and target must be in the same device. Got: {} and {}".format( 98 | input.device, target.device)) 99 | 100 | # compute softmax over the classes axis 101 | input_soft: torch.Tensor = F.softmax(input, dim=1) + eps 102 | 103 | # create the labels one hot tensor 104 | target_one_hot: torch.Tensor = one_hot(target, num_classes=input.shape[1], device=input.device, dtype=input.dtype) 105 | 106 | # compute the actual focal loss 107 | weight = torch.pow(-input_soft + 1., gamma) 108 | 109 | focal = -alpha * weight * torch.log(input_soft) 110 | loss_tmp = torch.sum(target_one_hot * focal, dim=1) 111 | 112 | if reduction == 'none': 113 | loss = loss_tmp 114 | elif reduction == 'mean': 115 | loss = torch.mean(loss_tmp) 116 | elif reduction == 'sum': 117 | loss = torch.sum(loss_tmp) 118 | else: 119 | raise NotImplementedError("Invalid reduction mode: {}".format(reduction)) 120 | return loss 121 | 122 | 123 | class FocalLoss(nn.Module): 124 | r"""Criterion that computes Focal loss. 125 | 126 | According to [1], the Focal loss is computed as follows: 127 | 128 | .. math:: 129 | 130 | \text{FL}(p_t) = -\alpha_t (1 - p_t)^{\gamma} \, \text{log}(p_t) 131 | 132 | where: 133 | - :math:`p_t` is the model's estimated probability for each class. 134 | 135 | 136 | Arguments: 137 | alpha (float): Weighting factor :math:`\alpha \in [0, 1]`. 138 | gamma (float): Focusing parameter :math:`\gamma >= 0`. 139 | reduction (str, optional): Specifies the reduction to apply to the 140 | output: ‘none’ | ‘mean’ | ‘sum’. ‘none’: no reduction will be applied, 141 | ‘mean’: the sum of the output will be divided by the number of elements 142 | in the output, ‘sum’: the output will be summed. Default: ‘none’. 143 | 144 | Shape: 145 | - Input: :math:`(N, C, *)` where C = number of classes. 146 | - Target: :math:`(N, *)` where each value is 147 | :math:`0 ≤ targets[i] ≤ C−1`. 148 | 149 | Examples: 150 | >>> N = 5 # num_classes 151 | >>> kwargs = {"alpha": 0.5, "gamma": 2.0, "reduction": 'mean'} 152 | >>> loss = kornia.losses.FocalLoss(**kwargs) 153 | >>> input = torch.randn(1, N, 3, 5, requires_grad=True) 154 | >>> target = torch.empty(1, 3, 5, dtype=torch.long).random_(N) 155 | >>> output = loss(input, target) 156 | >>> output.backward() 157 | 158 | References: 159 | [1] https://arxiv.org/abs/1708.02002 160 | """ 161 | def __init__(self, alpha: float, gamma: float = 2.0, reduction: str = 'none') -> None: 162 | super(FocalLoss, self).__init__() 163 | self.alpha: float = alpha 164 | self.gamma: float = gamma 165 | self.reduction: str = reduction 166 | self.eps: float = 1e-6 167 | 168 | def forward( # type: ignore 169 | self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: 170 | return focal_loss(input, target, self.alpha, self.gamma, self.reduction, self.eps) 171 | -------------------------------------------------------------------------------- /lanedet/datasets/process/alaug.py: -------------------------------------------------------------------------------- 1 | """ 2 | Alaug interface. 3 | """ 4 | import random 5 | import collections 6 | 7 | import albumentations as al 8 | import numpy as np 9 | 10 | from ..registry import PROCESS 11 | 12 | 13 | @PROCESS.register_module 14 | class Alaug(object): 15 | 16 | def __init__(self, transforms, cfg=None): 17 | assert isinstance(transforms, collections.abc.Sequence) 18 | # init as None 19 | self.__augmentor = None 20 | # put transforms in a list 21 | self.transforms = [] 22 | self.bbox_params = None 23 | self.keypoint_params = None 24 | 25 | for transform in transforms: 26 | if isinstance(transform, dict): 27 | if transform['type'] == 'Compose': 28 | self.get_al_params(transform['params']) 29 | else: 30 | transform = self.build_transforms(transform) 31 | if transform is not None: 32 | self.transforms.append(transform) 33 | else: 34 | raise TypeError('transform must be a dict') 35 | self.build() 36 | 37 | def get_al_params(self, compose): 38 | if compose['bboxes']: 39 | self.bbox_params = al.BboxParams( 40 | format='pascal_voc', 41 | min_area=0.0, 42 | min_visibility=0.0, 43 | label_fields=["bbox_labels"]) 44 | if compose['keypoints']: 45 | self.keypoint_params = al.KeypointParams( 46 | format='xy', remove_invisible=False) 47 | 48 | def build_transforms(self, transform): 49 | if transform['type'] == 'OneOf': 50 | transforms = transform['transforms'] 51 | choices = [] 52 | for t in transforms: 53 | parmas = { 54 | key: value 55 | for key, value in t.items() if key is not 'type' 56 | } 57 | choice = getattr(al, t['type'])(**parmas) 58 | choices.append(choice) 59 | return getattr(al, 'OneOf')(transforms=choices, p=transform['p']) 60 | 61 | parmas = { 62 | key: value 63 | for key, value in transform.items() if key is not 'type' 64 | } 65 | return getattr(al, transform['type'])(**parmas) 66 | 67 | def build(self): 68 | if len(self.transforms) == 0: 69 | return 70 | self.__augmentor = al.Compose( 71 | self.transforms, 72 | bbox_params=self.bbox_params, 73 | keypoint_params=self.keypoint_params, 74 | ) 75 | 76 | def cal_sum_list(self, itmes, index): 77 | sum = 0 78 | for i in range(index): 79 | sum += itmes[i] 80 | return sum 81 | 82 | def __call__(self, data): 83 | if self.__augmentor is None: 84 | return data 85 | img = data['img'] 86 | bboxes = None 87 | keypoints = None 88 | masks = None 89 | if 'gt_bboxes' in data: 90 | gt_bboxes = data['gt_bboxes'] 91 | bboxes = [] 92 | bbox_labels = [] 93 | for i in range(np.shape(gt_bboxes)[0]): 94 | if (gt_bboxes[i, 0] == gt_bboxes[i, 2]) | ( 95 | gt_bboxes[i, 1] == gt_bboxes[i, 3]): 96 | pass 97 | else: 98 | b = gt_bboxes[i, :] 99 | b = np.concatenate((b, [i])) 100 | bboxes.append(b) 101 | bbox_labels.append(data['gt_labels'][i]) 102 | else: 103 | bboxes = None 104 | bbox_labels = None 105 | if 'gt_masks' in data: 106 | masks = data['gt_masks'] 107 | else: 108 | masks = None 109 | 110 | if 'mask' in data: 111 | masks = data['mask'] 112 | 113 | if 'gt_keypoints' in data: 114 | keypoints = data["gt_keypoints"] 115 | kp_group_num = len(keypoints) 116 | # run aug 117 | keypoints_index = [] 118 | for k in keypoints: 119 | keypoints_index.append(int(len(k) / 2)) 120 | keypoints_val = [] 121 | for kps in keypoints: 122 | num = int(len(kps) / 2) 123 | for i in range(num): 124 | keypoints_val.append(kps[2 * i:2 * i + 2]) 125 | num_keypoints = len(kps) // 2 126 | else: 127 | keypoints_val = None 128 | 129 | if 'gt_points' in data: 130 | points = data["gt_points"] 131 | p_group_num = len(points) 132 | # run aug 133 | points_index = [] 134 | for k in points: 135 | points_index.append(int(len(k) / 2)) 136 | points_val = [] 137 | for pts in points: 138 | num = int(len(pts) / 2) 139 | for i in range(num): 140 | points_val.append(pts[2 * i:2 * i + 2]) 141 | num_keypoints = len(points_val) // 2 142 | if keypoints_val is None: 143 | keypoints_val = points_val 144 | else: 145 | keypoints_val = keypoints_val + points_val 146 | 147 | if 'lanes' in data: 148 | points_val = [] 149 | for lane in data['lanes']: 150 | points_val.extend(lane) 151 | 152 | points_index = [len(lane) for lane in data['lanes']] 153 | keypoints_val = points_val 154 | 155 | aug = self.__augmentor( 156 | image=img, 157 | keypoints=keypoints_val, 158 | bboxes=bboxes, 159 | mask=masks, 160 | bbox_labels=bbox_labels) 161 | 162 | data['img'] = aug['image'] 163 | data['img_shape'] = data['img'].shape 164 | if 'gt_bboxes' in data: 165 | if aug['bboxes']: 166 | data['gt_bboxes'] = np.array(aug['bboxes'])[:, :4] 167 | data['gt_labels'] = np.array(aug['bbox_labels']) 168 | else: 169 | return None 170 | if 'mask' in data: 171 | data['mask'] = np.array(aug['mask']) 172 | 173 | if 'gt_masks' in data: 174 | data['gt_masks'] = [np.array(aug['mask'])] 175 | 176 | if 'gt_points' in data or 'lanes' in data: 177 | start_idx = num_keypoints if 'gt_keypoints' in data else 0 178 | points = aug['keypoints'][start_idx:] 179 | kp_list = [[0 for j in range(i * 2)] for i in points_index] 180 | for i in range(len(points_index)): 181 | for j in range(points_index[i]): 182 | kp_list[i][2 * 183 | j] = points[self.cal_sum_list(points_index, i) + 184 | j][0] 185 | kp_list[i][2 * j + 186 | 1] = points[self.cal_sum_list(points_index, i) + 187 | j][1] 188 | data['gt_points'] = kp_list 189 | 190 | if 'gt_bboxes' in data and kp_group_num == 0: 191 | return None 192 | return data 193 | 194 | def __repr__(self): 195 | format_string = self.__class__.__name__ + '(' 196 | for t in self.transforms: 197 | format_string += '\n' 198 | format_string += ' {0}'.format(t) 199 | format_string += '\n)' 200 | return format_string 201 | 202 | -------------------------------------------------------------------------------- /lanedet/ops/csrc/nms_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | // Hard-coded maximum. Increase if needed. 10 | #define MAX_COL_BLOCKS 1000 11 | #define STRIDE 4 12 | #define N_OFFSETS 72 // if you use more than 73 offsets you will have to adjust this value 13 | #define N_STRIPS (N_OFFSETS - 1) 14 | #define PROP_SIZE (5 + N_OFFSETS) 15 | #define DATASET_OFFSET 0 16 | 17 | #define DIVUP(m,n) (((m)+(n)-1) / (n)) 18 | int64_t const threadsPerBlock = sizeof(unsigned long long) * 8; 19 | 20 | // The functions below originates from Fast R-CNN 21 | // See https://github.com/rbgirshick/py-faster-rcnn 22 | // Copyright (c) 2015 Microsoft 23 | // Licensed under The MIT License 24 | // Written by Shaoqing Ren 25 | 26 | template 27 | // __device__ inline scalar_t devIoU(scalar_t const * const a, scalar_t const * const b) { 28 | __device__ inline bool devIoU(scalar_t const * const a, scalar_t const * const b, const float threshold) { 29 | const int start_a = (int) (a[2] * N_STRIPS - DATASET_OFFSET + 0.5); // 0.5 rounding trick 30 | const int start_b = (int) (b[2] * N_STRIPS - DATASET_OFFSET + 0.5); 31 | const int start = max(start_a, start_b); 32 | const int end_a = start_a + a[4] - 1 + 0.5 - ((a[4] - 1) < 0); // - (x<0) trick to adjust for negative numbers (in case length is 0) 33 | const int end_b = start_b + b[4] - 1 + 0.5 - ((b[4] - 1) < 0); 34 | const int end = min(min(end_a, end_b), N_OFFSETS - 1); 35 | // if (end < start) return 1e9; 36 | if (end < start) return false; 37 | scalar_t dist = 0; 38 | for(unsigned char i = 5 + start; i <= 5 + end; ++i) { 39 | if (a[i] < b[i]) { 40 | dist += b[i] - a[i]; 41 | } else { 42 | dist += a[i] - b[i]; 43 | } 44 | } 45 | // return (dist / (end - start + 1)) < threshold; 46 | return dist < (threshold * (end - start + 1)); 47 | // return dist / (end - start + 1); 48 | } 49 | 50 | template 51 | __global__ void nms_kernel(const int64_t n_boxes, const scalar_t nms_overlap_thresh, 52 | const scalar_t *dev_boxes, const int64_t *idx, int64_t *dev_mask) { 53 | const int64_t row_start = blockIdx.y; 54 | const int64_t col_start = blockIdx.x; 55 | 56 | if (row_start > col_start) return; 57 | 58 | const int row_size = 59 | min(n_boxes - row_start * threadsPerBlock, threadsPerBlock); 60 | const int col_size = 61 | min(n_boxes - col_start * threadsPerBlock, threadsPerBlock); 62 | 63 | __shared__ scalar_t block_boxes[threadsPerBlock * PROP_SIZE]; 64 | if (threadIdx.x < col_size) { 65 | for (int i = 0; i < PROP_SIZE; ++i) { 66 | block_boxes[threadIdx.x * PROP_SIZE + i] = dev_boxes[idx[(threadsPerBlock * col_start + threadIdx.x)] * PROP_SIZE + i]; 67 | } 68 | // block_boxes[threadIdx.x * 4 + 0] = 69 | // dev_boxes[idx[(threadsPerBlock * col_start + threadIdx.x)] * 4 + 0]; 70 | // block_boxes[threadIdx.x * 4 + 1] = 71 | // dev_boxes[idx[(threadsPerBlock * col_start + threadIdx.x)] * 4 + 1]; 72 | // block_boxes[threadIdx.x * 4 + 2] = 73 | // dev_boxes[idx[(threadsPerBlock * col_start + threadIdx.x)] * 4 + 2]; 74 | // block_boxes[threadIdx.x * 4 + 3] = 75 | // dev_boxes[idx[(threadsPerBlock * col_start + threadIdx.x)] * 4 + 3]; 76 | } 77 | __syncthreads(); 78 | 79 | if (threadIdx.x < row_size) { 80 | const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x; 81 | const scalar_t *cur_box = dev_boxes + idx[cur_box_idx] * PROP_SIZE; 82 | int i = 0; 83 | unsigned long long t = 0; 84 | int start = 0; 85 | if (row_start == col_start) { 86 | start = threadIdx.x + 1; 87 | } 88 | for (i = start; i < col_size; i++) { 89 | if (devIoU(cur_box, block_boxes + i * PROP_SIZE, nms_overlap_thresh)) { 90 | t |= 1ULL << i; 91 | } 92 | } 93 | const int col_blocks = DIVUP(n_boxes, threadsPerBlock); 94 | dev_mask[cur_box_idx * col_blocks + col_start] = t; 95 | } 96 | } 97 | 98 | 99 | __global__ void nms_collect(const int64_t boxes_num, const int64_t col_blocks, int64_t top_k, const int64_t *idx, const int64_t *mask, int64_t *keep, int64_t *parent_object_index, int64_t *num_to_keep) { 100 | int64_t remv[MAX_COL_BLOCKS]; 101 | int64_t num_to_keep_ = 0; 102 | 103 | for (int i = 0; i < col_blocks; i++) { 104 | remv[i] = 0; 105 | } 106 | 107 | for (int i = 0; i < boxes_num; ++i) { 108 | parent_object_index[i] = 0; 109 | } 110 | 111 | for (int i = 0; i < boxes_num; i++) { 112 | int nblock = i / threadsPerBlock; 113 | int inblock = i % threadsPerBlock; 114 | 115 | 116 | if (!(remv[nblock] & (1ULL << inblock))) { 117 | int64_t idxi = idx[i]; 118 | keep[num_to_keep_] = idxi; 119 | const int64_t *p = &mask[0] + i * col_blocks; 120 | for (int j = nblock; j < col_blocks; j++) { 121 | remv[j] |= p[j]; 122 | } 123 | for (int j = i; j < boxes_num; j++) { 124 | int nblockj = j / threadsPerBlock; 125 | int inblockj = j % threadsPerBlock; 126 | if (p[nblockj] & (1ULL << inblockj)) 127 | parent_object_index[idx[j]] = num_to_keep_+1; 128 | } 129 | parent_object_index[idx[i]] = num_to_keep_+1; 130 | 131 | num_to_keep_++; 132 | 133 | if (num_to_keep_==top_k) 134 | break; 135 | } 136 | } 137 | 138 | // Initialize the rest of the keep array to avoid uninitialized values. 139 | for (int i = num_to_keep_; i < boxes_num; ++i) 140 | keep[i] = 0; 141 | 142 | *num_to_keep = min(top_k,num_to_keep_); 143 | } 144 | 145 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 146 | 147 | std::vector nms_cuda_forward( 148 | at::Tensor boxes, 149 | at::Tensor idx, 150 | float nms_overlap_thresh, 151 | unsigned long top_k) { 152 | 153 | const auto boxes_num = boxes.size(0); 154 | TORCH_CHECK(boxes.size(1) == PROP_SIZE, "Wrong number of offsets. Please adjust `PROP_SIZE`"); 155 | 156 | const int col_blocks = DIVUP(boxes_num, threadsPerBlock); 157 | 158 | AT_ASSERTM (col_blocks < MAX_COL_BLOCKS, "The number of column blocks must be less than MAX_COL_BLOCKS. Increase the MAX_COL_BLOCKS constant if needed."); 159 | 160 | auto longOptions = torch::TensorOptions().device(torch::kCUDA).dtype(torch::kLong); 161 | auto mask = at::empty({boxes_num * col_blocks}, longOptions); 162 | 163 | dim3 blocks(DIVUP(boxes_num, threadsPerBlock), 164 | DIVUP(boxes_num, threadsPerBlock)); 165 | dim3 threads(threadsPerBlock); 166 | 167 | CHECK_CONTIGUOUS(boxes); 168 | CHECK_CONTIGUOUS(idx); 169 | CHECK_CONTIGUOUS(mask); 170 | 171 | AT_DISPATCH_FLOATING_TYPES(boxes.type(), "nms_cuda_forward", ([&] { 172 | nms_kernel<<>>(boxes_num, 173 | (scalar_t)nms_overlap_thresh, 174 | boxes.data(), 175 | idx.data(), 176 | mask.data()); 177 | })); 178 | 179 | auto keep = at::empty({boxes_num}, longOptions); 180 | auto parent_object_index = at::empty({boxes_num}, longOptions); 181 | auto num_to_keep = at::empty({}, longOptions); 182 | 183 | nms_collect<<<1, 1>>>(boxes_num, col_blocks, top_k, 184 | idx.data(), 185 | mask.data(), 186 | keep.data(), 187 | parent_object_index.data(), 188 | num_to_keep.data()); 189 | 190 | 191 | return {keep,num_to_keep,parent_object_index}; 192 | } 193 | 194 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # LVLane 2 | ## Introduction 3 | This repository is the official implementation of the paper "[LVLane: Lane Detection and Classification in Challenging Conditions](https://arxiv.org/abs/2307.06853)", accpeted in 2023 IEEE International Conference on Intelligent Trabsportation Systems (ITSC). 4 | 5 | ![demo image](.github/test-class-lvlane-ufld2.jpg) 6 | 7 | ## Table of Contents 8 | * [Introduction](#Introduction) 9 | * [Benchmark and model zoo](#Benchmark-and-model-zoo) 10 | * [Installation](#Installation) 11 | * [Getting Started](#Getting-started) 12 | * [Contributing](#Contributing) 13 | * [Licenses](#Licenses) 14 | * [Acknowledgement](#Acknowledgement) 15 | 16 | ## Benchmark and model zoo 17 | Supported backbones: 18 | - [x] ResNet 19 | - [x] ERFNet 20 | - [x] VGG 21 | - [x] MobileNet 22 | 23 | Supported detectors: 24 | - [x] [UFLD](configs/ufld) 25 | - [x] [RESA](configs/resa) 26 | 27 | 28 | ## Installation 29 | This repository is a modified version of [lanedet](https://github.com/Turoad/lanedet.git); so, it you installed that, no need to install this one. Just clone this and use the same conda environment. 30 | 33 | 34 | ### Clone this repository 35 | ``` 36 | git clone https://github.com/zillur-av/LVLane.git 37 | ``` 38 | We call this directory as `$LANEDET_ROOT` 39 | 40 | ### Create a conda virtual environment and activate it (conda is optional) 41 | 42 | ```Shell 43 | conda create -n lanedet python=3.8 -y 44 | conda activate lanedet 45 | ``` 46 | 47 | ### Install dependencies 48 | 49 | ```Shell 50 | # Install pytorch firstly, the cudatoolkit version should be same in your system. 51 | 52 | conda install pytorch==1.8.0 torchvision==0.9.0 cudatoolkit=10.1 -c pytorch 53 | 54 | # Or you can install via pip 55 | pip install torch==1.8.0 torchvision==0.9.0 56 | 57 | # Install python packages 58 | python setup.py build develop 59 | ``` 60 | 61 | ## Data preparation 62 | 63 | ### Tusimple 64 | Download [Tusimple](https://github.com/TuSimple/tusimple-benchmark/issues/3). Then extract them to `$DATASETROOT`. Create link to `data` directory. 65 | 66 | ```Shell 67 | cd $LANEDET_ROOT 68 | mkdir -p data 69 | ln -s $DATASETROOT data/tusimple 70 | ``` 71 | 72 | For Tusimple, you should have structure like this: 73 | ``` 74 | $DATASETROOT/clips # data folders 75 | $DATASETROOT/lable_data_xxxx.json # label json file 76 | $DATASETROOT/test_label.json # test label json file 77 | 78 | ``` 79 | ### LVLane 80 | Download [LVLaneV1](https://drive.google.com/file/d/1lRhne-d87A4b0gLjf6quipDQ4MYvP7ky/view?usp=sharing). Then extract them to `$DATASETROOT` just like TuSimple dataset. This link contains class annotations for TuSimple dataset, so replace the orginal labels ones with the new ones. Lane annotations and class labels of Caltech dataset are also available in TuSimple format. Download the dataset from original site and resize them to 1280x720 to use with this model. 81 | 82 | ``` 83 | $DATASETROOT/clips/0531/ 84 | . 85 | . 86 | $DATASETROOT/clips/LVLane_train_sunny/ 87 | $DATASETROOT/label_data_xxxx.json 88 | $DATASETROOT/test_label.json 89 | $DATASETROOT/LVLane_test_sunny.json 90 | $DATASETROOT/LVLane_train_sunny.json 91 | 92 | ``` 93 | If you want to create a dataset in tusimple format, please follow instructions on [tusimple-annotation](https://github.com/zillur-av/tusimple-annotation) 94 | We need to generate segmentation from the json annotation. 95 | ### Generate masks 96 | ```Shell 97 | python tools/generate_seg_tusimple.py --root $DATASETROOT --filename 'LVLane_test_sunny' 98 | # this will generate seg_label directory 99 | ``` 100 | Then you will find new `json` annotations files that have both lane location and class id in `$DATASETROOT/seg_label/list/`. Replace the old annotation files in `$DATASETROOT` by these new files. 101 | 102 | ## Getting Started 103 | If we want just detection, no lane classification, switch to `detection` branch by running `git checkout detection`. 104 | ### Training 105 | 106 | For training, run 107 | 108 | ```Shell 109 | python main.py [configs/path_to_your_config] --gpus [gpu_ids] 110 | ``` 111 | 112 | 113 | For example, run 114 | ```Shell 115 | python main.py configs/ufld/resnet18_tusimple.py --gpus 0 116 | ``` 117 | Modifications before you run training script: 118 | * Check image resolution in here https://github.com/zillur-av/LVLane/blob/f89d53d63b45069fdae6689157c7f33caa6c8652/configs/ufld/resnet18_tusimple.py#L56-L61 119 | If your images have different resolution, try to resize them to 1280x720. Modify the annotations proportionately as well. This will be the best way to handle that situation. 120 | * Modify batch size, number of training samples, epochs in https://github.com/zillur-av/LVLane/blob/f89d53d63b45069fdae6689157c7f33caa6c8652/configs/ufld/resnet18_tusimple.py#L40-L43 121 | 122 | 123 | 124 | ### Testing 125 | For testing, run 126 | ```Shell 127 | python main.py [configs/path_to_your_config] --test --load_from [path_to_your_model] [gpu_num] 128 | ``` 129 | 130 | For example, run 131 | ```Shell 132 | python main.py configs/ufld/resnet18_tusimple.py --test --load_from ufld_tusimple.pth --gpus 0 133 | ``` 134 | 135 | Currently, this code can output the visualization result when testing, just add `--view`. 136 | We will get the visualization result in `work_dirs/xxx/xxx/visualization`. 137 | 138 | I am providing a sample weights for quick testing. You can download it from [here](https://drive.google.com/file/d/1YYWE-KiihE2c4BtYHeR2BuzXQkwjk0J3/view?usp=sharing) and put it on `$LANEDET_ROOT`. If you want to test your own images, create the json file and image folders following above instructions. Then edit `val` and `test` in https://github.com/zillur-av/LVLane/blob/943dbd3ac043bcee64c061b2db8e55e802bfc07f/lanedet/datasets/tusimple.py#L21 and in configs file https://github.com/zillur-av/LVLane/blob/943dbd3ac043bcee64c061b2db8e55e802bfc07f/configs/ufld/resnet18_tusimple.py#L120. 139 | 140 | For example, run 141 | ```Shell 142 | python main.py configs/ufld/resnet18_tusimple.py --test --load_from best-ufld.pth --gpus 0 --view 143 | ``` 144 | 145 | ### Inference 146 | See `tools/detect.py` for detailed information. 147 | ``` 148 | python tools/detect.py --help 149 | 150 | usage: detect.py [-h] [--img IMG] [--show] [--savedir SAVEDIR] 151 | [--load_from LOAD_FROM] 152 | config 153 | 154 | positional arguments: 155 | config The path of config file 156 | 157 | optional arguments: 158 | -h, --help show this help message and exit 159 | --img IMG The path of the img (img file or img_folder), for 160 | example: data/*.png 161 | --show Whether to show the image 162 | --savedir SAVEDIR The root of save directory 163 | --load_from LOAD_FROM 164 | The path of model 165 | ``` 166 | To run inference on example images in `./images` and save the visualization images in `vis` folder: 167 | ``` 168 | python tools/detect.py configs/ufld/resnet18_tusimple.py --img images\ 169 | --load_from best-ufld.pth --savedir ./show 170 | ``` 171 | 172 | 173 | ## Contributing 174 | We appreciate all contributions to improve LVLane. Any pull requests or issues are welcomed. 175 | 176 | ## Licenses 177 | This project is released under the [Apache 2.0 license](LICNESE). 178 | 179 | 180 | ## Acknowledgement 181 | 182 | * [Turoad/lanedet](https://github.com/Turoad/lanedet) 183 | * [pytorch/vision](https://github.com/pytorch/vision) 184 | * [ZJULearning/resa](https://github.com/ZJULearning/resa) 185 | * [cfzd/Ultra-Fast-Lane-Detection](https://github.com/cfzd/Ultra-Fast-Lane-Detection) 186 | 187 | 188 | 189 | ## Citation 190 | If you use our work or dataset, please cite the following paper: 191 | ``` 192 | @article{rahman2023lvlane, 193 | title={LVLane: Deep Learning for Lane Detection and Classification in Challenging Conditions}, 194 | author={Rahman, Zillur and Morris, Brendan Tran}, 195 | journal={2023 IEEE International Conference on Intelligent Trabsportation Systems (ITSC)}, 196 | year={2023} 197 | } 198 | 199 | ``` 200 | -------------------------------------------------------------------------------- /lanedet/models/backbones/mobilenet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch import Tensor 4 | from torchvision.models.utils import load_state_dict_from_url 5 | from typing import Callable, Any, Optional, List 6 | 7 | from lanedet.models.registry import BACKBONES 8 | 9 | 10 | model_urls = { 11 | 'MobileNetV2': 'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth', 12 | } 13 | 14 | 15 | def _make_divisible(v: float, divisor: int, min_value: Optional[int] = None) -> int: 16 | """ 17 | This function is taken from the original tf repo. 18 | It ensures that all layers have a channel number that is divisible by 8 19 | It can be seen here: 20 | https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py 21 | """ 22 | if min_value is None: 23 | min_value = divisor 24 | new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) 25 | # Make sure that round down does not go down by more than 10%. 26 | if new_v < 0.9 * v: 27 | new_v += divisor 28 | return new_v 29 | 30 | 31 | class ConvBNActivation(nn.Sequential): 32 | def __init__( 33 | self, 34 | in_planes: int, 35 | out_planes: int, 36 | kernel_size: int = 3, 37 | stride: int = 1, 38 | groups: int = 1, 39 | norm_layer: Optional[Callable[..., nn.Module]] = None, 40 | activation_layer: Optional[Callable[..., nn.Module]] = None, 41 | dilation: int = 1, 42 | ) -> None: 43 | padding = (kernel_size - 1) // 2 * dilation 44 | if norm_layer is None: 45 | norm_layer = nn.BatchNorm2d 46 | if activation_layer is None: 47 | activation_layer = nn.ReLU6 48 | super().__init__( 49 | nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, dilation=dilation, groups=groups, 50 | bias=False), 51 | norm_layer(out_planes), 52 | activation_layer(inplace=True) 53 | ) 54 | self.out_channels = out_planes 55 | 56 | 57 | # necessary for backwards compatibility 58 | ConvBNReLU = ConvBNActivation 59 | 60 | 61 | class InvertedResidual(nn.Module): 62 | def __init__( 63 | self, 64 | inp: int, 65 | oup: int, 66 | stride: int, 67 | expand_ratio: int, 68 | norm_layer: Optional[Callable[..., nn.Module]] = None 69 | ) -> None: 70 | super(InvertedResidual, self).__init__() 71 | self.stride = stride 72 | assert stride in [1, 2] 73 | 74 | if norm_layer is None: 75 | norm_layer = nn.BatchNorm2d 76 | 77 | hidden_dim = int(round(inp * expand_ratio)) 78 | self.use_res_connect = self.stride == 1 and inp == oup 79 | 80 | layers: List[nn.Module] = [] 81 | if expand_ratio != 1: 82 | # pw 83 | layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1, norm_layer=norm_layer)) 84 | layers.extend([ 85 | # dw 86 | ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim, norm_layer=norm_layer), 87 | # pw-linear 88 | nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), 89 | norm_layer(oup), 90 | ]) 91 | self.conv = nn.Sequential(*layers) 92 | self.out_channels = oup 93 | self._is_cn = stride > 1 94 | 95 | def forward(self, x: Tensor) -> Tensor: 96 | if self.use_res_connect: 97 | return x + self.conv(x) 98 | else: 99 | return self.conv(x) 100 | 101 | 102 | class MobileNetV2(nn.Module): 103 | def __init__( 104 | self, 105 | width_mult: float = 1.0, 106 | inverted_residual_setting: Optional[List[List[int]]] = None, 107 | round_nearest: int = 8, 108 | block: Optional[Callable[..., nn.Module]] = None, 109 | norm_layer: Optional[Callable[..., nn.Module]] = None 110 | ) -> None: 111 | """ 112 | MobileNet V2 main class 113 | Args: 114 | num_classes (int): Number of classes 115 | width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount 116 | inverted_residual_setting: Network structure 117 | round_nearest (int): Round the number of channels in each layer to be a multiple of this number 118 | Set to 1 to turn off rounding 119 | block: Module specifying inverted residual building block for mobilenet 120 | norm_layer: Module specifying the normalization layer to use 121 | """ 122 | super(MobileNetV2, self).__init__() 123 | 124 | if block is None: 125 | block = InvertedResidual 126 | 127 | if norm_layer is None: 128 | norm_layer = nn.BatchNorm2d 129 | 130 | input_channel = 32 131 | last_channel = 1280 132 | 133 | if inverted_residual_setting is None: 134 | inverted_residual_setting = [ 135 | # t, c, n, s 136 | [1, 16, 1, 1], 137 | [6, 24, 2, 2], 138 | [6, 32, 3, 2], 139 | [6, 64, 4, 2], 140 | [6, 96, 3, 1], 141 | [6, 160, 3, 2], 142 | [6, 320, 1, 1], 143 | ] 144 | 145 | # only check the first element, assuming user knows t,c,n,s are required 146 | if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4: 147 | raise ValueError("inverted_residual_setting should be non-empty " 148 | "or a 4-element list, got {}".format(inverted_residual_setting)) 149 | 150 | # building first layer 151 | input_channel = _make_divisible(input_channel * width_mult, round_nearest) 152 | self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest) 153 | features: List[nn.Module] = [ConvBNReLU(3, input_channel, stride=2, norm_layer=norm_layer)] 154 | # building inverted residual blocks 155 | for t, c, n, s in inverted_residual_setting: 156 | output_channel = _make_divisible(c * width_mult, round_nearest) 157 | for i in range(n): 158 | stride = s if i == 0 else 1 159 | features.append(block(input_channel, output_channel, stride, expand_ratio=t, norm_layer=norm_layer)) 160 | input_channel = output_channel 161 | # building last several layers 162 | features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1, norm_layer=norm_layer)) 163 | # make it nn.Sequential 164 | self.features = nn.Sequential(*features) 165 | 166 | # weight initialization 167 | for m in self.modules(): 168 | if isinstance(m, nn.Conv2d): 169 | nn.init.kaiming_normal_(m.weight, mode='fan_out') 170 | if m.bias is not None: 171 | nn.init.zeros_(m.bias) 172 | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): 173 | nn.init.ones_(m.weight) 174 | nn.init.zeros_(m.bias) 175 | elif isinstance(m, nn.Linear): 176 | nn.init.normal_(m.weight, 0, 0.01) 177 | nn.init.zeros_(m.bias) 178 | 179 | def forward(self, x: Tensor) -> Tensor: 180 | # This exists since TorchScript doesn't support inheritance, so the superclass method 181 | # (this one) needs to have a name other than `forward` that can be accessed in a subclass 182 | x = self.features(x) 183 | return x 184 | 185 | @BACKBONES.register_module 186 | class MobileNet(nn.Module): 187 | def __init__(self, 188 | net = 'MobileNetV2', 189 | pretrained=True, 190 | out_conv=False, 191 | cfg=None): 192 | super(MobileNet, self).__init__() 193 | self.cfg = cfg 194 | self.model = eval(net)() 195 | if pretrained: 196 | state_dict = load_state_dict_from_url(model_urls[net], 197 | progress=False) 198 | self.model.load_state_dict(state_dict, strict=False) 199 | 200 | def forward(self, x): 201 | x = self.model(x) 202 | return [x] 203 | -------------------------------------------------------------------------------- /lanedet/engine/runner.py: -------------------------------------------------------------------------------- 1 | import time 2 | import torch 3 | from tqdm import tqdm 4 | import pytorch_warmup as warmup 5 | import numpy as np 6 | import random 7 | import cv2 8 | 9 | from lanedet.models.registry import build_net 10 | from .registry import build_trainer, build_evaluator 11 | from .optimizer import build_optimizer 12 | from .scheduler import build_scheduler 13 | from lanedet.datasets import build_dataloader 14 | from lanedet.utils.recorder import build_recorder 15 | from lanedet.utils.net_utils import save_model, load_network 16 | from mmcv.parallel import MMDataParallel 17 | import torch.nn.functional as F 18 | 19 | class Runner(object): 20 | def __init__(self, cfg): 21 | torch.manual_seed(cfg.seed) 22 | np.random.seed(cfg.seed) 23 | random.seed(cfg.seed) 24 | self.cfg = cfg 25 | self.recorder = build_recorder(self.cfg) 26 | self.net = build_net(self.cfg) 27 | # self.net.to(torch.device('cuda')) 28 | # self.net = torch.nn.parallel.DataParallel( 29 | # self.net, device_ids = range(self.cfg.gpus)).cuda() 30 | self.net = MMDataParallel( 31 | self.net, device_ids = range(self.cfg.gpus)).cuda() 32 | total_params = sum(param.numel() for param in self.net.parameters()) 33 | print(f"total number of params: {total_params}") 34 | self.recorder.logger.info('Network: \n' + str(self.net)) 35 | self.resume() 36 | self.optimizer = build_optimizer(self.cfg, self.net) 37 | self.scheduler = build_scheduler(self.cfg, self.optimizer) 38 | self.warmup_scheduler = None 39 | # TODO(zhengtu): remove this hard code 40 | if self.cfg.optimizer.type == 'SGD': 41 | self.warmup_scheduler = warmup.LinearWarmup( 42 | self.optimizer, warmup_period=5000) 43 | self.detection_metric = 0. 44 | self.classification_metric = 0. 45 | self.val_loader = None 46 | self.test_loader = None 47 | 48 | def resume(self): 49 | if not self.cfg.load_from and not self.cfg.finetune_from: 50 | return 51 | load_network(self.net, self.cfg.load_from, 52 | finetune_from=self.cfg.finetune_from, logger=self.recorder.logger) 53 | 54 | def to_cuda(self, batch): 55 | for k in batch: 56 | if not isinstance(k, torch.Tensor): 57 | continue 58 | batch[k] = batch[k].cuda() 59 | return batch 60 | 61 | def train_epoch(self, epoch, train_loader): 62 | self.net.train() 63 | end = time.time() 64 | max_iter = len(train_loader) 65 | for i, data in enumerate(train_loader): 66 | if self.recorder.step >= self.cfg.total_iter: 67 | break 68 | date_time = time.time() - end 69 | #print(data['img'].shape, data['cls_label'].shape) 70 | 71 | self.recorder.step += 1 72 | data = self.to_cuda(data) 73 | self.optimizer.zero_grad() 74 | 75 | output = self.net(data) 76 | loss = output['loss'] 77 | loss.backward() 78 | self.optimizer.step() 79 | 80 | if not self.cfg.lr_update_by_epoch: 81 | self.scheduler.step() 82 | if self.warmup_scheduler: 83 | self.warmup_scheduler.dampen() 84 | batch_time = time.time() - end 85 | end = time.time() 86 | self.recorder.update_loss_stats(output['loss_stats']) 87 | self.recorder.batch_time.update(batch_time) 88 | self.recorder.data_time.update(date_time) 89 | 90 | if i % self.cfg.log_interval == 0 or i == max_iter - 1: 91 | lr = self.optimizer.param_groups[0]['lr'] 92 | self.recorder.lr = lr 93 | self.recorder.record('train') 94 | 95 | def train(self): 96 | self.recorder.logger.info('Build train loader...') 97 | train_loader = build_dataloader(self.cfg.dataset.train, self.cfg, is_train=True) 98 | 99 | self.recorder.logger.info('Start training...') 100 | for epoch in range(self.cfg.epochs): 101 | self.recorder.epoch = epoch 102 | self.train_epoch(epoch, train_loader) 103 | if (epoch + 1) % self.cfg.save_ep == 0 or epoch == self.cfg.epochs - 1: 104 | self.save_ckpt() 105 | if (epoch + 1) % self.cfg.eval_ep == 0 or epoch == self.cfg.epochs - 1: 106 | self.validate() 107 | if self.recorder.step >= self.cfg.total_iter: 108 | break 109 | if self.cfg.lr_update_by_epoch: 110 | self.scheduler.step() 111 | 112 | def validate(self): 113 | if not self.val_loader: 114 | self.val_loader = build_dataloader(self.cfg.dataset.val, self.cfg, is_train=False) 115 | self.net.eval() 116 | detection_predictions = [] 117 | classification_acc = 0 118 | for i, data in enumerate(tqdm(self.val_loader, desc=f'Validate')): 119 | data = self.to_cuda(data) 120 | with torch.no_grad(): 121 | output = self.net(data) 122 | detection_output = self.net.module.get_lanes(output)['lane_output'] 123 | detection_predictions.extend(detection_output) 124 | if self.cfg.classification: 125 | classification_acc += self.val_loader.dataset.evaluate_classification(output['category'].cuda(), data['category'].cuda()) 126 | 127 | if self.cfg.view: 128 | self.val_loader.dataset.view(detection_output, data['meta']) 129 | 130 | detection_out = self.val_loader.dataset.evaluate_detection(detection_predictions, self.cfg.work_dir) 131 | detection_metric = detection_out 132 | if detection_metric > self.detection_metric: 133 | self.detection_metric = detection_metric 134 | 135 | if self.cfg.classification: 136 | classification_acc /= len(self.val_loader) 137 | self.recorder.logger.info("Detection: " +str(detection_out) + " "+ "classification accuracy: " + str(classification_acc)) 138 | classification_metric = classification_acc 139 | if classification_metric > self.classification_metric: 140 | self.classification_metric = classification_metric 141 | self.save_ckpt(is_best=True) 142 | self.recorder.logger.info('Best detection metric: ' + str(self.detection_metric) + " " + 'Best classification metric: ' + str(self.classification_metric)) 143 | else: 144 | self.recorder.logger.info("Detection: " +str(detection_out)) 145 | self.recorder.logger.info('Best detection metric: ' + str(self.detection_metric)) 146 | 147 | def test(self): 148 | if not self.test_loader: 149 | self.test_loader = build_dataloader(self.cfg.dataset.test, self.cfg, is_train=False) 150 | self.recorder.logger.info('Start testing...') 151 | classification_acc = 0 152 | y_true = [] 153 | y_pred = [] 154 | self.net.eval() 155 | detection_predictions = [] 156 | start = torch.cuda.Event(enable_timing=True) 157 | end = torch.cuda.Event(enable_timing=True) 158 | start.record() 159 | for i, data in enumerate(tqdm(self.test_loader, desc=f'test')): 160 | data = self.to_cuda(data) 161 | with torch.no_grad(): 162 | output = self.net(data) 163 | detection_output = self.net.module.get_lanes(output)['lane_output'] 164 | detection_predictions.extend(detection_output) 165 | 166 | if self.cfg.classification: 167 | y_true.extend((data['category'].cpu().numpy()).flatten('C').tolist()) 168 | score = F.softmax(output['category'].cuda(), dim=1) 169 | score = score.argmax(dim=1) 170 | y_pred.extend((score.cpu().numpy()).flatten('C').tolist()) 171 | 172 | classification_acc += self.test_loader.dataset.evaluate_classification(output['category'].cuda(), data['category'].cuda()) 173 | if self.cfg.view: 174 | self.test_loader.dataset.view(detection_output, data['meta']) 175 | 176 | end.record() 177 | torch.cuda.synchronize() 178 | print('execution time in milliseconds per image: {}'. format(start.elapsed_time(end)/2782)) 179 | 180 | detection_out = self.test_loader.dataset.evaluate_detection(detection_predictions, self.cfg.work_dir) 181 | 182 | if self.cfg.classification: 183 | classification_acc /= len(self.test_loader) 184 | self.recorder.logger.info("Detection: " +str(detection_out) + " "+ "classification accuracy: " + str(classification_acc)) 185 | self.test_loader.dataset.plot_confusion_matrix(y_true, y_pred) 186 | else: 187 | self.recorder.logger.info("Detection: " +str(detection_out)) 188 | 189 | 190 | def save_ckpt(self, is_best=False): 191 | save_model(self.net, self.optimizer, self.scheduler, 192 | self.recorder, is_best) 193 | --------------------------------------------------------------------------------