├── lib ├── __init__.py ├── nms │ ├── src │ │ ├── nms.egg-info │ │ │ ├── top_level.txt │ │ │ ├── dependency_links.txt │ │ │ ├── SOURCES.txt │ │ │ └── PKG-INFO │ │ ├── nms │ │ │ └── __init__.py │ │ ├── nms.cpp │ │ └── nms_kernel.cu │ ├── setup.py │ └── LICENSE ├── models │ ├── __init__.py │ ├── readme.txt │ ├── __pycache__ │ │ ├── laneatt.cpython-36.pyc │ │ ├── resnet.cpython-36.pyc │ │ ├── __init__.cpython-36.pyc │ │ ├── matching.cpython-36.pyc │ │ ├── laneattoox.cpython-36.pyc │ │ └── laneattscri.cpython-36.pyc │ ├── matching.py │ └── resnet.py ├── datasets │ ├── __init__.py │ ├── __pycache__ │ │ ├── culane.cpython-36.pyc │ │ ├── ehl_wx.cpython-36.pyc │ │ ├── llamas.cpython-36.pyc │ │ ├── __init__.cpython-36.pyc │ │ ├── tusimple.cpython-36.pyc │ │ ├── lane_dataset.cpython-36.pyc │ │ ├── nolabel_dataset.cpython-36.pyc │ │ └── lane_dataset_loader.cpython-36.pyc │ ├── nolabel_dataset.py │ ├── lane_dataset_loader.py │ ├── ehl_wx.py │ ├── tusimple.py │ ├── llamas.py │ ├── culane.py │ └── lane_dataset.py ├── __pycache__ │ ├── lane.cpython-36.pyc │ ├── config.cpython-36.pyc │ ├── runner.cpython-36.pyc │ ├── __init__.cpython-36.pyc │ ├── configoox.cpython-36.pyc │ ├── experiment.cpython-36.pyc │ ├── focal_loss.cpython-36.pyc │ ├── distributed_utils.cpython-36.pyc │ └── runner_distributed.cpython-36.pyc ├── lane.py ├── configoox.py ├── config.py ├── distributed_utils.py ├── export.py ├── infer.py ├── runner.py ├── focal_loss.py ├── experiment.py ├── infer_torchscript.py └── runner_distributed.py ├── utils ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── llamas_utils.cpython-36.pyc │ ├── culane_metric.cpython-36.pyc │ ├── llamas_metric.cpython-36.pyc │ └── tusimple_metric.cpython-36.pyc ├── viz_dataset.py ├── speed.py ├── gen_anchor_mask.py ├── gen_video.py ├── tusimple_metric.py ├── culane_metric.py ├── llamas_metric.py └── llamas_utils.py ├── data ├── culane_anchors_freq.pt ├── ehl_wx_anchors_freq.pt ├── llamas_anchors_freq.pt └── tusimple_anchors_freq.pt ├── requirements.txt ├── README.md ├── cfgs ├── laneatt_llamas_resnet122.yml ├── laneatt_culane_resnet18.yml ├── laneatt_llamas_resnet18.yml ├── laneatt_llamas_resnet34.yml ├── laneatt_culane_resnet122.yml ├── laneatt_tusimple_resnet18.yml ├── laneatt_culane_resnet34.yml ├── laneatt_ehl_wx_resnet34.yml ├── laneatt_tusimple_resnet122.yml └── laneatt_tusimple_resnet34.yml ├── main.py └── main_distributed.py /lib/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /lib/nms/src/nms.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | nms 2 | -------------------------------------------------------------------------------- /lib/nms/src/nms.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /lib/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .laneattscri import LaneATT 2 | 3 | __all__ = ["LaneATT"] 4 | -------------------------------------------------------------------------------- /lib/models/readme.txt: -------------------------------------------------------------------------------- 1 | 训练 模型转换 修改__init__.py 首行 2 | 训 laneatt.py 3 | 转onnx 4 | 转torchscript -------------------------------------------------------------------------------- /lib/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .lane_dataset import LaneDataset 2 | 3 | __all__ = ["LaneDataset"] 4 | -------------------------------------------------------------------------------- /data/culane_anchors_freq.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/data/culane_anchors_freq.pt -------------------------------------------------------------------------------- /data/ehl_wx_anchors_freq.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/data/ehl_wx_anchors_freq.pt -------------------------------------------------------------------------------- /data/llamas_anchors_freq.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/data/llamas_anchors_freq.pt -------------------------------------------------------------------------------- /data/tusimple_anchors_freq.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/data/tusimple_anchors_freq.pt -------------------------------------------------------------------------------- /lib/__pycache__/lane.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/__pycache__/lane.cpython-36.pyc -------------------------------------------------------------------------------- /lib/__pycache__/config.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/__pycache__/config.cpython-36.pyc -------------------------------------------------------------------------------- /lib/__pycache__/runner.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/__pycache__/runner.cpython-36.pyc -------------------------------------------------------------------------------- /lib/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /lib/__pycache__/configoox.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/__pycache__/configoox.cpython-36.pyc -------------------------------------------------------------------------------- /lib/__pycache__/experiment.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/__pycache__/experiment.cpython-36.pyc -------------------------------------------------------------------------------- /lib/__pycache__/focal_loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/__pycache__/focal_loss.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/utils/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /lib/models/__pycache__/laneatt.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/models/__pycache__/laneatt.cpython-36.pyc -------------------------------------------------------------------------------- /lib/models/__pycache__/resnet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/models/__pycache__/resnet.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/llamas_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/utils/__pycache__/llamas_utils.cpython-36.pyc -------------------------------------------------------------------------------- /lib/datasets/__pycache__/culane.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/datasets/__pycache__/culane.cpython-36.pyc -------------------------------------------------------------------------------- /lib/datasets/__pycache__/ehl_wx.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/datasets/__pycache__/ehl_wx.cpython-36.pyc -------------------------------------------------------------------------------- /lib/datasets/__pycache__/llamas.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/datasets/__pycache__/llamas.cpython-36.pyc -------------------------------------------------------------------------------- /lib/models/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/models/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /lib/models/__pycache__/matching.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/models/__pycache__/matching.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/culane_metric.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/utils/__pycache__/culane_metric.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/llamas_metric.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/utils/__pycache__/llamas_metric.cpython-36.pyc -------------------------------------------------------------------------------- /lib/__pycache__/distributed_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/__pycache__/distributed_utils.cpython-36.pyc -------------------------------------------------------------------------------- /lib/__pycache__/runner_distributed.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/__pycache__/runner_distributed.cpython-36.pyc -------------------------------------------------------------------------------- /lib/datasets/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/datasets/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /lib/datasets/__pycache__/tusimple.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/datasets/__pycache__/tusimple.cpython-36.pyc -------------------------------------------------------------------------------- /lib/models/__pycache__/laneattoox.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/models/__pycache__/laneattoox.cpython-36.pyc -------------------------------------------------------------------------------- /lib/models/__pycache__/laneattscri.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/models/__pycache__/laneattscri.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/tusimple_metric.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/utils/__pycache__/tusimple_metric.cpython-36.pyc -------------------------------------------------------------------------------- /lib/datasets/__pycache__/lane_dataset.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/datasets/__pycache__/lane_dataset.cpython-36.pyc -------------------------------------------------------------------------------- /lib/datasets/__pycache__/nolabel_dataset.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/datasets/__pycache__/nolabel_dataset.cpython-36.pyc -------------------------------------------------------------------------------- /lib/datasets/__pycache__/lane_dataset_loader.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Hust-wayne/Multi_Lanes_Detection/HEAD/lib/datasets/__pycache__/lane_dataset_loader.cpython-36.pyc -------------------------------------------------------------------------------- /lib/nms/src/nms.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | setup.py 2 | src/nms.cpp 3 | src/nms_kernel.cu 4 | src/nms/__init__.py 5 | src/nms.egg-info/PKG-INFO 6 | src/nms.egg-info/SOURCES.txt 7 | src/nms.egg-info/dependency_links.txt 8 | src/nms.egg-info/top_level.txt -------------------------------------------------------------------------------- /lib/nms/src/nms.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 1.0 2 | Name: nms 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: UNKNOWN 7 | Author-email: UNKNOWN 8 | License: UNKNOWN 9 | Description: UNKNOWN 10 | Platform: UNKNOWN 11 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | opencv_python==4.2.0.32 2 | Shapely==1.7.0 3 | xmljson==0.2.0 4 | thop==0.0.31.post2005241907 5 | matplotlib==3.1.3 6 | imgaug==0.4.0 7 | scipy==1.4.1 8 | p_tqdm==1.3.3 9 | lxml==4.5.0 10 | tqdm==4.43.0 11 | ujson==1.35 12 | PyYAML==5.3.1 13 | scikit_learn==0.23.2 14 | tensorboard==2.3.0 15 | gdown==3.12.2 16 | -------------------------------------------------------------------------------- /lib/nms/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | from torch.utils.cpp_extension import CUDAExtension, BuildExtension 4 | 5 | setup(name='nms', packages=['nms'], 6 | package_dir={'':'src'}, 7 | ext_modules=[CUDAExtension('nms.details', ['src/nms.cpp', 'src/nms_kernel.cu'])], 8 | cmdclass={'build_ext': BuildExtension}) 9 | -------------------------------------------------------------------------------- /utils/viz_dataset.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | import cv2 4 | import torch 5 | import random 6 | import numpy as np 7 | 8 | from lib.config import Config 9 | 10 | 11 | def parse_args(): 12 | parser = argparse.ArgumentParser(description="Visualize a dataset") 13 | parser.add_argument("--cfg", help="Config file") 14 | parser.add_argument("--split", 15 | choices=["train", "test", "val"], 16 | default='train', 17 | help="Dataset split to visualize") 18 | args = parser.parse_args() 19 | 20 | return args 21 | 22 | 23 | def main(): 24 | np.random.seed(0) 25 | torch.manual_seed(0) 26 | random.seed(0) 27 | args = parse_args() 28 | cfg = Config(args.cfg) 29 | train_dataset = cfg.get_dataset(args.split) 30 | for idx in range(len(train_dataset)): 31 | img, _, _ = train_dataset.draw_annotation(idx) 32 | cv2.imshow('sample', img) 33 | cv2.waitKey(0) 34 | 35 | 36 | if __name__ == "__main__": 37 | main() 38 | -------------------------------------------------------------------------------- /lib/lane.py: -------------------------------------------------------------------------------- 1 | from scipy.interpolate import InterpolatedUnivariateSpline 2 | 3 | 4 | class Lane: 5 | def __init__(self, points=None, invalid_value=-2., metadata=None): 6 | super(Lane, self).__init__() 7 | self.curr_iter = 0 8 | self.points = points 9 | self.invalid_value = invalid_value 10 | self.function = InterpolatedUnivariateSpline(points[:, 1], points[:, 0], k=min(3, len(points) - 1)) 11 | self.min_y = points[:, 1].min() - 0.01 12 | self.max_y = points[:, 1].max() + 0.01 13 | 14 | self.metadata = metadata or {} 15 | 16 | def __repr__(self): 17 | return '[Lane]\n' + str(self.points) + '\n[/Lane]' 18 | 19 | def __call__(self, lane_ys): 20 | lane_xs = self.function(lane_ys) 21 | 22 | lane_xs[(lane_ys < self.min_y) | (lane_ys > self.max_y)] = self.invalid_value 23 | return lane_xs 24 | 25 | def __iter__(self): 26 | return self 27 | 28 | def __next__(self): 29 | if self.curr_iter < len(self.points): 30 | self.curr_iter += 1 31 | return self.points[self.curr_iter - 1] 32 | self.curr_iter = 0 33 | raise StopIteration 34 | -------------------------------------------------------------------------------- /lib/datasets/nolabel_dataset.py: -------------------------------------------------------------------------------- 1 | import glob 2 | 3 | from .lane_dataset_loader import LaneDatasetLoader 4 | 5 | 6 | class NoLabelDataset(LaneDatasetLoader): 7 | def __init__(self, img_h=720, img_w=1280, max_lanes=None, root=None, img_ext='.jpg', **_): 8 | """Use this loader if you want to test a model on an image without annotations or implemented loader.""" 9 | self.root = root 10 | if root is None: 11 | raise Exception('Please specify the root directory') 12 | 13 | self.img_w, self.img_h = img_w, img_h 14 | self.img_ext = img_ext 15 | self.annotations = [] 16 | self.load_annotations() 17 | 18 | # Force max_lanes, used when evaluating testing with models trained on other datasets 19 | # On NoLabelDataset, always force it 20 | self.max_lanes = max_lanes 21 | 22 | def get_img_heigth(self, _): 23 | return self.img_h 24 | 25 | def get_img_width(self, _): 26 | return self.img_w 27 | 28 | def get_metrics(self, lanes, _): 29 | return 0, 0, [1] * len(lanes), [1] * len(lanes) 30 | 31 | def load_annotations(self): 32 | self.annotations = [] 33 | pattern = '{}/**/*{}'.format(self.root, self.img_ext) 34 | print('Looking for image files with the pattern', pattern) 35 | for file in glob.glob(pattern, recursive=True): 36 | self.annotations.append({'lanes': [], 'path': file}) 37 | 38 | def eval(self, _, __, ___, ____, _____): 39 | return "", None 40 | 41 | def __getitem__(self, idx): 42 | return self.annotations[idx] 43 | 44 | def __len__(self): 45 | return len(self.annotations) 46 | -------------------------------------------------------------------------------- /lib/nms/src/nms/__init__.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 details 30 | 31 | 32 | def nms(boxes, scores, overlap, top_k): 33 | return details.nms_forward(boxes, scores, overlap, top_k) 34 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Multi_Lanes_Detection 2 | 本项目基于LaneATT，做出了一些改进，具体见下文 3 | # 2D车道线识别相关技术 4 | 5 | ## 1. 车道线表示 6 | 7 | - 输出类型：掩码/点集/矢量线条 8 | - 实例化：每个车道线是否形成实例 9 | - 分类：是否对车道线进行了分类（单白、双黄等） 10 | - 提前定义的参数：是否只能检测固定数量的车道线 11 | - 车道标记：是否对车道上的行车标记也做了检测 12 | 13 | ## 2. 目前方法 14 | 15 | - 传统方法-边缘检测滤波等方式分割出车道线区域，然后结合霍夫变换、RANSAC等算法进行车道线检测 16 | 17 | - 基于霍夫变换的车道线检测； 18 | 19 | - 基于LSD直线的车道线检测； 20 | 21 | - 基于俯视图变换的车道线检测； 22 | 23 | - 基于拟合的车道线检测； 24 | 25 | - 基于平行透视灭点的车道线检测； 26 | 27 | - 基于语义分割 28 | 29 | - 二值语义分割主要采用CNN方法并引入一些方式提高语义分割精度，在线的拟合阶段可以采用学习到的转换矩阵先将分割结果转换为鸟瞰图视角，然后，采用均匀取点+最小二乘法拟合，拟合方程可选三次方程（Lanenet） 30 | 31 | - 基于点的表示方法，车道线用固定数量的点表示 32 | 33 | - raw表示方法-将图像每行分成若干个grid，对每行grids进行分类选中存在车道线的grid(Ultra Fast Lane Detection) 34 | - 类似目标检测的anchor，生成lane anchor, 直接回归点与anchor做匹配(lLineCNN，LaneATT) 35 | 36 | - 利用多项式进行建模(Polynet) 37 | 38 | ## 3. 相关数据集 39 | 40 | - tusimple： 共72k张图片，位于高速路，天气晴朗，车道线清晰，特点是车道线以点来标注；图片大小：1280x720 41 | - culane ：共98k张图片，包含拥挤，黑夜，无线，暗影等八种难以检测的情况，最多标记4条车道线；图片大小：1640x590 ； 42 | - 百度ApolloScape ：140k张图片，特点是车道线以掩码的形式标注，包含2/3维28个类别；图片大小：3384x2710 43 | - CurveLanes：华为弯道检测数据集 135k张图片, 采用三次样条曲线手动标注所有车道线,包括很多较为复杂的场景,如S路、Y车道,还有夜间和多车道的场景。分为训练集10万张，验证集2万张，测试级3万张；图片大小：2650x1440 44 | - LLAMAS 45 | - 自定义数据集（ehl_wx无锡车联网公交车载2D）: 40W张图片，位于市区、郊区等，天气有变化，白天居多，车道线复杂，车道线以点标注，图片大小：1920*1080 46 | 47 | ## 4. 技术改进及实现 48 | 49 | - [x] LaneATT多类别车道线识别实现，代码重构，支持多类别参数化输入 50 | 51 | - [x] LaneATT分布式训练DDP实现（提速明显） 52 | 53 | - [x] LaneATT模型转torchscript及相关推理代码实现 54 | 55 | - [x] LaneATT模型转ONNX及相关推理代码实现 56 | - [ ] LaneATT tensorrt（python） 推理 57 | - [ ] LaneATT C++ 推理（libtorch, tensorrt） 58 | - [ ] LaneATT 结构改进优化提点 59 | 60 | - [x] 针对Culane和Tusimple公开数据集进行车道线类别标注及简易标注工具开发实现,[相关数据集已开源](https://blog.csdn.net/hustwayne/article/details/121139364?spm=1001.2014.3001.5501) 61 | 62 | - [x] 根据实际场景自定义车道线数据集 63 | 64 | - [x] Ultra Fast Lane Detection ONNX 及推理代码实现 65 | 66 | 67 | 68 | -------------------------------------------------------------------------------- /lib/configoox.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | import torch 3 | import models as models 4 | #import datasets as datasets 5 | 6 | 7 | class Config: 8 | def __init__(self, config_path): 9 | self.config = {} 10 | self.config_str = "" 11 | self.load(config_path) 12 | 13 | def load(self, path): 14 | with open(path, 'r') as file: 15 | self.config_str = file.read() 16 | self.config = yaml.load(self.config_str, Loader=yaml.FullLoader) 17 | 18 | def __repr__(self): 19 | return self.config_str 20 | 21 | def get_dataset(self, split): 22 | return getattr(datasets, 23 | self.config['datasets'][split]['type'])(**self.config['datasets'][split]['parameters']) 24 | 25 | def get_model(self, **kwargs): 26 | name = self.config['model']['name'] 27 | parameters = self.config['model']['parameters'] 28 | return getattr(models, name)(**parameters, **kwargs) 29 | 30 | def get_optimizer(self, model_parameters): 31 | return getattr(torch.optim, self.config['optimizer']['name'])(model_parameters, 32 | **self.config['optimizer']['parameters']) 33 | 34 | def get_lr_scheduler(self, optimizer): 35 | return getattr(torch.optim.lr_scheduler, 36 | self.config['lr_scheduler']['name'])(optimizer, **self.config['lr_scheduler']['parameters']) 37 | 38 | def get_loss_parameters(self): 39 | return self.config['loss_parameters'] 40 | 41 | def get_train_parameters(self): 42 | return self.config['train_parameters'] 43 | 44 | def get_test_parameters(self): 45 | return self.config['test_parameters'] 46 | 47 | def __getitem__(self, item): 48 | return self.config[item] 49 | 50 | def __contains__(self, item): 51 | return item in self.config 52 | -------------------------------------------------------------------------------- /lib/config.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | import torch 3 | import lib.models as models 4 | import lib.datasets as datasets 5 | 6 | 7 | class Config: 8 | def __init__(self, config_path): 9 | self.config = {} 10 | self.config_str = "" 11 | self.load(config_path) 12 | 13 | def load(self, path): 14 | with open(path, 'r') as file: 15 | self.config_str = file.read() 16 | self.config = yaml.load(self.config_str, Loader=yaml.FullLoader) 17 | 18 | def __repr__(self): 19 | return self.config_str 20 | 21 | def get_dataset(self, split): 22 | return getattr(datasets, 23 | self.config['datasets'][split]['type'])(**self.config['datasets'][split]['parameters']) 24 | 25 | def get_model(self, **kwargs): 26 | name = self.config['model']['name'] 27 | parameters = self.config['model']['parameters'] 28 | return getattr(models, name)(**parameters, **kwargs) 29 | 30 | def get_optimizer(self, model_parameters): 31 | return getattr(torch.optim, self.config['optimizer']['name'])(model_parameters, 32 | **self.config['optimizer']['parameters']) 33 | 34 | def get_lr_scheduler(self, optimizer): 35 | return getattr(torch.optim.lr_scheduler, 36 | self.config['lr_scheduler']['name'])(optimizer, **self.config['lr_scheduler']['parameters']) 37 | 38 | def get_loss_parameters(self): 39 | return self.config['loss_parameters'] 40 | 41 | def get_train_parameters(self): 42 | return self.config['train_parameters'] 43 | 44 | def get_test_parameters(self): 45 | return self.config['test_parameters'] 46 | 47 | def __getitem__(self, item): 48 | return self.config[item] 49 | 50 | def __contains__(self, item): 51 | return item in self.config 52 | -------------------------------------------------------------------------------- /utils/speed.py: -------------------------------------------------------------------------------- 1 | import time 2 | import argparse 3 | 4 | import torch 5 | from thop import profile, clever_format 6 | 7 | from lib.config import Config 8 | 9 | 10 | def parse_args(): 11 | parser = argparse.ArgumentParser(description="Tool to measure a model's speed") 12 | parser.add_argument("--cfg", default="config.yaml", help="Config file") 13 | parser.add_argument("--model_path", help="Model checkpoint path (optional)") 14 | parser.add_argument('--iters', default=100, type=int, help="Number of times to run the model and get the average") 15 | 16 | return parser.parse_args() 17 | 18 | 19 | # torch.backends.cudnn.benchmark = True 20 | 21 | 22 | def main(): 23 | args = parse_args() 24 | cfg = Config(args.cfg) 25 | device = torch.device('cuda') 26 | model = cfg.get_model() 27 | model = model.to(device) 28 | test_parameters = cfg.get_test_parameters() 29 | height, width = cfg['datasets']['test']['parameters']['img_size'] 30 | 31 | if args.model_path is not None: 32 | model.load_state_dict(torch.load(args.model_path)['model']) 33 | 34 | model.eval() 35 | 36 | x = torch.zeros((1, 3, height, width)).to(device) + 1 37 | 38 | # Benchmark MACs and params 39 | 40 | macs, params = profile(model, inputs=(x,)) 41 | macs, params = clever_format([macs, params], "%.3f") 42 | print('MACs: {}'.format(macs)) 43 | print('Params: {}'.format(params)) 44 | 45 | # GPU warmup 46 | for _ in range(100): 47 | model(x) 48 | 49 | # Benchmark latency and FPS 50 | t_all = 0 51 | for _ in range(args.iters): 52 | t1 = time.time() 53 | model(x, **test_parameters) 54 | t2 = time.time() 55 | t_all += t2 - t1 56 | 57 | print('Average latency (ms): {:.2f}'.format(t_all * 1000 / args.iters)) 58 | print('Average FPS: {:.2f}'.format(args.iters / t_all)) 59 | 60 | 61 | if __name__ == '__main__': 62 | main() 63 | -------------------------------------------------------------------------------- /lib/distributed_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import torch 4 | import torch.distributed as dist 5 | 6 | 7 | def init_distributed_mode(args): 8 | if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: 9 | args.rank = int(os.environ["RANK"]) 10 | args.world_size = int(os.environ['WORLD_SIZE']) 11 | args.gpu = int(os.environ['LOCAL_RANK']) 12 | elif 'SLURM_PROCID' in os.environ: 13 | args.rank = int(os.environ['SLURM_PROCID']) 14 | args.gpu = args.rank % torch.cuda.device_count() 15 | else: 16 | print('Not using distributed mode') 17 | args.distributed = False 18 | return 19 | 20 | args.distributed = True 21 | 22 | torch.cuda.set_device(args.gpu) 23 | args.dist_backend = 'nccl' # 通信后端，nvidia GPU推荐使用NCCL 24 | print('| distributed init (rank {}): {}'.format( 25 | args.rank, args.dist_url), flush=True) 26 | dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, 27 | world_size=args.world_size, rank=args.rank) 28 | dist.barrier() 29 | 30 | 31 | def cleanup(): 32 | dist.destroy_process_group() 33 | 34 | 35 | def is_dist_avail_and_initialized(): 36 | """检查是否支持分布式环境""" 37 | if not dist.is_available(): 38 | return False 39 | if not dist.is_initialized(): 40 | return False 41 | return True 42 | 43 | 44 | def get_world_size(): 45 | if not is_dist_avail_and_initialized(): 46 | return 1 47 | return dist.get_world_size() 48 | 49 | 50 | def get_rank(): 51 | if not is_dist_avail_and_initialized(): 52 | return 0 53 | return dist.get_rank() 54 | 55 | 56 | def is_main_process(): 57 | return get_rank() == 0 58 | 59 | 60 | def reduce_value(value, average=True): 61 | world_size = get_world_size() 62 | if world_size < 2: # 单GPU的情况 63 | return value 64 | 65 | with torch.no_grad(): 66 | dist.all_reduce(value) 67 | if average: 68 | value /= world_size 69 | 70 | return value 71 | -------------------------------------------------------------------------------- /cfgs/laneatt_llamas_resnet122.yml: -------------------------------------------------------------------------------- 1 | # Training settings 2 | val_every: 1000 3 | model_checkpoint_interval: 1 4 | seed: 0 5 | model: 6 | name: LaneATT 7 | parameters: 8 | backbone: resnet122 9 | S: &S 72 10 | topk_anchors: 1000 11 | anchors_freq_path: 'data/llamas_anchors_freq.pt' 12 | img_h: &img_h 360 13 | img_w: &img_w 640 14 | batch_size: 16 15 | epochs: 15 16 | loss_parameters: {} 17 | train_parameters: 18 | conf_threshold: 19 | nms_thres: 15. 20 | nms_topk: 3000 21 | test_parameters: 22 | conf_threshold: 0.5 23 | nms_thres: 50. 24 | nms_topk: &max_lanes 4 25 | optimizer: 26 | name: Adam 27 | parameters: 28 | lr: 0.0003 29 | lr_scheduler: 30 | name: CosineAnnealingLR 31 | parameters: 32 | T_max: 54615 # 15 * 29135 iterations 33 | 34 | # Dataset settings 35 | datasets: 36 | train: 37 | type: LaneDataset 38 | parameters: 39 | S: *S 40 | dataset: llamas 41 | split: train 42 | img_size: [*img_h, *img_w] 43 | max_lanes: *max_lanes 44 | normalize: false 45 | aug_chance: 1.0 46 | augmentations: 47 | - name: Affine 48 | parameters: 49 | translate_px: 50 | x: !!python/tuple [-25, 25] 51 | y: !!python/tuple [-10, 10] 52 | rotate: !!python/tuple [-6, 6] 53 | scale: !!python/tuple [0.85, 1.15] 54 | - name: HorizontalFlip 55 | parameters: 56 | p: 0.5 57 | 58 | root: "datasets/llamas" 59 | 60 | test: 61 | type: LaneDataset 62 | parameters: 63 | S: *S 64 | dataset: llamas 65 | split: test 66 | img_size: [*img_h, *img_w] 67 | normalize: false 68 | aug_chance: 0 69 | augmentations: 70 | root: "datasets/llamas" 71 | 72 | 73 | val: 74 | type: LaneDataset 75 | parameters: 76 | S: *S 77 | dataset: llamas 78 | split: val 79 | img_size: [img_h, *img_w] 80 | max_lanes: *max_lanes 81 | normalize: false 82 | aug_chance: 0 83 | augmentations: 84 | root: "datasets/llamas" 85 | -------------------------------------------------------------------------------- /cfgs/laneatt_culane_resnet18.yml: -------------------------------------------------------------------------------- 1 | # Model settings 2 | val_every: 1000 3 | model_checkpoint_interval: 1 4 | seed: 0 5 | model: 6 | name: LaneATT 7 | parameters: 8 | backbone: resnet18 9 | S: &S 72 10 | topk_anchors: 1000 11 | anchors_freq_path: 'data/culane_anchors_freq.pt' 12 | img_h: &img_h 360 13 | img_w: &img_w 640 14 | batch_size: 8 15 | epochs: 15 16 | loss_parameters: {} 17 | train_parameters: 18 | conf_threshold: 19 | nms_thres: 15. 20 | nms_topk: 3000 21 | test_parameters: 22 | conf_threshold: 0.5 23 | nms_thres: 50. 24 | nms_topk: &max_lanes 4 25 | optimizer: 26 | name: Adam 27 | parameters: 28 | lr: 0.0003 29 | lr_scheduler: 30 | name: CosineAnnealingLR 31 | parameters: 32 | T_max: 166650 # 15 * 11110 iterations 33 | 34 | # Dataset settings 35 | datasets: 36 | train: 37 | type: LaneDataset 38 | parameters: 39 | S: *S 40 | dataset: culane 41 | split: train 42 | img_size: [*img_h, *img_w] 43 | max_lanes: *max_lanes 44 | normalize: false 45 | aug_chance: 1.0 46 | augmentations: 47 | - name: Affine 48 | parameters: 49 | translate_px: 50 | x: !!python/tuple [-25, 25] 51 | y: !!python/tuple [-10, 10] 52 | rotate: !!python/tuple [-6, 6] 53 | scale: !!python/tuple [0.85, 1.15] 54 | - name: HorizontalFlip 55 | parameters: 56 | p: 0.5 57 | 58 | root: "datasets/culane" 59 | 60 | test: 61 | type: LaneDataset 62 | parameters: 63 | S: *S 64 | dataset: culane 65 | split: test 66 | img_size: [*img_h, *img_w] 67 | max_lanes: *max_lanes 68 | normalize: false 69 | aug_chance: 0 70 | augmentations: 71 | root: "datasets/culane" 72 | 73 | val: 74 | type: LaneDataset 75 | parameters: 76 | S: *S 77 | dataset: culane 78 | split: val 79 | img_size: [288, 512] 80 | max_lanes: *max_lanes 81 | normalize: false 82 | aug_chance: 0 83 | augmentations: 84 | root: "datasets/culane" 85 | -------------------------------------------------------------------------------- /cfgs/laneatt_llamas_resnet18.yml: -------------------------------------------------------------------------------- 1 | # Training settings 2 | val_every: 1000 3 | model_checkpoint_interval: 1 4 | seed: 0 5 | model: 6 | name: LaneATT 7 | parameters: 8 | backbone: resnet18 9 | S: &S 72 10 | topk_anchors: 1000 11 | anchors_freq_path: 'data/llamas_anchors_freq.pt' 12 | img_h: &img_h 360 13 | img_w: &img_w 640 14 | batch_size: 8 15 | epochs: 15 16 | loss_parameters: {} 17 | train_parameters: 18 | conf_threshold: 19 | nms_thres: 15. 20 | nms_topk: 3000 21 | test_parameters: 22 | conf_threshold: 0.5 23 | nms_thres: 50. 24 | nms_topk: &max_lanes 4 25 | optimizer: 26 | name: Adam 27 | parameters: 28 | lr: 0.0003 29 | lr_scheduler: 30 | name: CosineAnnealingLR 31 | parameters: 32 | T_max: 109260 # 15 * 7284 iterations 33 | 34 | # Dataset settings 35 | datasets: 36 | train: 37 | type: LaneDataset 38 | parameters: 39 | S: *S 40 | dataset: llamas 41 | split: train 42 | img_size: [*img_h, *img_w] 43 | max_lanes: *max_lanes 44 | normalize: false 45 | aug_chance: 1.0 46 | augmentations: 47 | - name: Affine 48 | parameters: 49 | translate_px: 50 | x: !!python/tuple [-25, 25] 51 | y: !!python/tuple [-10, 10] 52 | rotate: !!python/tuple [-6, 6] 53 | scale: !!python/tuple [0.85, 1.15] 54 | - name: HorizontalFlip 55 | parameters: 56 | p: 0.5 57 | 58 | root: "datasets/llamas" 59 | 60 | test: 61 | type: LaneDataset 62 | parameters: 63 | S: *S 64 | dataset: llamas 65 | split: test 66 | img_size: [*img_h, *img_w] 67 | max_lanes: 5 68 | normalize: false 69 | aug_chance: 0 70 | augmentations: 71 | root: "datasets/llamas" 72 | 73 | 74 | val: 75 | type: LaneDataset 76 | parameters: 77 | S: *S 78 | dataset: llamas 79 | split: val 80 | img_size: [img_h, *img_w] 81 | max_lanes: *max_lanes 82 | normalize: false 83 | aug_chance: 0 84 | augmentations: 85 | root: "datasets/llamas" 86 | -------------------------------------------------------------------------------- /cfgs/laneatt_llamas_resnet34.yml: -------------------------------------------------------------------------------- 1 | # Training settings 2 | val_every: 1000 3 | model_checkpoint_interval: 1 4 | seed: 0 5 | model: 6 | name: LaneATT 7 | parameters: 8 | backbone: resnet34 9 | S: &S 72 10 | topk_anchors: 1000 11 | anchors_freq_path: 'data/llamas_anchors_freq.pt' 12 | img_h: &img_h 360 13 | img_w: &img_w 640 14 | batch_size: 8 15 | epochs: 15 16 | loss_parameters: {} 17 | train_parameters: 18 | conf_threshold: 19 | nms_thres: 15. 20 | nms_topk: 3000 21 | test_parameters: 22 | conf_threshold: 0.5 23 | nms_thres: 50. 24 | nms_topk: &max_lanes 4 25 | optimizer: 26 | name: Adam 27 | parameters: 28 | lr: 0.0003 29 | lr_scheduler: 30 | name: CosineAnnealingLR 31 | parameters: 32 | T_max: 109260 # 15 * 7284 iterations 33 | 34 | # Dataset settings 35 | datasets: 36 | train: 37 | type: LaneDataset 38 | parameters: 39 | S: *S 40 | dataset: llamas 41 | split: train 42 | img_size: [*img_h, *img_w] 43 | max_lanes: *max_lanes 44 | normalize: false 45 | aug_chance: 1.0 46 | augmentations: 47 | - name: Affine 48 | parameters: 49 | translate_px: 50 | x: !!python/tuple [-25, 25] 51 | y: !!python/tuple [-10, 10] 52 | rotate: !!python/tuple [-6, 6] 53 | scale: !!python/tuple [0.85, 1.15] 54 | - name: HorizontalFlip 55 | parameters: 56 | p: 0.5 57 | 58 | root: "datasets/llamas" 59 | 60 | test: 61 | type: LaneDataset 62 | parameters: 63 | S: *S 64 | dataset: llamas 65 | split: test 66 | img_size: [*img_h, *img_w] 67 | max_lanes: 5 68 | normalize: false 69 | aug_chance: 0 70 | augmentations: 71 | root: "datasets/llamas" 72 | 73 | 74 | val: 75 | type: LaneDataset 76 | parameters: 77 | S: *S 78 | dataset: llamas 79 | split: val 80 | img_size: [img_h, *img_w] 81 | max_lanes: *max_lanes 82 | normalize: false 83 | aug_chance: 0 84 | augmentations: 85 | root: "datasets/llamas" 86 | -------------------------------------------------------------------------------- /cfgs/laneatt_culane_resnet122.yml: -------------------------------------------------------------------------------- 1 | # Model settings 2 | val_every: 1000 3 | model_checkpoint_interval: 1 4 | seed: 0 5 | model: 6 | name: LaneATT 7 | parameters: 8 | backbone: resnet122 9 | S: &S 72 10 | topk_anchors: 1000 11 | anchors_freq_path: 'data/culane_anchors_freq.pt' 12 | img_h: &img_h 360 13 | img_w: &img_w 640 14 | batch_size: 4 15 | epochs: 15 16 | loss_parameters: {} 17 | train_parameters: 18 | conf_threshold: 19 | nms_thres: 15. 20 | nms_topk: 3000 21 | test_parameters: 22 | conf_threshold: 0.5 23 | nms_thres: 50. 24 | nms_topk: &max_lanes 4 25 | optimizer: 26 | name: Adam 27 | parameters: 28 | lr: 0.0003 29 | lr_scheduler: 30 | name: CosineAnnealingLR 31 | parameters: 32 | T_max: 333300 # 15 * 22220 iterations 33 | 34 | # Dataset settings 35 | datasets: 36 | train: 37 | type: LaneDataset 38 | parameters: 39 | S: *S 40 | dataset: culane 41 | split: train 42 | img_size: [*img_h, *img_w] 43 | max_lanes: *max_lanes 44 | normalize: false 45 | aug_chance: 1.0 46 | augmentations: 47 | - name: Affine 48 | parameters: 49 | translate_px: 50 | x: !!python/tuple [-25, 25] 51 | y: !!python/tuple [-10, 10] 52 | rotate: !!python/tuple [-6, 6] 53 | scale: !!python/tuple [0.85, 1.15] 54 | - name: HorizontalFlip 55 | parameters: 56 | p: 0.5 57 | 58 | root: "datasets/culane" 59 | 60 | test: 61 | type: LaneDataset 62 | parameters: 63 | S: *S 64 | dataset: culane 65 | split: test 66 | img_size: [*img_h, *img_w] 67 | max_lanes: *max_lanes 68 | normalize: false 69 | aug_chance: 0 70 | augmentations: 71 | root: "datasets/culane" 72 | 73 | val: 74 | type: LaneDataset 75 | parameters: 76 | S: *S 77 | dataset: culane 78 | split: val 79 | img_size: [*img_h, *img_w] 80 | max_lanes: *max_lanes 81 | normalize: false 82 | aug_chance: 0 83 | augmentations: 84 | root: "datasets/culane" 85 | -------------------------------------------------------------------------------- /lib/datasets/lane_dataset_loader.py: -------------------------------------------------------------------------------- 1 | class LaneDatasetLoader: 2 | def get_img_heigth(self, path): 3 | """Returns the image's height in pixels""" 4 | raise NotImplementedError() 5 | 6 | def get_img_width(self, path): 7 | """Returns the image's width in pixels""" 8 | raise NotImplementedError() 9 | 10 | def get_metrics(self, lanes, idx): 11 | """Returns dataset's metrics for a prediction `lanes` 12 | 13 | A tuple `(fp, fn, matches, accs)` should be returned, where `fp` and `fn` indicate the number of false-positives 14 | and false-negatives, respectively, matches` is a list with a boolean value for each 15 | prediction in `lanes` indicating if the prediction is a true positive and `accs` is a metric indicating the 16 | quality of each prediction (e.g., the IoU with an annotation) 17 | 18 | If the metrics can't be computed, placeholder values should be returned. 19 | """ 20 | raise NotImplementedError() 21 | 22 | def load_annotations(self): 23 | """Loads all annotations from the dataset 24 | 25 | Should return a list where each item is a dictionary with keys `path` and `lanes`, where `path` is the path to 26 | the image and `lanes` is a list of lanes, represented by a list of points for example: 27 | 28 | return [{ 29 | 'path': 'example/path.png' # path to the image 30 | 'lanes': [[10, 20], [20, 25]] 31 | }] 32 | """ 33 | raise NotImplementedError() 34 | 35 | def eval_predictions(self, predictions, output_basedir): 36 | """Should return a dictionary with each metric's results 37 | Example: 38 | return { 39 | 'F1': 0.9 40 | 'Acc': 0.95 41 | } 42 | """ 43 | raise NotImplementedError() 44 | 45 | def __getitem__(self, idx): 46 | """Should return the annotation with index idx""" 47 | raise NotImplementedError() 48 | 49 | def __len__(self): 50 | """Should return the number of samples in the dataset""" 51 | raise NotImplementedError() 52 | -------------------------------------------------------------------------------- /cfgs/laneatt_tusimple_resnet18.yml: -------------------------------------------------------------------------------- 1 | # Model settings 2 | val_every: 10 3 | model_checkpoint_interval: 1 4 | seed: 0 5 | model: 6 | name: LaneATT 7 | parameters: 8 | backbone: resnet18 9 | S: &S 72 10 | topk_anchors: 1000 11 | anchors_freq_path: 'data/tusimple_anchors_freq.pt' 12 | img_h: &img_h 360 13 | img_w: &img_w 640 14 | batch_size: 8 15 | epochs: 100 16 | loss_parameters: {} 17 | train_parameters: 18 | conf_threshold: 19 | nms_thres: 15. 20 | nms_topk: 3000 21 | test_parameters: 22 | conf_threshold: 0.2 23 | nms_thres: 45. 24 | nms_topk: &max_lanes 5 25 | optimizer: 26 | name: Adam 27 | parameters: 28 | lr: 0.0003 29 | lr_scheduler: 30 | name: CosineAnnealingLR 31 | parameters: 32 | T_max: 45400 # 100 * 454 iterations 33 | 34 | # Dataset settings 35 | datasets: 36 | train: &train 37 | type: LaneDataset 38 | parameters: 39 | S: *S 40 | dataset: tusimple 41 | split: train+val 42 | img_size: [*img_h, *img_w] 43 | max_lanes: *max_lanes 44 | normalize: false 45 | aug_chance: 1.0 46 | augmentations: 47 | - name: Affine 48 | parameters: 49 | translate_px: 50 | x: !!python/tuple [-25, 25] 51 | y: !!python/tuple [-10, 10] 52 | rotate: !!python/tuple [-6, 6] 53 | scale: !!python/tuple [0.85, 1.15] 54 | - name: HorizontalFlip 55 | parameters: 56 | p: 0.5 57 | 58 | root: "datasets/tusimple" 59 | 60 | test: 61 | type: LaneDataset 62 | parameters: 63 | S: *S 64 | dataset: tusimple 65 | split: test 66 | img_size: [*img_h, *img_w] 67 | max_lanes: *max_lanes 68 | normalize: false 69 | aug_chance: 0 70 | augmentations: 71 | root: "datasets/tusimple-test" 72 | 73 | val: 74 | type: LaneDataset 75 | parameters: 76 | S: *S 77 | dataset: tusimple 78 | split: val 79 | img_size: [*img_h, *img_w] 80 | max_lanes: *max_lanes 81 | normalize: false 82 | aug_chance: 0 83 | augmentations: 84 | root: "datasets/tusimple" 85 | -------------------------------------------------------------------------------- /cfgs/laneatt_culane_resnet34.yml: -------------------------------------------------------------------------------- 1 | # Model settings 2 | val_every: 1000 3 | model_checkpoint_interval: 1 4 | seed: 0 5 | model: 6 | name: LaneATT 7 | parameters: 8 | backbone: resnet34 9 | S: &S 72 10 | topk_anchors: 1000 11 | anchors_freq_path: '/home/LaneATT-main/data/culane_anchors_freq.pt' 12 | img_h: &img_h 360 13 | img_w: &img_w 640 14 | lane_num_types: &lane_num_types 5 15 | batch_size: 16 16 | epochs: 16 17 | loss_parameters: {} 18 | train_parameters: 19 | conf_threshold: 20 | nms_thres: 15. 21 | nms_topk: 3000 22 | test_parameters: 23 | conf_threshold: 0.5 24 | nms_thres: 50. 25 | nms_topk: &max_lanes 4 26 | optimizer: 27 | name: Adam 28 | parameters: 29 | lr: 0.0003 30 | lr_scheduler: 31 | name: CosineAnnealingLR 32 | parameters: 33 | T_max: 166650 # 15 * 11110 iterations 34 | 35 | # Dataset settings 36 | datasets: 37 | train: 38 | type: LaneDataset 39 | parameters: 40 | lane_num_types: *lane_num_types 41 | S: *S 42 | dataset: culane 43 | split: train 44 | img_size: [*img_h, *img_w] 45 | max_lanes: *max_lanes 46 | normalize: false 47 | aug_chance: 1.0 48 | augmentations: 49 | - name: Affine 50 | parameters: 51 | translate_px: 52 | x: !!python/tuple [-25, 25] 53 | y: !!python/tuple [-10, 10] 54 | rotate: !!python/tuple [-6, 6] 55 | scale: !!python/tuple [0.85, 1.15] 56 | - name: HorizontalFlip 57 | parameters: 58 | p: 0.5 59 | 60 | root: "/notebooks/CULane" 61 | 62 | test: 63 | type: LaneDataset 64 | parameters: 65 | S: *S 66 | dataset: culane 67 | split: test 68 | img_size: [*img_h, *img_w] 69 | max_lanes: *max_lanes 70 | normalize: false 71 | aug_chance: 0 72 | augmentations: 73 | root: "datasets/culane" 74 | 75 | val: 76 | type: LaneDataset 77 | parameters: 78 | S: *S 79 | dataset: culane 80 | split: val 81 | img_size: [*img_h, *img_w] 82 | max_lanes: *max_lanes 83 | normalize: false 84 | aug_chance: 0 85 | augmentations: 86 | root: "datasets/culane" 87 | -------------------------------------------------------------------------------- /cfgs/laneatt_ehl_wx_resnet34.yml: -------------------------------------------------------------------------------- 1 | # Model settings 2 | val_every: 10000000000000000000 3 | model_checkpoint_interval: 1 4 | seed: 0 5 | model: 6 | name: LaneATT 7 | parameters: 8 | backbone: resnet34 9 | S: &S 72 10 | topk_anchors: 1000 11 | anchors_freq_path: '/home/LaneATT-main/data/ehl_wx_anchors_freq.pt' 12 | img_h: &img_h 360 13 | img_w: &img_w 640 14 | lane_num_types: &lane_num_types 13 15 | batch_size: 16 16 | epochs: 60 17 | loss_parameters: {} 18 | train_parameters: 19 | conf_threshold: 20 | nms_thres: 15. 21 | nms_topk: 3000 22 | test_parameters: 23 | conf_threshold: 0.5 24 | nms_thres: 50. 25 | nms_topk: &max_lanes 10 26 | optimizer: 27 | name: Adam 28 | parameters: 29 | lr: 0.0003 30 | lr_scheduler: 31 | name: CosineAnnealingLR 32 | parameters: 33 | T_max: 30000 # 60 * 500 iterations 34 | 35 | # Dataset settings 36 | datasets: 37 | train: 38 | type: LaneDataset 39 | parameters: 40 | lane_num_types: *lane_num_types 41 | S: *S 42 | dataset: ehl_wx 43 | split: train 44 | img_size: [*img_h, *img_w] 45 | max_lanes: *max_lanes 46 | normalize: false 47 | aug_chance: 1.0 48 | augmentations: 49 | - name: Affine 50 | parameters: 51 | translate_px: 52 | x: !!python/tuple [-25, 25] 53 | y: !!python/tuple [-10, 10] 54 | rotate: !!python/tuple [-6, 6] 55 | scale: !!python/tuple [0.85, 1.15] 56 | - name: HorizontalFlip 57 | parameters: 58 | p: 0.5 59 | 60 | root: "/notebooks/ehl_wx/train_set" 61 | 62 | test: 63 | type: LaneDataset 64 | parameters: 65 | S: *S 66 | dataset: ehl_wx 67 | split: test 68 | img_size: [*img_h, *img_w] 69 | max_lanes: *max_lanes 70 | normalize: false 71 | aug_chance: 0 72 | augmentations: 73 | root: "E:\\ehl_wx" 74 | 75 | val: 76 | type: LaneDataset 77 | parameters: 78 | S: *S 79 | dataset: ehl_wx 80 | split: val 81 | img_size: [*img_h, *img_w] 82 | max_lanes: *max_lanes 83 | normalize: false 84 | aug_chance: 0 85 | augmentations: 86 | root: "E:\\ehl_wx" 87 | -------------------------------------------------------------------------------- /cfgs/laneatt_tusimple_resnet122.yml: -------------------------------------------------------------------------------- 1 | # Model settings 2 | val_every: 1000000 3 | model_checkpoint_interval: 1 4 | seed: 0 5 | model: 6 | name: LaneATT 7 | parameters: 8 | backbone: resnet122 9 | lane_num_types: &lane_num_types 5 10 | S: &S 72 11 | topk_anchors: 1000 12 | anchors_freq_path: 'data/tusimple_anchors_freq.pt' 13 | img_h: &img_h 360 14 | img_w: &img_w 640 15 | batch_size: 2 16 | epochs: 101 17 | loss_parameters: {} 18 | train_parameters: 19 | conf_threshold: 20 | nms_thres: 15. 21 | nms_topk: 3000 22 | test_parameters: 23 | conf_threshold: 0.2 24 | nms_thres: 45. 25 | nms_topk: &max_lanes 5 26 | optimizer: 27 | name: Adam 28 | parameters: 29 | lr: 0.0003 30 | lr_scheduler: 31 | name: CosineAnnealingLR 32 | parameters: 33 | T_max: 181300 # 100 * 1813 iterations 34 | 35 | # Dataset settings 36 | datasets: 37 | train: &train 38 | type: LaneDataset 39 | parameters: 40 | lane_num_types: *lane_num_types 41 | S: *S 42 | dataset: tusimple 43 | split: train+val 44 | img_size: [*img_h, *img_w] 45 | max_lanes: *max_lanes 46 | normalize: false 47 | aug_chance: 1.0 48 | augmentations: 49 | - name: Affine 50 | parameters: 51 | translate_px: 52 | x: !!python/tuple [-25, 25] 53 | y: !!python/tuple [-10, 10] 54 | rotate: !!python/tuple [-6, 6] 55 | scale: !!python/tuple [0.85, 1.15] 56 | - name: HorizontalFlip 57 | parameters: 58 | p: 0.5 59 | 60 | root: "/notebooks/tusimple/train_set" 61 | 62 | test: 63 | type: LaneDataset 64 | parameters: 65 | S: *S 66 | dataset: tusimple 67 | split: test 68 | img_size: [*img_h, *img_w] 69 | max_lanes: *max_lanes 70 | normalize: false 71 | aug_chance: 0 72 | augmentations: 73 | root: "/notebooks/tusimple/test_set" 74 | 75 | val: 76 | type: LaneDataset 77 | parameters: 78 | S: *S 79 | dataset: tusimple 80 | split: val 81 | img_size: [*img_h, *img_w] 82 | max_lanes: *max_lanes 83 | normalize: false 84 | aug_chance: 0 85 | augmentations: 86 | root: "datasets/tusimple" 87 | -------------------------------------------------------------------------------- /cfgs/laneatt_tusimple_resnet34.yml: -------------------------------------------------------------------------------- 1 | # Model settings 2 | val_every: 1000000 3 | model_checkpoint_interval: 1 4 | seed: 0 5 | model: 6 | name: LaneATT 7 | parameters: 8 | backbone: resnet34 9 | S: &S 72 10 | topk_anchors: 1000 11 | anchors_freq_path: '/home/LaneATT-main/data/tusimple_anchors_freq.pt' 12 | img_h: &img_h 360 13 | img_w: &img_w 640 14 | lane_num_types: &lane_num_types 5 15 | batch_size: 8 16 | epochs: 101 17 | loss_parameters: {} 18 | train_parameters: 19 | conf_threshold: 20 | nms_thres: 15. 21 | nms_topk: 3000 22 | test_parameters: 23 | conf_threshold: 0.2 24 | nms_thres: 45. 25 | nms_topk: &max_lanes 5 26 | optimizer: 27 | name: Adam 28 | parameters: 29 | lr: 0.0003 30 | lr_scheduler: 31 | name: CosineAnnealingLR 32 | parameters: 33 | T_max: 45400 # 100 * 454 iterations 34 | 35 | # Dataset settings 36 | datasets: 37 | train: &train 38 | type: LaneDataset 39 | parameters: 40 | lane_num_types: *lane_num_types 41 | S: *S 42 | dataset: tusimple 43 | split: train+val 44 | img_size: [*img_h, *img_w] 45 | max_lanes: *max_lanes 46 | normalize: false 47 | aug_chance: 1.0 48 | augmentations: 49 | - name: Affine 50 | parameters: 51 | translate_px: 52 | x: !!python/tuple [-25, 25] 53 | y: !!python/tuple [-10, 10] 54 | rotate: !!python/tuple [-6, 6] 55 | scale: !!python/tuple [0.85, 1.15] 56 | - name: HorizontalFlip 57 | parameters: 58 | p: 0.5 59 | 60 | root: "/notebooks/tusimple/train_set" 61 | 62 | test: 63 | type: LaneDataset 64 | parameters: 65 | S: *S 66 | dataset: tusimple 67 | split: test 68 | img_size: [*img_h, *img_w] 69 | max_lanes: *max_lanes 70 | normalize: false 71 | aug_chance: 0 72 | augmentations: 73 | root: "/notebooks/tusimple/test_set" 74 | 75 | val: 76 | type: LaneDataset 77 | parameters: 78 | S: *S 79 | dataset: tusimple 80 | split: val 81 | img_size: [*img_h, *img_w] 82 | max_lanes: *max_lanes 83 | normalize: false 84 | aug_chance: 0 85 | augmentations: 86 | root: "datasets/tusimple" 87 | -------------------------------------------------------------------------------- /utils/gen_anchor_mask.py: -------------------------------------------------------------------------------- 1 | import random 2 | import argparse 3 | 4 | import cv2 5 | import torch 6 | import numpy as np 7 | from tqdm import trange 8 | 9 | from lib.config import Config 10 | from lib.models.matching import match_proposals_with_targets 11 | 12 | 13 | def get_anchors_use_frequency(cfg, split='train', t_pos=15., t_neg=20.): 14 | model = cfg.get_model() 15 | anchors_frequency = torch.zeros(len(model.anchors), dtype=torch.int32) 16 | nb_unmatched_targets = 0 17 | dataset = cfg.get_dataset(split) 18 | for idx in trange(len(dataset)): 19 | _, targets, _ = dataset[idx] 20 | targets = targets[targets[:, 1] == 1] 21 | n_targets = len(targets) 22 | if n_targets == 0: 23 | continue 24 | targets = torch.tensor(targets) 25 | positives_mask, _, _, target_indices = match_proposals_with_targets(model, 26 | model.anchors, 27 | targets, 28 | t_pos=t_pos, 29 | t_neg=t_neg) 30 | n_matches = len(set(target_indices.tolist())) 31 | nb_unmatched_targets += n_targets - n_matches 32 | assert (n_targets - n_matches) >= 0 33 | 34 | anchors_frequency += positives_mask 35 | 36 | return anchors_frequency 37 | 38 | 39 | def save_mask(cfg_path, output_path): 40 | cfg = Config(cfg_path) 41 | frequency = get_anchors_use_frequency(cfg, split='train', t_pos=30., t_neg=35.) 42 | torch.save(frequency, output_path) 43 | 44 | 45 | def view_mask(): 46 | cfg = Config('config.yaml') 47 | model = cfg.get_model() 48 | img = model.draw_anchors(img_w=512, img_h=288) 49 | cv2.imshow('anchors', img) 50 | cv2.waitKey(0) 51 | 52 | 53 | def parse_args(): 54 | parser = argparse.ArgumentParser(description="Compute anchor frequency for later use as anchor mask") 55 | parser.add_argument("--output", help="Output path (e.g., `anchors_mask.pt`", required=True) 56 | parser.add_argument("--cfg", help="Config file (e.g., `config.yml`") 57 | args = parser.parse_args() 58 | 59 | return args 60 | 61 | 62 | if __name__ == '__main__': 63 | args = parse_args() 64 | # Fix seeds 65 | torch.manual_seed(0) 66 | np.random.seed(0) 67 | random.seed(0) 68 | torch.backends.cudnn.deterministic = True 69 | torch.backends.cudnn.benchmark = False 70 | 71 | save_mask(args.cfg, args.output) 72 | # view_mask() 73 | -------------------------------------------------------------------------------- /lib/export.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | from torchvision.transforms import ToTensor 4 | from lane import Lane 5 | from configoox import Config 6 | import torch 7 | import os 8 | os.environ["CUDA_VISIBLE_DEVICES"] = "6" 9 | 10 | 11 | #cfg = Config("/home/LaneATT-main/cfgs/laneatt_culane_resnet34.yml") 12 | cfg = Config("/home/LaneATT-main/cfgs/laneatt_ehl_wx_resnet34.yml") 13 | model_path = "/home/LaneATT-main/experiments/laneatt_r34_ehlwx/models/model_0014.pt" 14 | 15 | to_tensor = ToTensor() 16 | 17 | 18 | def get_img(img_path): 19 | img = cv2.resize(cv2.imread(img_path), (640, 360)) 20 | 21 | img = img / 255. 22 | # img = np.expand_dims(img, 0) 23 | img = to_tensor(img.astype(np.float32)) 24 | return img 25 | 26 | 27 | def get_epoch_model(model): 28 | return torch.load(model)['model'] 29 | 30 | 31 | model = cfg.get_model() 32 | #model.load_state_dict(get_epoch_model(model_path)) 33 | ####用ddp时保存模型用model.state_dict()(未用model.module.state_dict()) 导致加载时缺失”key“ 34 | model.load_state_dict({k.replace('module.', ''): v for k, v in torch.load(model_path)['model'].items()}) 35 | model = model.to(torch.device('cuda')) 36 | model.eval() 37 | test_parameters = cfg.get_test_parameters() 38 | 39 | img = torch.randn(1, 3, 360, 640) 40 | img = img.cuda() 41 | with torch.no_grad(): 42 | output1, ot2 = model(img, **test_parameters) 43 | #prediction = model.decode(output, as_lanes=True) 44 | 45 | 46 | # TorchScript export 47 | try: 48 | print('\nStarting TorchScript export with torch %s...' % torch.__version__) 49 | f = model_path.replace('.pt', '.torchscript.pt') # filename 50 | ts = torch.jit.trace(model, img) 51 | ts.save(f) 52 | print('TorchScript export success, saved as %s' % f) 53 | except Exception as e: 54 | print('TorchScript export failure: %s' % e) 55 | 56 | 57 | 58 | ####ONNX export -- 目前转出的onnx存在问题，修改多数网络层未解决 59 | """ 60 | try: 61 | import onnx 62 | from onnxsim import simplify 63 | 64 | print('\nStarting ONNX export with onnx %s...' % onnx.__version__) 65 | f = model_path.replace('.pt', '.onnx') # filename 66 | torch.onnx.export(model, img, f, verbose=True, opset_version=12, input_names=['images'], 67 | output_names=['output']) 68 | 69 | # Checks 70 | onnx_model = onnx.load(f) # load onnx model 71 | model_simp, check = simplify(onnx_model) 72 | assert check, "Simplified ONNX model could not be validated" 73 | onnx.save(model_simp, f) 74 | print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model 75 | print('ONNX export success, saved as %s' % f) 76 | except Exception as e: 77 | print('ONNX export failure: %s' % e) 78 | """ -------------------------------------------------------------------------------- /lib/nms/src/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 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import argparse 3 | 4 | import torch 5 | 6 | from lib.config import Config 7 | from lib.runner import Runner 8 | from lib.experiment import Experiment 9 | import os 10 | 11 | os.environ['CUDA_VISIBLE_DEVICES'] = "1" 12 | 13 | 14 | def parse_args(): 15 | parser = argparse.ArgumentParser(description="Train lane detector") 16 | parser.add_argument("--mode", choices=["train", "test"], default="train", help="Train or test?") 17 | parser.add_argument("--exp_name", default="laneatt_r34_culane", help="Experiment name") 18 | parser.add_argument("--cfg", default="cfgs/laneatt_culane_resnet34.yml", help="Config file") 19 | parser.add_argument("--resume", default=True, action="store_true", help="Resume training") 20 | parser.add_argument("--epoch", type=int, help="Epoch to test the model on") 21 | parser.add_argument("--cpu", action="store_true", help="(Unsupported) Use CPU instead of GPU") 22 | parser.add_argument("--save_predictions", action="store_true", help="Save predictions to pickle file") 23 | parser.add_argument("--view", choices=["all", "mistakes"], help="Show predictions") 24 | parser.add_argument("--deterministic", 25 | action="store_true", 26 | help="set cudnn.deterministic = True and cudnn.benchmark = False") 27 | args = parser.parse_args() 28 | if args.cfg is None and args.mode == "train": 29 | raise Exception("If you are training, you have to set a config file using --cfg /path/to/your/config.yaml") 30 | if args.resume and args.mode == "test": 31 | raise Exception("args.resume is set on `test` mode: can't resume testing") 32 | if args.epoch is not None and args.mode == 'train': 33 | raise Exception("The `epoch` parameter should not be set when training") 34 | if args.view is not None and args.mode != "test": 35 | raise Exception('Visualization is only available during evaluation') 36 | if args.cpu: 37 | raise Exception("CPU training/testing is not supported: the NMS procedure is only implemented for CUDA") 38 | 39 | return args 40 | 41 | 42 | def main(): 43 | args = parse_args() 44 | exp = Experiment(args.exp_name, args, mode=args.mode) 45 | if args.cfg is None: 46 | cfg_path = exp.cfg_path 47 | else: 48 | cfg_path = args.cfg 49 | cfg = Config(cfg_path) 50 | exp.set_cfg(cfg, override=False) 51 | device = torch.device('cpu') if not torch.cuda.is_available() or args.cpu else torch.device('cuda') 52 | runner = Runner(cfg, exp, device, view=args.view, resume=args.resume, deterministic=args.deterministic) 53 | if args.mode == 'train': 54 | try: 55 | runner.train() 56 | except KeyboardInterrupt: 57 | logging.info('Training interrupted.') 58 | runner.eval(epoch=args.epoch or exp.get_last_checkpoint_epoch(), save_predictions=args.save_predictions) 59 | 60 | 61 | if __name__ == '__main__': 62 | main() 63 | -------------------------------------------------------------------------------- /lib/nms/LICENSE: -------------------------------------------------------------------------------- 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 | 31 | THIRD-PARTY SOFTWARE NOTICES AND INFORMATION 32 | 33 | This project incorporates material from the project(s) 34 | listed below (collectively, "Third Party Code"). This Third Party Code is 35 | licensed to you under their original license terms set forth below. 36 | 37 | 1. Faster R-CNN, (https://github.com/rbgirshick/py-faster-rcnn) 38 | 39 | The MIT License (MIT) 40 | 41 | Copyright (c) 2015 Microsoft Corporation 42 | 43 | Permission is hereby granted, free of charge, to any person obtaining a copy 44 | of this software and associated documentation files (the "Software"), to deal 45 | in the Software without restriction, including without limitation the rights 46 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 47 | copies of the Software, and to permit persons to whom the Software is 48 | furnished to do so, subject to the following conditions: 49 | 50 | The above copyright notice and this permission notice shall be included in 51 | all copies or substantial portions of the Software. 52 | 53 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 54 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 55 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 56 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 57 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 58 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 59 | THE SOFTWARE. 60 | 61 | -------------------------------------------------------------------------------- /utils/gen_video.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import argparse 3 | 4 | import cv2 5 | import numpy as np 6 | from tqdm import tqdm 7 | 8 | from lib.config import Config 9 | 10 | 11 | def parse_args(): 12 | parser = argparse.ArgumentParser(description="Tool to generate qualitative results videos") 13 | parser.add_argument("--pred", help=".pkl file to load predictions from", required=True) 14 | parser.add_argument("--cfg", default="config.yaml", help="Config file") 15 | parser.add_argument("--cover", default="tusimple_cover.png", help="Cover image file") 16 | parser.add_argument("--out", default="video.avi", help="Output filename") 17 | parser.add_argument("--view", action="store_true", help="Show predictions instead of creating video") 18 | parser.add_argument("--length", type=int, help="Length of the output video (seconds)") 19 | parser.add_argument("--clips", type=int, help="Number of clips") 20 | parser.add_argument("--fps", default=5, type=int, help="Video FPS") 21 | parser.add_argument("--legend", help="Path to legend image file") 22 | 23 | return parser.parse_args() 24 | 25 | 26 | def add_cover_img(video, cover_path, frames=90): 27 | cover = cv2.imread(cover_path) 28 | for _ in range(frames): 29 | video.write(cover) 30 | 31 | 32 | def create_video(filename, width, height, fps=5): 33 | fourcc = cv2.VideoWriter_fourcc(*'MP42') 34 | video = cv2.VideoWriter(filename, fourcc, float(fps), (width, height)) 35 | 36 | return video 37 | 38 | 39 | def main(): 40 | np.random.seed(0) 41 | args = parse_args() 42 | cfg = Config(args.cfg) 43 | print('Loading dataset...') 44 | dataset = cfg.get_dataset('test') 45 | print('Done') 46 | height, width = cfg['datasets']['test']['parameters']['img_size'] 47 | print('Using resolution {}x{}'.format(width, height)) 48 | legend = cv2.imread(args.legend) if args.legend else None 49 | if not args.view: 50 | video = create_video(args.out, width, height + legend.shape[0] if legend is not None else 0, args.fps) 51 | 52 | print('Loading predictions...') 53 | with open(args.pred, "rb") as pred_file: 54 | predictions = np.array(pickle.load(pred_file)) 55 | print("Done.") 56 | 57 | if args.length is not None and args.clips is not None: 58 | video_length = args.length * args.fps 59 | assert video_length % args.clips == 0 60 | clip_length = video_length // args.clips 61 | all_clip_ids = np.arange(len(dataset) // clip_length) 62 | selected_clip_ids = np.random.choice(all_clip_ids, size=args.clips, replace=False) 63 | frame_idxs = (np.repeat(selected_clip_ids, clip_length).reshape(args.clips, clip_length) + np.arange( 64 | clip_length)).flatten() 65 | total = len(frame_idxs) 66 | else: 67 | total = len(dataset) 68 | frame_idxs = np.arange(len(dataset)) 69 | 70 | for idx, pred in tqdm(zip(frame_idxs, predictions[frame_idxs]), total=total): 71 | frame, _, _ = dataset.draw_annotation(idx, pred=pred) 72 | assert frame.shape[:2] == (height, width) 73 | if legend is not None: 74 | frame = np.vstack((legend, frame)) 75 | if args.view: 76 | cv2.imshow('frame', frame) 77 | cv2.waitKey(0) 78 | else: 79 | video.write(frame) 80 | 81 | if not args.view: 82 | video.release() 83 | print('Video saved as {}'.format(args.out)) 84 | 85 | 86 | if __name__ == '__main__': 87 | main() 88 | -------------------------------------------------------------------------------- /lib/models/matching.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import pdb 3 | 4 | INFINITY = 987654. 5 | 6 | 7 | def match_proposals_with_targets(model, proposals, targets, t_pos=15., t_neg=20.): 8 | # repeat proposals and targets to generate all combinations 9 | num_proposals = proposals.shape[0] 10 | num_targets = targets.shape[0] 11 | # pad proposals and target for the valid_offset_mask's trick 12 | proposals_pad = proposals.new_zeros(proposals.shape[0], proposals.shape[1] + 1) 13 | proposals_pad[:, :-1] = proposals 14 | proposals = proposals_pad 15 | targets_pad = targets.new_zeros(targets.shape[0], targets.shape[1] + 1) 16 | targets_pad[:, :-1] = targets 17 | targets = targets_pad 18 | 19 | proposals = torch.repeat_interleave(proposals, num_targets, 20 | dim=0) # repeat_interleave'ing [a, b] 2 times gives [a, a, b, b] 21 | 22 | targets = torch.cat(num_proposals * [targets]) # applying this 2 times on [c, d] gives [c, d, c, d] 23 | 24 | # get start and the intersection of offsets 25 | 26 | ############## multly lane 27 | targets_starts = targets[:, model.num_lane_type] * model.n_strips 28 | proposals_starts = proposals[:, model.num_lane_type] * model.n_strips 29 | starts = torch.max(targets_starts, proposals_starts).round().long() 30 | ends = (targets_starts + targets[:, model.num_lane_type+2] - 1).round().long() 31 | lengths = ends - starts + 1 32 | ends[lengths < 0] = starts[lengths < 0] - 1 33 | lengths[lengths < 0] = 0 # a negative number here means no intersection, 34 | valid_offsets_mask = targets.new_zeros(targets.shape) 35 | all_indices = torch.arange(valid_offsets_mask.shape[0], dtype=torch.long, device=targets.device) 36 | valid_offsets_mask[all_indices, model.num_lane_type + 3 + starts] = 1. 37 | valid_offsets_mask[all_indices, model.num_lane_type + 3 + ends + 1] -= 1. 38 | 39 | ## put a -1 on the `end` index, giving [0, 1, 0, -1, 0] 40 | ## if lenght is zero, the previous line would put a one where it shouldnt be. 41 | ## this -=1 (instead of =-1) fixes this 42 | ## the cumsum gives [0, 1, 1, 0, 0], the correct mask for the offsets 43 | valid_offsets_mask = valid_offsets_mask.cumsum(dim=1) != 0. 44 | invalid_offsets_mask = ~valid_offsets_mask 45 | 46 | # compute distances 47 | # this compares [ac, ad, bc, bd], i.e., all combinations 48 | distances = torch.abs((targets - proposals) * valid_offsets_mask.float()).sum(dim=1) / (lengths.float() + 1e-9 49 | ) # avoid division by zero 50 | distances[lengths == 0] = INFINITY 51 | invalid_offsets_mask = invalid_offsets_mask.view(num_proposals, num_targets, invalid_offsets_mask.shape[1]) 52 | distances = distances.view(num_proposals, num_targets) # d[i,j] = distance from proposal i to target j 53 | 54 | positives = distances.min(dim=1)[0] < t_pos # all arget lane 中对应满足条件的 proposal 中的一条 size : proposals lins, ->bool 55 | negatives = distances.min(dim=1)[0] > t_neg 56 | 57 | if positives.sum() == 0: 58 | target_positives_indices = torch.tensor([], device=positives.device, dtype=torch.long) 59 | else: 60 | target_positives_indices = distances[positives].argmin(dim=1) # proposal lane 对应的 target lane索引 61 | invalid_offsets_mask = invalid_offsets_mask[positives, target_positives_indices] 62 | 63 | return positives, invalid_offsets_mask[:, :-1], negatives, target_positives_indices 64 | -------------------------------------------------------------------------------- /lib/infer.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | from torchvision.transforms import ToTensor 4 | from lane import Lane 5 | from configoox import Config 6 | import torch 7 | 8 | 9 | cfg = Config("/home/LaneATT-main/cfgs/laneatt_tusimple_resnet34.yml") 10 | model_path = "/home/LaneATT-main/experiments/laneatt_r34_tusimple/models/model_0009.pt" 11 | 12 | to_tensor = ToTensor() 13 | 14 | 15 | 16 | GT_COLOR = (255, 0, 0) 17 | PRED_HIT_COLOR = (0, 255, 0) 18 | PRED_MISS_COLOR = (0, 0, 255) 19 | IMAGENET_MEAN = np.array([0.485, 0.456, 0.406]) 20 | IMAGENET_STD = np.array([0.229, 0.224, 0.225]) 21 | 22 | 23 | def get_img(img_path): 24 | img = cv2.resize(cv2.imread(img_path), (640, 360)) 25 | 26 | img = img / 255. 27 | # img = np.expand_dims(img, 0) 28 | img = to_tensor(img.astype(np.float32)) 29 | return img 30 | 31 | 32 | def get_epoch_model(model): 33 | return torch.load(model)['model'] 34 | 35 | 36 | def draw_annotation(pred=None, img=None): 37 | 38 | #img = cv2.resize(img, (self.img_w, self.img_h)) 39 | 40 | img_h, _, _ = img.shape 41 | # Pad image to visualize extrapolated predictions 42 | pad = 0 43 | if pad > 0: 44 | img_pad = np.zeros((self.img_h + 2 * pad, self.img_w + 2 * pad, 3), dtype=np.uint8) 45 | img_pad[pad:-pad, pad:-pad, :] = img 46 | img = img_pad 47 | 48 | for i, l in enumerate(pred): 49 | if l.metadata['type'] == 0: 50 | continue 51 | 52 | color = PRED_MISS_COLOR 53 | points = l.points 54 | points[:, 0] *= img.shape[1] 55 | points[:, 1] *= img.shape[0] 56 | points = points.round().astype(int) 57 | points += pad 58 | xs, ys = points[:, 0], points[:, 1] 59 | for curr_p, next_p in zip(points[:-1], points[1:]): 60 | img = cv2.line(img, 61 | tuple(curr_p), 62 | tuple(next_p), 63 | color=color, 64 | thickness=3) 65 | if 'start_x' in l.metadata: 66 | start_x = l.metadata['start_x'] * img.shape[1] 67 | start_y = l.metadata['start_y'] * img.shape[0] 68 | cv2.circle(img, (int(start_x + pad), int(img_h - 1 - start_y + pad)), 69 | radius=5, 70 | color=(0, 0, 255), 71 | thickness=-1) 72 | if len(xs) == 0: 73 | print("Empty pred") 74 | if len(xs) > 0: 75 | cv2.putText(img, 76 | '{}-{:.0f}'.format(l.metadata['type'],l.metadata['conf'] * 100), 77 | (int(xs[len(xs) // 2] + pad), int(ys[len(xs) // 2] + pad - 50)+(i+1)*20), 78 | fontFace=cv2.FONT_HERSHEY_COMPLEX, 79 | fontScale=0.7, 80 | color=(255, 0, 255)) 81 | return img 82 | 83 | 84 | 85 | def eval(image_path): 86 | model = cfg.get_model() 87 | model.load_state_dict(get_epoch_model(model_path)) 88 | model = model.to(torch.device('cuda')) 89 | model.eval() 90 | test_parameters = cfg.get_test_parameters() 91 | with torch.no_grad(): 92 | image0 = get_img(image_path) 93 | image = image0.unsqueeze(0) 94 | image = image.to(torch.device('cuda')) 95 | output = model(image, **test_parameters) 96 | prediction = model.decode3(output, as_lanes=True) 97 | img = (image0.cpu().permute(1, 2, 0).numpy() * 255).astype(np.uint8) 98 | img = draw_annotation(img=img, pred=prediction[0]) 99 | cv2.imwrite("nbv.jpg", img) 100 | #cv2.imshow('pred', img) 101 | #cv2.waitKey(0) 102 | kk = '/home/LaneATT-main/data/figures/1.png' 103 | eval(kk) -------------------------------------------------------------------------------- /main_distributed.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import argparse 3 | 4 | import torch 5 | 6 | from lib.config import Config 7 | from lib.runner import Runner 8 | from lib.runner_distributed import Runner_Distributed 9 | from lib.experiment import Experiment 10 | from lib.distributed_utils import * 11 | 12 | 13 | 14 | def parse_args(): 15 | parser = argparse.ArgumentParser(description="Train lane detector") 16 | parser.add_argument("--mode", choices=["train", "test"], default="train", help="Train or test?") 17 | parser.add_argument("--exp_name", default="laneatt_r34_ehlwx", help="Experiment name") 18 | parser.add_argument("--cfg", default="cfgs/laneatt_ehl_wx_resnet34.yml", help="Config file") 19 | parser.add_argument("--resume", default=True, action="store_true", help="Resume training") 20 | parser.add_argument("--epoch", type=int, help="Epoch to test the model on") 21 | parser.add_argument("--cpu", action="store_true", help="(Unsupported) Use CPU instead of GPU") 22 | parser.add_argument("--save_predictions", action="store_true", help="Save predictions to pickle file") 23 | parser.add_argument("--view", choices=["all", "mistakes"], help="Show predictions") 24 | parser.add_argument("--deterministic", 25 | action="store_true", 26 | help="set cudnn.deterministic = True and cudnn.benchmark = False") 27 | 28 | parser.add_argument('--sync_bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode') 29 | 30 | parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training') 31 | 32 | # 不要改该参数，系统会自动分配 33 | parser.add_argument('--device', default='cuda', help='device id (i.e. 0 or 0,1 or cpu)') 34 | # 开启的进程数(注意不是线程),不用设置该参数，会根据nproc_per_node自动设置 35 | parser.add_argument('--world-size', default=4, type=int, 36 | help='number of distributed processes') 37 | parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training') 38 | 39 | args = parser.parse_args() 40 | if args.cfg is None and args.mode == "train": 41 | raise Exception("If you are training, you have to set a config file using --cfg /path/to/your/config.yaml") 42 | if args.resume and args.mode == "test": 43 | raise Exception("args.resume is set on `test` mode: can't resume testing") 44 | if args.epoch is not None and args.mode == 'train': 45 | raise Exception("The `epoch` parameter should not be set when training") 46 | if args.view is not None and args.mode != "test": 47 | raise Exception('Visualization is only available during evaluation') 48 | if args.cpu: 49 | raise Exception("CPU training/testing is not supported: the NMS procedure is only implemented for CUDA") 50 | 51 | return args 52 | 53 | 54 | def main(): 55 | args = parse_args() 56 | exp = Experiment(args.exp_name, args, mode=args.mode) 57 | if args.cfg is None: 58 | cfg_path = exp.cfg_path 59 | else: 60 | cfg_path = args.cfg 61 | cfg = Config(cfg_path) 62 | 63 | ###分布式训练 根据world_size 成比例增加lr 64 | init_distributed_mode(args) 65 | cfg['optimizer']['parameters']['lr'] *= args.world_size 66 | exp.set_cfg(cfg, override=False) 67 | 68 | device = torch.device('cpu') if not torch.cuda.is_available() or args.cpu else torch.device('cuda') 69 | 70 | runner = Runner_Distributed(cfg, exp, device, view=args.view, resume=args.resume, deterministic=args.deterministic, 71 | device_id=args.gpu, syncBN=args.sync_bn, distributed=args.distributed, main_pod=args.rank) 72 | if args.mode == 'train': 73 | try: 74 | runner.train() 75 | except KeyboardInterrupt: 76 | logging.info('Training interrupted.') 77 | #runner.eval(epoch=args.epoch or exp.get_last_checkpoint_epoch(), save_predictions=args.save_predictions) 78 | 79 | 80 | if __name__ == '__main__': 81 | main() 82 | -------------------------------------------------------------------------------- /lib/datasets/ehl_wx.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import random 4 | import logging 5 | 6 | import numpy as np 7 | 8 | from utils.tusimple_metric import LaneEval 9 | 10 | from .lane_dataset_loader import LaneDatasetLoader 11 | 12 | SPLIT_FILES = { 13 | 'train+val': [], 14 | 'train': ['ehl_wuxi.json'], 15 | 'val': [], 16 | 'test': [], 17 | } 18 | 19 | 20 | class EHLWX(LaneDatasetLoader): 21 | def __init__(self, split='train', max_lanes=None, root=None): 22 | self.split = split 23 | self.root = root 24 | self.logger = logging.getLogger(__name__) 25 | 26 | if split not in SPLIT_FILES.keys(): 27 | raise Exception('Split `{}` does not exist.'.format(split)) 28 | 29 | self.anno_files = [os.path.join(self.root, path) for path in SPLIT_FILES[split]] 30 | 31 | if root is None: 32 | raise Exception('Please specify the root directory') 33 | 34 | self.img_w, self.img_h = 1920, 1080 35 | self.annotations = [] 36 | self.load_annotations() 37 | 38 | # Force max_lanes, used when evaluating testing with models trained on other datasets 39 | if max_lanes is not None: 40 | self.max_lanes = max_lanes 41 | 42 | def get_img_heigth(self): 43 | return 1080 44 | 45 | def get_img_width(self): 46 | return 1920 47 | 48 | def get_metrics(self, lanes, idx): 49 | label = self.annotations[idx] 50 | org_anno = label['old_anno'] 51 | pred = self.pred2lanes(org_anno['path'], lanes, org_anno['y_samples']) 52 | _, fp, fn, matches, accs, _ = LaneEval.bench(pred, org_anno['org_lanes'], org_anno['y_samples'], 0, True) 53 | return fp, fn, matches, accs 54 | 55 | def pred2lanes(self, path, pred, y_samples): 56 | ys = np.array(y_samples) / self.img_h 57 | lanes = [] 58 | for lane in pred: 59 | xs = lane(ys) 60 | invalid_mask = xs < 0 61 | lane = (xs * self.get_img_width(path)).astype(int) 62 | lane[invalid_mask] = -2 63 | lanes.append(lane.tolist()) 64 | 65 | return lanes 66 | 67 | def load_annotations(self): 68 | self.logger.info('Loading EHLWX annotations...') 69 | self.annotations = [] 70 | max_lanes = 0 71 | for anno_file in self.anno_files: 72 | with open(anno_file, 'r') as anno_obj: 73 | lines = anno_obj.readlines() 74 | for line in lines: 75 | data = json.loads(line) 76 | y_samples = data['h_samples'] 77 | gt_lanes = data['lanes'] 78 | lanes_type = data['lane_type'] 79 | lanes = [[(x, y) for (x, y) in zip(lane, y_samples) if x >= 0] for lane in gt_lanes] 80 | lanes = [lane for lane in lanes if len(lane) > 0] 81 | if len(lanes) > 0: 82 | assert len(lanes) == len(lanes_type) 83 | max_lanes = max(max_lanes, len(lanes)) 84 | self.annotations.append({ 85 | 'path': os.path.join(self.root, data['raw_file'].replace("\\", "/")), 86 | 'org_path': data['raw_file'], 87 | 'org_lanes': gt_lanes, 88 | 'lanes': lanes, 89 | 'lane_type':lanes_type, 90 | 'aug': False, 91 | 'y_samples': y_samples 92 | }) 93 | 94 | if self.split == 'train': 95 | random.shuffle(self.annotations) 96 | self.max_lanes = max_lanes 97 | self.logger.info('%d annotations loaded, with a maximum of %d lanes in an image.', len(self.annotations), 98 | self.max_lanes) 99 | 100 | def transform_annotations(self, transform): 101 | self.annotations = list(map(transform, self.annotations)) 102 | 103 | def pred2tusimpleformat(self, idx, pred, runtime): 104 | runtime *= 1000. # s to ms 105 | img_name = self.annotations[idx]['old_anno']['org_path'] 106 | h_samples = self.annotations[idx]['old_anno']['y_samples'] 107 | lanes = self.pred2lanes(img_name, pred, h_samples) 108 | output = {'raw_file': img_name, 'lanes': lanes, 'run_time': runtime} 109 | return json.dumps(output) 110 | 111 | def save_tusimple_predictions(self, predictions, filename, runtimes=None): 112 | if runtimes is None: 113 | runtimes = np.ones(len(predictions)) * 1.e-3 114 | lines = [] 115 | for idx, (prediction, runtime) in enumerate(zip(predictions, runtimes)): 116 | line = self.pred2tusimpleformat(idx, prediction, runtime) 117 | lines.append(line) 118 | with open(filename, 'w') as output_file: 119 | output_file.write('\n'.join(lines)) 120 | 121 | def eval_predictions(self, predictions, output_basedir, runtimes=None): 122 | pred_filename = os.path.join(output_basedir, 'tusimple_predictions.json') 123 | self.save_tusimple_predictions(predictions, pred_filename, runtimes) 124 | result = json.loads(LaneEval.bench_one_submit(pred_filename, self.anno_files[0])) 125 | table = {} 126 | for metric in result: 127 | table[metric['name']] = metric['value'] 128 | 129 | return table 130 | 131 | def __getitem__(self, idx): 132 | return self.annotations[idx] 133 | 134 | def __len__(self): 135 | return len(self.annotations) 136 | -------------------------------------------------------------------------------- /lib/datasets/tusimple.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import random 4 | import logging 5 | 6 | import numpy as np 7 | 8 | from utils.tusimple_metric import LaneEval 9 | 10 | from .lane_dataset_loader import LaneDatasetLoader 11 | 12 | SPLIT_FILES = { 13 | 'train+val': ['label_data_0313_new.json', 'label_data_0601_new.json', 'label_data_0531_new.json'], 14 | 'train': ['label_data_0313.json_new', 'label_data_0601_new.json'], 15 | 'val': ['label_data_0531_new.json'], 16 | 'test': ['test_tasks_0627_new.json'], 17 | } 18 | 19 | 20 | class TuSimple(LaneDatasetLoader): 21 | def __init__(self, split='train', max_lanes=None, root=None): 22 | self.split = split 23 | self.root = root 24 | self.logger = logging.getLogger(__name__) 25 | 26 | if split not in SPLIT_FILES.keys(): 27 | raise Exception('Split `{}` does not exist.'.format(split)) 28 | 29 | self.anno_files = [os.path.join(self.root, path) for path in SPLIT_FILES[split]] 30 | 31 | if root is None: 32 | raise Exception('Please specify the root directory') 33 | 34 | self.img_w, self.img_h = 1280, 720 35 | self.annotations = [] 36 | self.load_annotations() 37 | 38 | # Force max_lanes, used when evaluating testing with models trained on other datasets 39 | if max_lanes is not None: 40 | self.max_lanes = max_lanes 41 | 42 | def get_img_heigth(self, _): 43 | return 720 44 | 45 | def get_img_width(self, _): 46 | return 1280 47 | 48 | def get_metrics(self, lanes, idx): 49 | label = self.annotations[idx] 50 | org_anno = label['old_anno'] 51 | pred = self.pred2lanes(org_anno['path'], lanes, org_anno['y_samples']) 52 | _, fp, fn, matches, accs, _ = LaneEval.bench(pred, org_anno['org_lanes'], org_anno['y_samples'], 0, True) 53 | return fp, fn, matches, accs 54 | 55 | def pred2lanes(self, path, pred, y_samples): 56 | ys = np.array(y_samples) / self.img_h 57 | lanes = [] 58 | for lane in pred: 59 | xs = lane(ys) 60 | invalid_mask = xs < 0 61 | lane = (xs * self.get_img_width(path)).astype(int) 62 | lane[invalid_mask] = -2 63 | lanes.append(lane.tolist()) 64 | 65 | return lanes 66 | 67 | def load_annotations(self): 68 | self.logger.info('Loading TuSimple annotations...') 69 | self.annotations = [] 70 | max_lanes = 0 71 | for anno_file in self.anno_files: 72 | with open(anno_file, 'r') as anno_obj: 73 | lines = anno_obj.readlines() 74 | for line in lines: 75 | data = json.loads(line) 76 | y_samples = data['h_samples'] 77 | gt_lanes = data['lanes'] 78 | lanes_type = data['lane_type'] 79 | lanes = [[(x, y) for (x, y) in zip(lane, y_samples) if x >= 0] for lane in gt_lanes] 80 | lanes = [lane for lane in lanes if len(lane) > 0] 81 | if len(lanes) > 0: 82 | assert len(lanes) == len(lanes_type) 83 | max_lanes = max(max_lanes, len(lanes)) 84 | self.annotations.append({ 85 | 'path': os.path.join(self.root, data['raw_file']), 86 | 'org_path': data['raw_file'], 87 | 'org_lanes': gt_lanes, 88 | 'lanes': lanes, 89 | 'lane_type':lanes_type, 90 | 'aug': False, 91 | 'y_samples': y_samples 92 | }) 93 | 94 | if self.split == 'train': 95 | random.shuffle(self.annotations) 96 | self.max_lanes = max_lanes 97 | self.logger.info('%d annotations loaded, with a maximum of %d lanes in an image.', len(self.annotations), 98 | self.max_lanes) 99 | 100 | def transform_annotations(self, transform): 101 | self.annotations = list(map(transform, self.annotations)) 102 | 103 | def pred2tusimpleformat(self, idx, pred, runtime): 104 | runtime *= 1000. # s to ms 105 | img_name = self.annotations[idx]['old_anno']['org_path'] 106 | h_samples = self.annotations[idx]['old_anno']['y_samples'] 107 | lanes = self.pred2lanes(img_name, pred, h_samples) 108 | output = {'raw_file': img_name, 'lanes': lanes, 'run_time': runtime} 109 | return json.dumps(output) 110 | 111 | def save_tusimple_predictions(self, predictions, filename, runtimes=None): 112 | if runtimes is None: 113 | runtimes = np.ones(len(predictions)) * 1.e-3 114 | lines = [] 115 | for idx, (prediction, runtime) in enumerate(zip(predictions, runtimes)): 116 | line = self.pred2tusimpleformat(idx, prediction, runtime) 117 | lines.append(line) 118 | with open(filename, 'w') as output_file: 119 | output_file.write('\n'.join(lines)) 120 | 121 | def eval_predictions(self, predictions, output_basedir, runtimes=None): 122 | pred_filename = os.path.join(output_basedir, 'tusimple_predictions.json') 123 | self.save_tusimple_predictions(predictions, pred_filename, runtimes) 124 | result = json.loads(LaneEval.bench_one_submit(pred_filename, self.anno_files[0])) 125 | table = {} 126 | for metric in result: 127 | table[metric['name']] = metric['value'] 128 | 129 | return table 130 | 131 | def __getitem__(self, idx): 132 | return self.annotations[idx] 133 | 134 | def __len__(self): 135 | return len(self.annotations) 136 | -------------------------------------------------------------------------------- /utils/tusimple_metric.py: -------------------------------------------------------------------------------- 1 | # pylint: disable-all 2 | import numpy as np 3 | import ujson as json 4 | from sklearn.linear_model import LinearRegression 5 | 6 | 7 | class LaneEval(object): 8 | lr = LinearRegression() 9 | pixel_thresh = 20 10 | pt_thresh = 0.85 11 | 12 | @staticmethod 13 | def get_angle(xs, y_samples): 14 | xs, ys = xs[xs >= 0], y_samples[xs >= 0] 15 | if len(xs) > 1: 16 | LaneEval.lr.fit(ys[:, None], xs) 17 | k = LaneEval.lr.coef_[0] 18 | theta = np.arctan(k) 19 | else: 20 | theta = 0 21 | return theta 22 | 23 | @staticmethod 24 | def line_accuracy(pred, gt, thresh): 25 | pred = np.array([p if p >= 0 else -100 for p in pred]) 26 | gt = np.array([g if g >= 0 else -100 for g in gt]) 27 | return np.sum(np.where(np.abs(pred - gt) < thresh, 1., 0.)) / len(gt) 28 | 29 | @staticmethod 30 | def distances(pred, gt): 31 | return np.abs(pred - gt) 32 | 33 | @staticmethod 34 | def bench(pred, gt, y_samples, running_time, get_matches=False): 35 | if any(len(p) != len(y_samples) for p in pred): 36 | raise Exception('Format of lanes error.') 37 | if running_time > 20000 or len(gt) + 2 < len(pred): 38 | if get_matches: 39 | return 0., 0., 1., [False] * len(pred), [0] * len(pred), [None] * len(pred) 40 | return 0., 0., 1., 41 | angles = [LaneEval.get_angle(np.array(x_gts), np.array(y_samples)) for x_gts in gt] 42 | threshs = [LaneEval.pixel_thresh / np.cos(angle) for angle in angles] 43 | line_accs = [] 44 | fp, fn = 0., 0. 45 | matched = 0. 46 | my_matches = [False] * len(pred) 47 | my_accs = [0] * len(pred) 48 | my_dists = [None] * len(pred) 49 | for x_gts, thresh in zip(gt, threshs): 50 | accs = [LaneEval.line_accuracy(np.array(x_preds), np.array(x_gts), thresh) for x_preds in pred] 51 | my_accs = np.maximum(my_accs, accs) 52 | max_acc = np.max(accs) if len(accs) > 0 else 0. 53 | my_dist = [LaneEval.distances(np.array(x_preds), np.array(x_gts)) for x_preds in pred] 54 | if len(accs) > 0: 55 | my_dists[np.argmax(accs)] = { 56 | 'y_gts': list(np.array(y_samples)[np.array(x_gts) >= 0].astype(int)), 57 | 'dists': list(my_dist[np.argmax(accs)]) 58 | } 59 | 60 | if max_acc < LaneEval.pt_thresh: 61 | fn += 1 62 | else: 63 | my_matches[np.argmax(accs)] = True 64 | matched += 1 65 | line_accs.append(max_acc) 66 | fp = len(pred) - matched 67 | if len(gt) > 4 and fn > 0: 68 | fn -= 1 69 | s = sum(line_accs) 70 | if len(gt) > 4: 71 | s -= min(line_accs) 72 | if get_matches: 73 | return s / max(min(4.0, len(gt)), 1.), fp / len(pred) if len(pred) > 0 else 0., fn / max( 74 | min(len(gt), 4.), 1.), my_matches, my_accs, my_dists 75 | 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.) 76 | 77 | @staticmethod 78 | def bench_one_submit(pred_file, gt_file): 79 | try: 80 | json_pred = [json.loads(line) for line in open(pred_file).readlines()] 81 | except BaseException as e: 82 | raise Exception('Fail to load json file of the prediction.') 83 | json_gt = [json.loads(line) for line in open(gt_file).readlines()] 84 | if len(json_gt) != len(json_pred): 85 | raise Exception('We do not get the predictions of all the test tasks') 86 | gts = {img['raw_file']: img for img in json_gt} 87 | accuracy, fp, fn = 0., 0., 0. 88 | run_times = [] 89 | for pred in json_pred: 90 | if 'raw_file' not in pred or 'lanes' not in pred or 'run_time' not in pred: 91 | raise Exception('raw_file or lanes or run_time not in some predictions.') 92 | raw_file = pred['raw_file'] 93 | pred_lanes = pred['lanes'] 94 | run_time = pred['run_time'] 95 | run_times.append(run_time) 96 | if raw_file not in gts: 97 | raise Exception('Some raw_file from your predictions do not exist in the test tasks.') 98 | gt = gts[raw_file] 99 | gt_lanes = gt['lanes'] 100 | y_samples = gt['h_samples'] 101 | try: 102 | a, p, n = LaneEval.bench(pred_lanes, gt_lanes, y_samples, run_time) 103 | except BaseException as e: 104 | raise Exception('Format of lanes error.') 105 | accuracy += a 106 | fp += p 107 | fn += n 108 | num = len(gts) 109 | # the first return parameter is the default ranking parameter 110 | return json.dumps([{ 111 | 'name': 'Accuracy', 112 | 'value': accuracy / num, 113 | 'order': 'desc' 114 | }, { 115 | 'name': 'FP', 116 | 'value': fp / num, 117 | 'order': 'asc' 118 | }, { 119 | 'name': 'FN', 120 | 'value': fn / num, 121 | 'order': 'asc' 122 | }, { 123 | 'name': 'FPS', 124 | 'value': 1000. / np.mean(run_times) 125 | }]) 126 | 127 | 128 | if __name__ == '__main__': 129 | import sys 130 | 131 | try: 132 | if len(sys.argv) != 3: 133 | raise Exception('Invalid input arguments') 134 | print(LaneEval.bench_one_submit(sys.argv[1], sys.argv[2])) 135 | except Exception as e: 136 | print(e) 137 | # sys.exit(e.message) 138 | -------------------------------------------------------------------------------- /lib/models/resnet.py: -------------------------------------------------------------------------------- 1 | # pylint: disable-all 2 | ''' 3 | Source: https://github.com/akamaster/pytorch_resnet_cifar10 4 | 5 | Properly implemented ResNet-s for CIFAR10 as described in paper [1]. 6 | 7 | The implementation and structure of this file is hugely influenced by [2] 8 | which is implemented for ImageNet and doesn't have option A for identity. 9 | Moreover, most of the implementations on the web is copy-paste from 10 | torchvision's resnet and has wrong number of params. 11 | 12 | Proper ResNet-s for CIFAR10 (for fair comparision and etc.) has following 13 | number of layers and parameters: 14 | 15 | name | layers | params 16 | ResNet20 | 20 | 0.27M 17 | ResNet32 | 32 | 0.46M 18 | ResNet44 | 44 | 0.66M 19 | ResNet56 | 56 | 0.85M 20 | ResNet110 | 110 | 1.7M 21 | ResNet1202| 1202 | 19.4m 22 | 23 | which this implementation indeed has. 24 | 25 | Reference: 26 | [1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun 27 | Deep Residual Learning for Image Recognition. arXiv:1512.03385 28 | [2] https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py 29 | 30 | If you use this implementation in you work, please don't forget to mention the 31 | author, Yerlan Idelbayev. 32 | ''' 33 | import torch.nn as nn 34 | import torch.nn.init as init 35 | import torch.nn.functional as F 36 | 37 | __all__ = ['ResNet', 'resnet20', 'resnet32', 'resnet44', 'resnet56', 'resnet110', 'resnet122', 'resnet1202'] 38 | 39 | 40 | def _weights_init(m): 41 | if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d): 42 | init.kaiming_normal_(m.weight) 43 | 44 | 45 | class LambdaLayer(nn.Module): 46 | def __init__(self, lambd): 47 | super(LambdaLayer, self).__init__() 48 | self.lambd = lambd 49 | 50 | def forward(self, x): 51 | return self.lambd(x) 52 | 53 | 54 | class BasicBlock(nn.Module): 55 | expansion = 1 56 | 57 | def __init__(self, in_planes, planes, stride=1, option='A'): 58 | super(BasicBlock, self).__init__() 59 | self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) 60 | self.bn1 = nn.BatchNorm2d(planes) 61 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) 62 | self.bn2 = nn.BatchNorm2d(planes) 63 | 64 | self.shortcut = nn.Sequential() 65 | if stride != 1 or in_planes != planes: 66 | if option == 'A': 67 | """ 68 | For CIFAR10 ResNet paper uses option A. 69 | """ 70 | self.shortcut = LambdaLayer(lambda x: F.pad(x[:, :, ::2, ::2], 71 | (0, 0, 0, 0, planes // 4, planes // 4), "constant", 0)) 72 | elif option == 'B': 73 | self.shortcut = nn.Sequential( 74 | nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False), 75 | nn.BatchNorm2d(self.expansion * planes)) 76 | 77 | def forward(self, x): 78 | out = F.relu(self.bn1(self.conv1(x))) 79 | out = self.bn2(self.conv2(out)) 80 | out += self.shortcut(x) 81 | out = F.relu(out) 82 | return out 83 | 84 | 85 | class ResNet(nn.Module): 86 | def __init__(self, block, num_blocks): 87 | super(ResNet, self).__init__() 88 | self.in_planes = 16 89 | 90 | self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False) 91 | self.bn1 = nn.BatchNorm2d(16) 92 | self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1) 93 | self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2) 94 | self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2) 95 | # self.linear = nn.Linear(64, num_classes) 96 | 97 | self.apply(_weights_init) 98 | 99 | def _make_layer(self, block, planes, num_blocks, stride): 100 | strides = [stride] + [1] * (num_blocks - 1) 101 | layers = [] 102 | for stride in strides: 103 | layers.append(block(self.in_planes, planes, stride)) 104 | self.in_planes = planes * block.expansion 105 | 106 | return nn.Sequential(*layers) 107 | 108 | def forward(self, x): 109 | out = F.relu(self.bn1(self.conv1(x))) 110 | out = self.layer1(out) 111 | out = self.layer2(out) 112 | out = self.layer3(out) 113 | # out = F.avg_pool2d(out, out.size()[3]) 114 | # out = out.view(out.size(0), -1) 115 | # out = self.linear(out) 116 | return out 117 | 118 | 119 | def resnet20(): 120 | return ResNet(BasicBlock, [3, 3, 3]) 121 | 122 | 123 | def resnet32(): 124 | return ResNet(BasicBlock, [5, 5, 5]) 125 | 126 | 127 | def resnet44(): 128 | return ResNet(BasicBlock, [7, 7, 7]) 129 | 130 | 131 | def resnet50(): 132 | return ResNet(BasicBlock, [8, 8, 8]) 133 | 134 | 135 | def resnet56(): 136 | return ResNet(BasicBlock, [9, 9, 9]) 137 | 138 | 139 | def resnet110(): 140 | return ResNet(BasicBlock, [18, 18, 18]) 141 | 142 | 143 | def resnet122(): 144 | return ResNet(BasicBlock, [20, 20, 20]) 145 | 146 | 147 | def resnet1202(): 148 | return ResNet(BasicBlock, [200, 200, 200]) 149 | 150 | 151 | def test(net): 152 | import numpy as np 153 | total_params = 0 154 | 155 | for x in filter(lambda p: p.requires_grad, net.parameters()): 156 | total_params += np.prod(x.data.numpy().shape) 157 | print("Total number of params", total_params) 158 | print("Total layers", len(list(filter(lambda p: p.requires_grad and len(p.data.size()) > 1, net.parameters())))) 159 | 160 | 161 | if __name__ == "__main__": 162 | for net_name in __all__: 163 | if net_name.startswith('resnet'): 164 | print(net_name) 165 | test(globals()[net_name]()) 166 | print() 167 | -------------------------------------------------------------------------------- /lib/datasets/llamas.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle as pkl 3 | 4 | import numpy as np 5 | from tqdm import tqdm 6 | 7 | from .lane_dataset_loader import LaneDatasetLoader 8 | 9 | TRAIN_LABELS_DIR = 'labels/train' 10 | TEST_LABELS_DIR = 'labels/valid' 11 | TEST_IMGS_DIR = 'color_images/test' 12 | SPLIT_DIRECTORIES = {'train': 'labels/train', 'val': 'labels/valid'} 13 | from utils.llamas_utils import get_horizontal_values_for_four_lanes 14 | import utils.llamas_metric as llamas_metric 15 | 16 | 17 | class LLAMAS(LaneDatasetLoader): 18 | def __init__(self, split='train', max_lanes=None, root=None): 19 | self.split = split 20 | self.root = root 21 | if split != 'test' and split not in SPLIT_DIRECTORIES.keys(): 22 | raise Exception('Split `{}` does not exist.'.format(split)) 23 | if split != 'test': 24 | self.labels_dir = os.path.join(self.root, SPLIT_DIRECTORIES[split]) 25 | 26 | self.img_w, self.img_h = 1276, 717 27 | self.annotations = [] 28 | self.load_annotations() 29 | 30 | # Force max_lanes, used when evaluating testing with models trained on other datasets 31 | if max_lanes is not None: 32 | self.max_lanes = max_lanes 33 | 34 | def get_img_heigth(self, _): 35 | return self.img_h 36 | 37 | def get_img_width(self, _): 38 | return self.img_w 39 | 40 | def get_metrics(self, lanes, _): 41 | # Placeholders 42 | return [0] * len(lanes), [0] * len(lanes), [1] * len(lanes), [1] * len(lanes) 43 | 44 | def get_img_path(self, json_path): 45 | # /foo/bar/test/folder/image_label.ext --> test/folder/image_label.ext 46 | base_name = '/'.join(json_path.split('/')[-3:]) 47 | image_path = os.path.join('color_images', base_name.replace('.json', '_color_rect.png')) 48 | return image_path 49 | 50 | def get_json_paths(self): 51 | json_paths = [] 52 | for root, _, files in os.walk(self.labels_dir): 53 | for file in files: 54 | if file.endswith(".json"): 55 | json_paths.append(os.path.join(root, file)) 56 | return json_paths 57 | 58 | def load_annotations(self): 59 | # the labels are not public for the test set yet 60 | if self.split == 'test': 61 | imgs_dir = os.path.join(self.root, TEST_IMGS_DIR) 62 | self.annotations = [{ 63 | 'path': os.path.join(root, file), 64 | 'lanes': [], 65 | 'relative_path': file 66 | } for root, _, files in os.walk(imgs_dir) for file in files if file.endswith('.png')] 67 | self.annotations = sorted(self.annotations, key=lambda x: x['path']) 68 | return 69 | # Waiting for the dataset to load is tedious, let's cache it 70 | os.makedirs('cache', exist_ok=True) 71 | cache_path = 'cache/llamas_{}.pkl'.format(self.split) 72 | if os.path.exists(cache_path): 73 | with open(cache_path, 'rb') as cache_file: 74 | self.annotations = pkl.load(cache_file) 75 | self.max_lanes = max(len(anno['lanes']) for anno in self.annotations) 76 | return 77 | 78 | self.max_lanes = 0 79 | print("Searching annotation files...") 80 | json_paths = self.get_json_paths() 81 | print('{} annotations found.'.format(len(json_paths))) 82 | 83 | for json_path in tqdm(json_paths): 84 | lanes = get_horizontal_values_for_four_lanes(json_path) 85 | lanes = [[(x, y) for x, y in zip(lane, range(self.img_h)) if x >= 0] for lane in lanes] 86 | lanes = [lane for lane in lanes if len(lane) > 0] 87 | relative_path = self.get_img_path(json_path) 88 | img_path = os.path.join(self.root, relative_path) 89 | self.max_lanes = max(self.max_lanes, len(lanes)) 90 | self.annotations.append({'path': img_path, 'lanes': lanes, 'aug': False, 'relative_path': relative_path}) 91 | 92 | with open(cache_path, 'wb') as cache_file: 93 | pkl.dump(self.annotations, cache_file) 94 | 95 | def assign_class_to_lanes(self, lanes): 96 | return {label: value for label, value in zip(['l0', 'l1', 'r0', 'r1'], lanes)} 97 | 98 | def get_prediction_string(self, pred): 99 | ys = np.arange(self.img_h) / self.img_h 100 | out = [] 101 | for lane in pred: 102 | xs = lane(ys) 103 | valid_mask = (xs >= 0) & (xs < 1) 104 | xs = xs * self.img_w 105 | lane_xs = xs[valid_mask] 106 | lane_ys = ys[valid_mask] * self.img_h 107 | lane_xs, lane_ys = lane_xs[::-1], lane_ys[::-1] 108 | lane_str = ' '.join(['{:.5f} {:.5f}'.format(x, y) for x, y in zip(lane_xs, lane_ys)]) 109 | if lane_str != '': 110 | out.append(lane_str) 111 | 112 | return '\n'.join(out) 113 | 114 | def eval_predictions(self, predictions, output_basedir): 115 | print('Generating prediction output...') 116 | for idx, pred in enumerate(tqdm(predictions)): 117 | relative_path = self.annotations[idx]['old_anno']['relative_path'] 118 | output_filename = '/'.join(relative_path.split('/')[-2:]).replace('_color_rect.png', '.lines.txt') 119 | output_filepath = os.path.join(output_basedir, output_filename) 120 | os.makedirs(os.path.dirname(output_filepath), exist_ok=True) 121 | output = self.get_prediction_string(pred) 122 | with open(output_filepath, 'w') as out_file: 123 | out_file.write(output) 124 | if self.split == 'test': 125 | return {} 126 | return llamas_metric.eval_predictions(output_basedir, self.labels_dir, unofficial=False) 127 | 128 | def __getitem__(self, idx): 129 | return self.annotations[idx] 130 | 131 | def __len__(self): 132 | return len(self.annotations) 133 | -------------------------------------------------------------------------------- /lib/datasets/culane.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import logging 4 | 5 | import numpy as np 6 | from tqdm import tqdm 7 | 8 | import utils.culane_metric as culane_metric 9 | 10 | from .lane_dataset_loader import LaneDatasetLoader 11 | 12 | SPLIT_FILES = { 13 | 'train': "list/train.txt", 14 | 'val': 'list/val.txt', 15 | 'test': "list/test.txt", 16 | 'normal': 'list/test_split/test0_normal.txt', 17 | 'crowd': 'list/test_split/test1_crowd.txt', 18 | 'hlight': 'list/test_split/test2_hlight.txt', 19 | 'shadow': 'list/test_split/test3_shadow.txt', 20 | 'noline': 'list/test_split/test4_noline.txt', 21 | 'arrow': 'list/test_split/test5_arrow.txt', 22 | 'curve': 'list/test_split/test6_curve.txt', 23 | 'cross': 'list/test_split/test7_cross.txt', 24 | 'night': 'list/test_split/test8_night.txt', 25 | 'debug': 'list/debug.txt' 26 | } 27 | 28 | 29 | class CULane(LaneDatasetLoader): 30 | def __init__(self, max_lanes=None, split='train', root=None, official_metric=True): 31 | self.split = split 32 | self.root = root 33 | self.official_metric = official_metric 34 | self.logger = logging.getLogger(__name__) 35 | 36 | if root is None: 37 | raise Exception('Please specify the root directory') 38 | if split not in SPLIT_FILES: 39 | raise Exception('Split `{}` does not exist.'.format(split)) 40 | 41 | self.list = os.path.join(root, SPLIT_FILES[split]) 42 | 43 | self.img_w, self.img_h = 1640, 590 44 | self.annotations = [] 45 | self.load_annotations() 46 | self.max_lanes = 4 if max_lanes is None else max_lanes 47 | 48 | def get_img_heigth(self): 49 | return self.img_h 50 | 51 | def get_img_width(self): 52 | return self.img_w 53 | 54 | def get_metrics(self, raw_lanes, idx): 55 | lanes = [] 56 | pred_str = self.get_prediction_string(raw_lanes) 57 | for lane in pred_str.split('\n'): 58 | if lane == '': 59 | continue 60 | lane = list(map(float, lane.split())) 61 | lane = [(lane[i], lane[i + 1]) for i in range(0, len(lane), 2) if lane[i] >= 0 and lane[i + 1] >= 0] 62 | lanes.append(lane) 63 | anno = culane_metric.load_culane_img_data(self.annotations[idx]['path'].replace('.jpg', '.lines.txt')) 64 | _, fp, fn, ious, matches = culane_metric.culane_metric(lanes, anno) 65 | 66 | return fp, fn, matches, ious 67 | 68 | def load_annotation(self, img_path): 69 | anno_path = img_path[:-3] + 'lines.txt' # remove sufix jpg and add lines.txt 70 | 71 | with open(anno_path, 'r') as anno_file: 72 | data = [list(map(float, line.split())) for line in anno_file.readlines()] 73 | 74 | lanes = [[(lane[i], lane[i + 1]) for i in range(0, len(lane), 2) if lane[i] >= 0 and lane[i + 1] >= 0] 75 | for lane in data] 76 | lanes = [list(set(lane)) for lane in lanes] # remove duplicated points 77 | lanes = [lane for lane in lanes if len(lane) >= 2] # remove lanes with less than 2 points 78 | 79 | lanes = [sorted(lane, key=lambda x: x[1]) for lane in lanes] # sort by y 80 | 81 | return {'path': img_path, 'lanes': lanes} 82 | 83 | def load_annotations(self): 84 | self.annotations = [] 85 | self.max_lanes = 0 86 | os.makedirs('cache', exist_ok=True) 87 | cache_path = 'cache/culane_{}.json'.format(self.split) 88 | 89 | if os.path.exists(cache_path): 90 | self.logger.info('Loading CULane annotations (cached)...') 91 | with open(cache_path, 'r') as cache_file: 92 | data = json.load(cache_file) 93 | self.annotations = data['annotations'] 94 | self.max_lanes = data['max_lanes'] 95 | else: 96 | self.logger.info('Loading CULane annotations and caching...') 97 | with open(self.list, 'r') as list_file: 98 | files = [line.rstrip()[1 if line[0] == '/' else 0::] 99 | for line in list_file] # remove `/` from beginning if needed 100 | 101 | for file in tqdm(files): 102 | img_path = os.path.join(self.root, file) 103 | anno = self.load_annotation(img_path) 104 | anno['org_path'] = file 105 | 106 | if len(anno['lanes']) > 0: 107 | self.max_lanes = max(self.max_lanes, len(anno['lanes'])) 108 | self.annotations.append(anno) 109 | with open(cache_path, 'w') as cache_file: 110 | json.dump({'annotations': self.annotations, 'max_lanes': self.max_lanes}, cache_file) 111 | 112 | self.logger.info('%d annotations loaded, with a maximum of %d lanes in an image.', len(self.annotations), 113 | self.max_lanes) 114 | 115 | def get_prediction_string(self, pred): 116 | ys = np.arange(self.img_h) / self.img_h 117 | out = [] 118 | for lane in pred: 119 | xs = lane(ys) 120 | valid_mask = (xs >= 0) & (xs < 1) 121 | xs = xs * self.img_w 122 | lane_xs = xs[valid_mask] 123 | lane_ys = ys[valid_mask] * self.img_h 124 | lane_xs, lane_ys = lane_xs[::-1], lane_ys[::-1] 125 | lane_str = ' '.join(['{:.5f} {:.5f}'.format(x, y) for x, y in zip(lane_xs, lane_ys)]) 126 | if lane_str != '': 127 | out.append(lane_str) 128 | 129 | return '\n'.join(out) 130 | 131 | def eval_predictions(self, predictions, output_basedir): 132 | print('Generating prediction output...') 133 | for idx, pred in enumerate(tqdm(predictions)): 134 | output_dir = os.path.join(output_basedir, os.path.dirname(self.annotations[idx]['old_anno']['org_path'])) 135 | output_filename = os.path.basename(self.annotations[idx]['old_anno']['org_path'])[:-3] + 'lines.txt' 136 | os.makedirs(output_dir, exist_ok=True) 137 | output = self.get_prediction_string(pred) 138 | with open(os.path.join(output_dir, output_filename), 'w') as out_file: 139 | out_file.write(output) 140 | return culane_metric.eval_predictions(output_basedir, self.root, self.list, official=self.official_metric) 141 | 142 | def transform_annotations(self, transform): 143 | self.annotations = list(map(transform, self.annotations)) 144 | 145 | def __getitem__(self, idx): 146 | return self.annotations[idx] 147 | 148 | def __len__(self): 149 | return len(self.annotations) 150 | -------------------------------------------------------------------------------- /lib/runner.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import random 3 | import logging 4 | 5 | import cv2 6 | import torch 7 | import numpy as np 8 | from tqdm import tqdm, trange 9 | import pdb 10 | 11 | 12 | class Runner: 13 | def __init__(self, cfg, exp, device, resume=False, view=None, deterministic=False): 14 | self.cfg = cfg 15 | self.exp = exp 16 | self.device = device 17 | self.resume = resume 18 | self.view = view 19 | self.logger = logging.getLogger(__name__) 20 | 21 | # Fix seeds 22 | torch.manual_seed(cfg['seed']) 23 | np.random.seed(cfg['seed']) 24 | random.seed(cfg['seed']) 25 | 26 | if deterministic: 27 | torch.backends.cudnn.deterministic = True 28 | torch.backends.cudnn.benchmark = False 29 | 30 | def train(self): 31 | self.exp.train_start_callback(self.cfg) 32 | starting_epoch = 1 33 | model = self.cfg.get_model() 34 | model = model.to(self.device) 35 | optimizer = self.cfg.get_optimizer(model.parameters()) 36 | scheduler = self.cfg.get_lr_scheduler(optimizer) 37 | if self.resume: 38 | last_epoch, model, optimizer, scheduler = self.exp.load_last_train_state(model, optimizer, scheduler) 39 | starting_epoch = last_epoch + 1 40 | max_epochs = self.cfg['epochs'] 41 | train_loader = self.get_train_dataloader() 42 | loss_parameters = self.cfg.get_loss_parameters() 43 | for epoch in trange(starting_epoch, max_epochs + 1, initial=starting_epoch - 1, total=max_epochs): 44 | self.exp.epoch_start_callback(epoch, max_epochs) 45 | model.train() 46 | pbar = tqdm(train_loader) 47 | for i, (images, labels, _) in enumerate(pbar): 48 | images = images.to(self.device) 49 | labels = labels.to(self.device) 50 | 51 | # Forward pass 52 | outputs = model(images, **self.cfg.get_train_parameters()) 53 | 54 | ######multy lane 55 | loss, loss_dict_i = model.loss(outputs, labels, **loss_parameters) 56 | 57 | # loss, loss_dict_i = model.loss(outputs, labels, **loss_parameters) 58 | 59 | # Backward and optimize 60 | optimizer.zero_grad() 61 | loss.backward() 62 | optimizer.step() 63 | 64 | # Scheduler step (iteration based) 65 | scheduler.step() 66 | 67 | # Log 68 | postfix_dict = {key: float(value) for key, value in loss_dict_i.items()} 69 | postfix_dict['lr'] = optimizer.param_groups[0]["lr"] 70 | self.exp.iter_end_callback(epoch, max_epochs, i, len(train_loader), loss.item(), postfix_dict) 71 | postfix_dict['loss'] = loss.item() 72 | pbar.set_postfix(ordered_dict=postfix_dict) 73 | self.exp.epoch_end_callback(epoch, max_epochs, model, optimizer, scheduler) 74 | 75 | # Validate 76 | #if (epoch + 1) % self.cfg['val_every'] == 0: 77 | # self.eval(epoch, on_val=True) 78 | self.exp.train_end_callback() 79 | 80 | def eval(self, epoch, on_val=False, save_predictions=False): 81 | model = self.cfg.get_model() 82 | model_path = self.exp.get_checkpoint_path(epoch) 83 | self.logger.info('Loading model %s', model_path) 84 | model.load_state_dict(self.exp.get_epoch_model(epoch)) 85 | model = model.to(self.device) 86 | model.eval() 87 | if on_val: 88 | dataloader = self.get_val_dataloader() 89 | else: 90 | dataloader = self.get_test_dataloader() 91 | test_parameters = self.cfg.get_test_parameters() 92 | predictions = [] 93 | self.exp.eval_start_callback(self.cfg) 94 | with torch.no_grad(): 95 | for idx, (images, _, _) in enumerate(tqdm(dataloader)): 96 | images = images.to(self.device) 97 | output = model(images, **test_parameters) 98 | prediction = model.decode(output, as_lanes=True) 99 | predictions.extend(prediction) 100 | 101 | for bb in range(len(images)): 102 | img = (images[bb].cpu().permute(1, 2, 0).numpy() * 255).astype(np.uint8) 103 | img, fp, fn = dataloader.dataset.draw_annotation(idx, img=img, pred=prediction[bb]) 104 | cv2.imwrite('/home/LaneATT-main/XXX/'+'pred_'+ str(idx)+ "_"+str(bb)+".jpg", img) 105 | 106 | 107 | #if self.view: 108 | # img = (images[0].cpu().permute(1, 2, 0).numpy() * 255).astype(np.uint8) 109 | # img, fp, fn = dataloader.dataset.draw_annotation(idx, img=img, pred=prediction[0]) 110 | # if self.view == 'mistakes' and fp == 0 and fn == 0: 111 | # continue 112 | #cv2.imshow('pred', img) 113 | #cv2.waitKey(0) 114 | 115 | if save_predictions: 116 | with open('predictions.pkl', 'wb') as handle: 117 | pickle.dump(predictions, handle, protocol=pickle.HIGHEST_PROTOCOL) 118 | self.exp.eval_end_callback(dataloader.dataset.dataset, predictions, epoch) 119 | 120 | def get_train_dataloader(self): 121 | train_dataset = self.cfg.get_dataset('train') 122 | train_loader = torch.utils.data.DataLoader(dataset=train_dataset, 123 | batch_size=self.cfg['batch_size'], 124 | shuffle=True, 125 | num_workers=0, 126 | worker_init_fn=self._worker_init_fn_) 127 | return train_loader 128 | 129 | def get_test_dataloader(self): 130 | test_dataset = self.cfg.get_dataset('test') 131 | test_loader = torch.utils.data.DataLoader(dataset=test_dataset, 132 | batch_size=self.cfg['batch_size'] if not self.view else 1, 133 | shuffle=False, 134 | num_workers=8, 135 | worker_init_fn=self._worker_init_fn_) 136 | return test_loader 137 | 138 | def get_val_dataloader(self): 139 | val_dataset = self.cfg.get_dataset('val') 140 | val_loader = torch.utils.data.DataLoader(dataset=val_dataset, 141 | batch_size=self.cfg['batch_size'], 142 | shuffle=False, 143 | num_workers=8, 144 | worker_init_fn=self._worker_init_fn_) 145 | return val_loader 146 | 147 | @staticmethod 148 | def _worker_init_fn_(_): 149 | torch_seed = torch.initial_seed() 150 | np_seed = torch_seed // 2**32 - 1 151 | random.seed(torch_seed) 152 | np.random.seed(np_seed) 153 | -------------------------------------------------------------------------------- /lib/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 | 11 | def one_hot(labels: torch.Tensor, 12 | num_classes: int, 13 | device: Optional[torch.device] = None, 14 | dtype: Optional[torch.dtype] = None, 15 | eps: Optional[float] = 1e-6) -> torch.Tensor: 16 | r"""Converts an integer label x-D tensor to a one-hot (x+1)-D tensor. 17 | 18 | Args: 19 | labels (torch.Tensor) : tensor with labels of shape :math:`(N, *)`, 20 | where N is batch size. Each value is an integer 21 | representing correct classification. 22 | num_classes (int): number of classes in labels. 23 | device (Optional[torch.device]): the desired device of returned tensor. 24 | Default: if None, uses the current device for the default tensor type 25 | (see torch.set_default_tensor_type()). device will be the CPU for CPU 26 | tensor types and the current CUDA device for CUDA tensor types. 27 | dtype (Optional[torch.dtype]): the desired data type of returned 28 | tensor. Default: if None, infers data type from values. 29 | 30 | Returns: 31 | torch.Tensor: the labels in one hot tensor of shape :math:`(N, C, *)`, 32 | 33 | Examples:: 34 | # >>> labels = torch.LongTensor([[[0, 1], [2, 0]]]) 35 | # >>> kornia.losses.one_hot(labels, num_classes=3) 36 | tensor([[[[1., 0.], 37 | [0., 1.]], 38 | [[0., 1.], 39 | [0., 0.]], 40 | [[0., 0.], 41 | [1., 0.]]]] 42 | """ 43 | if not torch.is_tensor(labels): 44 | raise TypeError("Input labels type is not a torch.Tensor. Got {}".format(type(labels))) 45 | if not labels.dtype == torch.int64: 46 | raise ValueError("labels must be of the same dtype torch.int64. Got: {}".format(labels.dtype)) 47 | if num_classes < 1: 48 | raise ValueError("The number of classes must be bigger than one." " Got: {}".format(num_classes)) 49 | shape = labels.shape 50 | one_hot = torch.zeros(shape[0], num_classes, *shape[1:], device=device, dtype=dtype) 51 | return one_hot.scatter_(1, labels.unsqueeze(1), 1.0) + eps 52 | 53 | 54 | def focal_loss(input: torch.Tensor, 55 | target: torch.Tensor, 56 | alpha: float, 57 | gamma: float = 2.0, 58 | reduction: str = 'none', 59 | eps: float = 1e-8) -> torch.Tensor: 60 | r"""Function that computes Focal loss. 61 | 62 | See :class:`~kornia.losses.FocalLoss` for details. 63 | """ 64 | if not torch.is_tensor(input): 65 | raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(input))) 66 | 67 | if not len(input.shape) >= 2: 68 | raise ValueError("Invalid input shape, we expect BxCx*. Got: {}".format(input.shape)) 69 | 70 | if input.size(0) != target.size(0): 71 | raise ValueError('Expected input batch_size ({}) to match target batch_size ({}).'.format( 72 | input.size(0), target.size(0))) 73 | 74 | n = input.size(0) 75 | out_size = (n, ) + input.size()[2:] 76 | if target.size()[1:] != input.size()[2:]: 77 | raise ValueError('Expected target size {}, got {}'.format(out_size, target.size())) 78 | 79 | if not input.device == target.device: 80 | raise ValueError("input and target must be in the same device. Got: {} and {}".format( 81 | input.device, target.device)) 82 | 83 | ####################### key code ######################### 84 | # compute softmax over the classes axis 85 | input_soft: torch.Tensor = F.softmax(input, dim=1) + eps 86 | 87 | # create the labels one hot tensor 88 | target_one_hot: torch.Tensor = one_hot(target, num_classes=input.shape[1], device=input.device, dtype=input.dtype) 89 | 90 | # compute the actual focal loss 91 | weight = torch.pow(-input_soft + 1., gamma) 92 | 93 | focal = -alpha * weight * torch.log(input_soft) 94 | loss_tmp = torch.sum(target_one_hot * focal, dim=1) 95 | 96 | if reduction == 'none': 97 | loss = loss_tmp 98 | elif reduction == 'mean': 99 | loss = torch.mean(loss_tmp) 100 | elif reduction == 'sum': 101 | loss = torch.sum(loss_tmp) 102 | else: 103 | raise NotImplementedError("Invalid reduction mode: {}".format(reduction)) 104 | return loss 105 | 106 | 107 | class FocalLoss(nn.Module): 108 | r"""Criterion that computes Focal loss. 109 | 110 | According to [1], the Focal loss is computed as follows: 111 | 112 | .. math:: 113 | 114 | \text{FL}(p_t) = -\alpha_t (1 - p_t)^{\gamma} \, \text{log}(p_t) 115 | 116 | where: 117 | - :math:`p_t` is the model's estimated probability for each class. 118 | 119 | 120 | Arguments: 121 | alpha (float): Weighting factor :math:`\alpha \in [0, 1]`. 122 | gamma (float): Focusing parameter :math:`\gamma >= 0`. 123 | reduction (str, optional): Specifies the reduction to apply to the 124 | output: ‘none’ | ‘mean’ | ‘sum’. ‘none’: no reduction will be applied, 125 | ‘mean’: the sum of the output will be divided by the number of elements 126 | in the output, ‘sum’: the output will be summed. Default: ‘none’. 127 | 128 | Shape: 129 | - Input: :math:`(N, C, *)` where C = number of classes. 130 | - Target: :math:`(N, *)` where each value is 131 | :math:`0 ≤ targets[i] ≤ C−1`. 132 | 133 | Examples: 134 | # >>> N = 5 # num_classes 135 | # >>> kwargs = {"alpha": 0.5, "gamma": 2.0, "reduction": 'mean'} 136 | # >>> loss = kornia.losses.FocalLoss(**kwargs) 137 | # >>> input = torch.randn(1, N, 3, 5, requires_grad=True) 138 | # >>> target = torch.empty(1, 3, 5, dtype=torch.long).random_(N) 139 | # >>> output = loss(input, target) 140 | # >>> output.backward() 141 | 142 | References: 143 | [1] https://arxiv.org/abs/1708.02002 144 | """ 145 | def __init__(self, alpha: float, gamma: float = 2.0, reduction: str = 'none') -> None: 146 | super(FocalLoss, self).__init__() 147 | self.alpha: float = alpha 148 | self.gamma: float = gamma 149 | self.reduction: str = reduction 150 | self.eps: float = 1e-6 151 | 152 | def forward( # type: ignore 153 | self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: 154 | return focal_loss(input, target, self.alpha, self.gamma, self.reduction, self.eps) 155 | 156 | 157 | if __name__ == '__main__': 158 | import numpy as np 159 | 160 | focal_losshh = FocalLoss(alpha=0.25, gamma=2.) 161 | aa = np.array([0,1,1], dtype=np.int64) 162 | cls_target = torch.from_numpy(aa) 163 | bb = np.array([[0.01, 0.46, 0.09],[0.01, 0.46, 0.09],[0.01, 0.46, 0.09]],dtype=np.float) 164 | cls_pred = torch.from_numpy(bb) 165 | cls_loss = focal_losshh(cls_pred, cls_target).sum() 166 | print(cls_loss) 167 | 168 | 169 | 170 | 171 | 172 | 173 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /lib/nms/src/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 (16 + 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[13] * N_STRIPS - DATASET_OFFSET + 0.5); // 0.5 rounding trick 30 | const int start_b = (int) (b[13] * N_STRIPS - DATASET_OFFSET + 0.5); 31 | const int start = max(start_a, start_b); 32 | const int end_a = start_a + a[15] - 1 + 0.5 - ((a[15] - 1) < 0); // - (x<0) trick to adjust for negative numbers (in case length is 0) 33 | const int end_b = start_b + b[15] - 1 + 0.5 - ((b[15] - 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 = 16 + start; i <= 16 + 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 | -------------------------------------------------------------------------------- /lib/experiment.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | import json 4 | import logging 5 | import subprocess 6 | 7 | import torch 8 | from torch.utils.tensorboard import SummaryWriter 9 | 10 | 11 | class Experiment: 12 | def __init__(self, exp_name, args=None, mode='train', exps_basedir='experiments', tensorboard_dir='tensorboard'): 13 | self.name = exp_name 14 | self.exp_dirpath = os.path.join(exps_basedir, exp_name) 15 | self.models_dirpath = os.path.join(self.exp_dirpath, 'models') 16 | self.results_dirpath = os.path.join(self.exp_dirpath, 'results') 17 | self.cfg_path = os.path.join(self.exp_dirpath, 'config.yaml') 18 | self.code_state_path = os.path.join(self.exp_dirpath, 'code_state.txt') 19 | self.log_path = os.path.join(self.exp_dirpath, 'log_{}.txt'.format(mode)) 20 | self.tensorboard_writer = SummaryWriter(os.path.join(tensorboard_dir, exp_name)) 21 | self.cfg = None 22 | self.setup_exp_dir() 23 | self.setup_logging() 24 | 25 | if args is not None: 26 | self.log_args(args) 27 | 28 | def setup_exp_dir(self): 29 | if not os.path.exists(self.exp_dirpath): 30 | os.makedirs(self.exp_dirpath) 31 | os.makedirs(self.models_dirpath) 32 | os.makedirs(self.results_dirpath) 33 | self.save_code_state() 34 | 35 | def save_code_state(self): 36 | state = "Git hash: {}".format( 37 | subprocess.run(['git', 'rev-parse', 'HEAD'], stdout=subprocess.PIPE, check=False).stdout.decode('utf-8')) 38 | state += '\n*************\nGit diff:\n*************\n' 39 | state += subprocess.run(['git', 'diff'], stdout=subprocess.PIPE, check=False).stdout.decode('utf-8') 40 | with open(self.code_state_path, 'w') as code_state_file: 41 | code_state_file.write(state) 42 | 43 | def setup_logging(self): 44 | formatter = logging.Formatter("[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s") 45 | file_handler = logging.FileHandler(self.log_path) 46 | file_handler.setLevel(logging.DEBUG) 47 | file_handler.setFormatter(formatter) 48 | stream_handler = logging.StreamHandler() 49 | stream_handler.setLevel(logging.INFO) 50 | stream_handler.setFormatter(formatter) 51 | logging.basicConfig(level=logging.DEBUG, handlers=[file_handler, stream_handler]) 52 | self.logger = logging.getLogger(__name__) 53 | 54 | def log_args(self, args): 55 | self.logger.debug('CLI Args:\n %s', str(args)) 56 | 57 | def set_cfg(self, cfg, override=False): 58 | assert 'model_checkpoint_interval' in cfg 59 | self.cfg = cfg 60 | if not os.path.exists(self.cfg_path) or override: 61 | with open(self.cfg_path, 'w') as cfg_file: 62 | cfg_file.write(str(cfg)) 63 | 64 | def get_last_checkpoint_epoch(self): 65 | pattern = re.compile('model_(\\d+).pt') 66 | last_epoch = -1 67 | for ckpt_file in os.listdir(self.models_dirpath): 68 | result = pattern.match(ckpt_file) 69 | if result is not None: 70 | epoch = int(result.groups()[0]) 71 | if epoch > last_epoch: 72 | last_epoch = epoch 73 | 74 | return last_epoch 75 | 76 | def get_checkpoint_path(self, epoch): 77 | return os.path.join(self.models_dirpath, 'model_{:04d}.pt'.format(epoch)) 78 | 79 | def get_epoch_model(self, epoch): 80 | return torch.load(self.get_checkpoint_path(epoch))['model'] 81 | 82 | def load_last_train_state(self, model, optimizer, scheduler, device=None): 83 | epoch = self.get_last_checkpoint_epoch() 84 | train_state_path = self.get_checkpoint_path(epoch) 85 | train_state = torch.load(train_state_path, map_location=device) 86 | #train_state = torch.load(train_state_path) 87 | 88 | # for k, v in train_state['model'].items(): 89 | # print(k) 90 | ###### pop cls_layer for multly lane type 91 | model_dict = model.state_dict() 92 | pop_cls_layer = {k: v for k, v in train_state['model'].items() if 'cls_layer' not in k} 93 | model_dict.update(pop_cls_layer) 94 | model.load_state_dict(model_dict) 95 | 96 | # model.load_state_dict(train_state['model']) 97 | # optimizer.load_state_dict(train_state['optimizer']) 98 | # scheduler.load_state_dict(train_state['scheduler']) 99 | 100 | return epoch, model, optimizer, scheduler 101 | 102 | def save_train_state(self, epoch, model, optimizer, scheduler): 103 | train_state_path = self.get_checkpoint_path(epoch) 104 | torch.save( 105 | { 106 | 'epoch': epoch, 107 | 'model': model.state_dict(), 108 | 'optimizer': optimizer.state_dict(), 109 | 'scheduler': scheduler.state_dict() 110 | }, train_state_path) 111 | 112 | def iter_end_callback(self, epoch, max_epochs, iter_nb, max_iter, loss, loss_components): 113 | line = 'Epoch [{}/{}] - Iter [{}/{}] - Loss: {:.5f} - '.format(epoch, max_epochs, iter_nb, max_iter, loss) 114 | line += ' - '.join( 115 | ['{}: {:.5f}'.format(component, loss_components[component]) for component in loss_components]) 116 | self.logger.debug(line) 117 | overall_iter = (epoch * max_iter) + iter_nb 118 | self.tensorboard_writer.add_scalar('loss/total_loss', loss, overall_iter) 119 | for key in loss_components: 120 | self.tensorboard_writer.add_scalar('loss/{}'.format(key), loss_components[key], overall_iter) 121 | 122 | def epoch_start_callback(self, epoch, max_epochs): 123 | self.logger.debug('Epoch [%d/%d] starting.', epoch, max_epochs) 124 | 125 | def epoch_end_callback(self, epoch, max_epochs, model, optimizer, scheduler): 126 | self.logger.debug('Epoch [%d/%d] finished.', epoch, max_epochs) 127 | if epoch % self.cfg['model_checkpoint_interval'] == 0: 128 | self.save_train_state(epoch, model, optimizer, scheduler) 129 | 130 | def train_start_callback(self, cfg): 131 | self.logger.debug('Beginning training session. CFG used:\n%s', str(cfg)) 132 | 133 | def train_end_callback(self): 134 | self.logger.debug('Training session finished.') 135 | 136 | def eval_start_callback(self, cfg): 137 | self.logger.debug('Beginning testing session. CFG used:\n%s', str(cfg)) 138 | 139 | def eval_end_callback(self, dataset, predictions, epoch_evaluated): 140 | metrics = self.save_epoch_results(dataset, predictions, epoch_evaluated) 141 | self.logger.debug('Testing session finished on model after epoch %d.', epoch_evaluated) 142 | self.logger.info('Results:\n %s', str(metrics)) 143 | 144 | def save_epoch_results(self, dataset, predictions, epoch): 145 | # setup dirs 146 | epoch_results_path = os.path.join(self.results_dirpath, 'epoch_{:04d}'.format(epoch)) 147 | predictions_dir = os.path.join(epoch_results_path, '{}_predictions'.format(dataset.split)) 148 | os.makedirs(predictions_dir, exist_ok=True) 149 | # eval metrics 150 | metrics = dataset.eval_predictions(predictions, output_basedir=predictions_dir) 151 | # log tensorboard metrics 152 | for key in metrics: 153 | self.tensorboard_writer.add_scalar('{}_metrics/{}'.format(dataset.split, key), metrics[key], epoch) 154 | # save metrics 155 | metrics_path = os.path.join(epoch_results_path, '{}_metrics.json'.format(dataset.split)) 156 | with open(metrics_path, 'w') as results_file: 157 | json.dump(metrics, results_file) 158 | # save the cfg used 159 | with open(os.path.join(epoch_results_path, 'config.yaml'), 'w') as cfg_file: 160 | cfg_file.write(str(self.cfg)) 161 | 162 | return metrics 163 | -------------------------------------------------------------------------------- /lib/infer_torchscript.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | from torchvision.transforms import ToTensor 4 | from lane import Lane 5 | import torch 6 | from nms import nms 7 | import torch.nn as nn 8 | import time 9 | import os 10 | os.environ["CUDA_VISIBLE_DEVICES"] = "6" 11 | 12 | model_path = "/home/LaneATT-main/experiments/laneatt_r34_ehlwx/models/model_0014.torchscript.pt" 13 | 14 | to_tensor = ToTensor() 15 | 16 | GT_COLOR = {"1":(255, 0, 0), "2":(55, 28, 206),"3":(35, 198, 26), "4":(35, 78, 210), 17 | "5":(255, 255, 0), "6":(55, 255, 206),"7":(255, 18, 26), "8":(255, 78, 210), 18 | "9":(255, 0, 255), "10":(55, 28, 255),"11":(35, 198, 255), "12":(35, 78, 255)} 19 | PRED_HIT_COLOR = (0, 255, 0) 20 | PRED_MISS_COLOR = (0, 0, 255) 21 | IMAGENET_MEAN = np.array([0.485, 0.456, 0.406]) 22 | IMAGENET_STD = np.array([0.229, 0.224, 0.225]) 23 | img_w, img_h = 640, 360 24 | num_lane_types = 13 # bg + 12 25 | n_strips = 71 26 | anchor_ys = torch.linspace(1, 0, steps=72, dtype=torch.float32) 27 | 28 | 29 | def draw_annotation(pred=None, img=None): 30 | pad = 0 31 | for i, l in enumerate(pred): 32 | if l.metadata['type'] == 0: 33 | continue 34 | color = PRED_MISS_COLOR 35 | points = l.points 36 | points[:, 0] *= img.shape[1] 37 | points[:, 1] *= img.shape[0] 38 | points = points.round().astype(int) 39 | points += pad 40 | xs, ys = points[:, 0], points[:, 1] 41 | 42 | if int(l.metadata['type']) in [1,3,5,7,9]: ##虚线 43 | for curr_p in points: 44 | img = cv2.circle(img, tuple(curr_p), 5, GT_COLOR[str(int(l.metadata['type']))], -1) 45 | else: 46 | for curr_p, next_p in zip(points[:-1], points[1:]): 47 | img = cv2.line(img, 48 | tuple(curr_p), 49 | tuple(next_p), 50 | color=GT_COLOR[str(int(l.metadata['type']))], 51 | thickness=3) 52 | """ 53 | if 'start_x' in l.metadata: 54 | start_x = l.metadata['start_x'] * img.shape[1] 55 | start_y = l.metadata['start_y'] * img.shape[0] 56 | cv2.circle(img, (int(start_x + pad), int(img_h - 1 - start_y + pad)), 57 | radius=5, 58 | color=(0, 0, 255), 59 | thickness=-1) 60 | """ 61 | if len(xs) == 0: 62 | print("Empty pred") 63 | if len(xs) > 0: 64 | cv2.putText(img, 65 | '{}-{:.0f}'.format(l.metadata['type'], l.metadata['conf'] * 100), 66 | (int(xs[len(xs) // 2] + pad), int(ys[len(xs) // 2] + pad - 50) + (i + 1) * 20), 67 | fontFace=cv2.FONT_HERSHEY_COMPLEX, 68 | fontScale=0.7, 69 | color=(255, 0, 255)) 70 | return img 71 | 72 | 73 | def nms_func(proposals, scores, nms_thres, nms_topk, conf_threshold=0.2): 74 | if proposals.shape[0] == 0: 75 | return None 76 | above_threshold = scores > conf_threshold 77 | proposals = proposals[above_threshold] 78 | scores = scores[above_threshold] 79 | keep, num_to_keep, _ = nms(proposals, scores, overlap=nms_thres, top_k=nms_topk) 80 | keep = keep[:num_to_keep] 81 | proposals = proposals[keep] 82 | return proposals 83 | 84 | 85 | def proposals_to_pred(proposals): 86 | global anchor_ys 87 | anchor_ys = anchor_ys.to(proposals.device) 88 | anchor_ys = anchor_ys.double() 89 | lanes = [] 90 | for lane in proposals: 91 | lane_xs = lane[(num_lane_types + 3):] / img_w 92 | start = int(round(lane[num_lane_types].item() * n_strips)) 93 | length = int(round(lane[num_lane_types + 2].item())) 94 | end = start + length - 1 95 | end = min(end, len(anchor_ys) - 1) 96 | # end = label_end 97 | # if the proposal does not start at the bottom of the image, 98 | # extend its proposal until the x is outside the image 99 | mask = ~((((lane_xs[:start] >= 0.) & 100 | (lane_xs[:start] <= 1.)).cpu().numpy()[::-1].cumprod()[::-1]).astype(np.bool)) 101 | lane_xs[end + 1:] = -2 102 | lane_xs[:start][mask] = -2 103 | lane_ys = anchor_ys[lane_xs >= 0] 104 | lane_xs = lane_xs[lane_xs >= 0] 105 | lane_xs = lane_xs.flip(0).double() 106 | lane_ys = lane_ys.flip(0) 107 | if len(lane_xs) <= 1: 108 | continue 109 | points = torch.stack((lane_xs.reshape(-1, 1), lane_ys.reshape(-1, 1)), dim=1).squeeze(2) 110 | lane = Lane(points=points.cpu().numpy(), 111 | metadata={ 112 | 'start_x': lane[num_lane_types + 1]/img_w, 113 | 'start_y': lane[num_lane_types]/img_h, 114 | 'conf': lane[:num_lane_types].max(), 115 | 'type': lane[:num_lane_types].argmax() 116 | }) 117 | lanes.append(lane) 118 | return lanes 119 | 120 | 121 | def decode(proposals): 122 | if proposals.shape[0] == 0: 123 | return None 124 | softmax = nn.Softmax(dim=1) 125 | proposals[:, :num_lane_types] = softmax(proposals[:, :num_lane_types]) 126 | proposals[:, num_lane_types + 2] = torch.round(proposals[:, num_lane_types + 2]) 127 | pred = proposals_to_pred(proposals) 128 | return pred 129 | 130 | 131 | def create_video(filename, size, fps=5): 132 | fourcc = cv2.VideoWriter_fourcc(*'MP42') 133 | video = cv2.VideoWriter(filename, fourcc, float(fps), size) 134 | 135 | return video 136 | 137 | 138 | def eval_video(video_path): 139 | model = torch.jit.load(model_path) 140 | videoCapture = cv2.VideoCapture(video_path) 141 | size = (int(videoCapture.get(cv2.CAP_PROP_FRAME_WIDTH)), 142 | int(videoCapture.get(cv2.CAP_PROP_FRAME_HEIGHT))) 143 | video = create_video(os.path.join(os.path.dirname(video_path), "out34.avi"), size=size, fps=15) 144 | with torch.no_grad(): 145 | success, img_data = videoCapture.read() 146 | st = time.time() 147 | while success: 148 | pst = time.time() 149 | image0 = img_data 150 | image = cv2.resize(image0, (640, 360)) 151 | image = image / 255. 152 | image = to_tensor(image.astype(np.float32)) 153 | image = image.unsqueeze(0) 154 | image = image.to(torch.device('cuda')) 155 | proposal, scores = model(image) 156 | result = nms_func(proposal, scores, 50, 4, 0.5) 157 | prediction = decode(result) 158 | print(time.time() -pst) 159 | img = image0 160 | if prediction: 161 | img = draw_annotation(img=img, pred=prediction) 162 | else: 163 | print("NO NO NO") 164 | video.write(img) 165 | 166 | success, img_data = videoCapture.read() 167 | print(time.time() -st) 168 | 169 | 170 | def eval(image_path, save_path): 171 | model = torch.jit.load(model_path) 172 | #video = create_video("wxout.avi", size=(1640,590), fps=25) #culane 173 | #video = create_video("out.avi", size=(1276,717), fps=5) #LLAMAS 174 | with torch.no_grad(): 175 | i = 0 176 | for file in os.listdir(image_path): 177 | i+=1 178 | st = time.time() 179 | file_path = os.path.join(image_path, file) 180 | image0 = cv2.imread(file_path) 181 | image = cv2.resize(image0, (640, 360)) 182 | image = image / 255. 183 | image = to_tensor(image.astype(np.float32)) 184 | image = image.unsqueeze(0) 185 | image = image.to(torch.device('cuda')) 186 | 187 | proposal, scores = model(image) 188 | result = nms_func(proposal, scores, 50, 4, 0.2) 189 | 190 | prediction = decode(result) 191 | print(time.time() -st) 192 | #img = (image.cpu().permute(2, 1, 0).numpy() * 255).astype(np.uint8) 193 | img = image0 194 | if prediction: 195 | img = draw_annotation(img=img, pred=prediction) 196 | cv2.imwrite(os.path.join(save_path,str(i) + "nbv.jpg"), img) 197 | #video.write(img) 198 | 199 | 200 | 201 | #kk = '/home/LaneATT-main/data/video_example' 202 | #kk = "/notebooks/LLAMAS/train/images-2014-12-22-14-01-36_mapping_280N_3rd_lane" 203 | #kk = '/home/LaneATT-main/data/images2' 204 | #save = '/home/LaneATT-main/data/result' 205 | #eval(kk,save) 206 | 207 | vd = '/home/LaneATT-main/data/fig/34.mp4' 208 | eval_video(vd) 209 | -------------------------------------------------------------------------------- /lib/runner_distributed.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import random 3 | import logging 4 | 5 | import cv2 6 | import torch 7 | import numpy as np 8 | from tqdm import tqdm, trange 9 | from lib.distributed_utils import * 10 | import sys 11 | 12 | 13 | class Runner_Distributed: 14 | def __init__(self, cfg, exp, device, resume=False, view=None, deterministic=False, 15 | syncBN=False, distributed=False, device_id=0, main_pod=0): 16 | self.cfg = cfg 17 | self.exp = exp 18 | self.device = device 19 | self.resume = resume 20 | self.view = view 21 | self.SyncBN = syncBN 22 | self.device_id = device_id 23 | self.main_pod = main_pod 24 | self.distributed=distributed 25 | self.logger = logging.getLogger(__name__) 26 | 27 | # Fix seeds 28 | torch.manual_seed(cfg['seed']) 29 | np.random.seed(cfg['seed']) 30 | random.seed(cfg['seed']) 31 | 32 | if deterministic: 33 | torch.backends.cudnn.deterministic = True 34 | torch.backends.cudnn.benchmark = False 35 | 36 | def train(self): 37 | self.exp.train_start_callback(self.cfg) 38 | starting_epoch = 1 39 | model = self.cfg.get_model() 40 | model = model.to(self.device) 41 | optimizer = self.cfg.get_optimizer(model.parameters()) 42 | scheduler = self.cfg.get_lr_scheduler(optimizer) 43 | if self.resume: 44 | last_epoch, model, optimizer, scheduler = self.exp.load_last_train_state(model, optimizer, scheduler, self.device) 45 | starting_epoch = last_epoch + 1 46 | 47 | if self.SyncBN: 48 | model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(self.device) 49 | 50 | # 转为DDP模型 51 | model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[self.device_id]) 52 | 53 | max_epochs = self.cfg['epochs'] 54 | train_loader, train_sampler = self.get_train_dataloader() 55 | loss_parameters = self.cfg.get_loss_parameters() 56 | for epoch in trange(starting_epoch, max_epochs + 1, initial=starting_epoch - 1, total=max_epochs): 57 | train_sampler.set_epoch(epoch) 58 | self.exp.epoch_start_callback(epoch, max_epochs) 59 | model.train() 60 | pbar = train_loader 61 | 62 | if is_main_process(): 63 | pbar = tqdm(train_loader) 64 | 65 | #mean_loss = torch.zeros(1).to(self.device) 66 | for i, (images, labels, _) in enumerate(pbar): 67 | images = images.to(self.device) 68 | labels = labels.to(self.device) 69 | 70 | # Forward pass 71 | outputs = model(images, **self.cfg.get_train_parameters()) 72 | 73 | 74 | """ 75 | DDP出现如AttributeError: 'DDP' object has no attribute 'xxx'的错误。 76 | 77 | 原因：在使用net = torch.nn.DDP(net)之后，原来的net会被封装为新的net的module属性里。 78 | 79 | 解决方案：所有在net = DDP(net)后调用了不是初始化与forward的属性，需要将net替换为net.module.XXX XXX为模型中自定义得函数 80 | """ 81 | ######multy lane 82 | #loss, loss_dict_i = model.loss(outputs, labels, **loss_parameters) 83 | loss, loss_dict_i = model.module.loss(outputs, labels, **loss_parameters) 84 | 85 | # Backward and optimize 86 | optimizer.zero_grad() 87 | loss.backward() 88 | 89 | loss = reduce_value(loss, average=True) 90 | #mean_loss = (mean_loss * i + loss.detach()) / (i + 1) # update mean losses 91 | optimizer.step() 92 | # Scheduler step (iteration based) 93 | scheduler.step() 94 | 95 | # Log 96 | if self.main_pod == 0: 97 | postfix_dict = {key: float(value) for key, value in loss_dict_i.items()} 98 | postfix_dict['lr'] = optimizer.param_groups[0]["lr"] 99 | self.exp.iter_end_callback(epoch, max_epochs, i, len(train_loader), loss.item(), postfix_dict) 100 | postfix_dict['loss'] = loss.item() 101 | pbar.set_postfix(ordered_dict=postfix_dict) 102 | 103 | # 等待所有进程计算完毕 104 | if self.device != torch.device("cpu"): 105 | torch.cuda.synchronize(self.device) 106 | 107 | if self.main_pod == 0: 108 | self.exp.epoch_end_callback(epoch, max_epochs, model, optimizer, scheduler) 109 | 110 | self.exp.train_end_callback() 111 | 112 | 113 | def eval(self, epoch, on_val=False, save_predictions=False): 114 | model = self.cfg.get_model() 115 | model_path = self.exp.get_checkpoint_path(epoch) 116 | self.logger.info('Loading model %s', model_path) 117 | model.load_state_dict(self.exp.get_epoch_model(epoch)) 118 | model = model.to(self.device) 119 | model.eval() 120 | if on_val: 121 | dataloader = self.get_val_dataloader() 122 | else: 123 | dataloader = self.get_test_dataloader() 124 | test_parameters = self.cfg.get_test_parameters() 125 | predictions = [] 126 | self.exp.eval_start_callback(self.cfg) 127 | with torch.no_grad(): 128 | for idx, (images, _, _) in enumerate(tqdm(dataloader)): 129 | images = images.to(self.device) 130 | output = model(images, **test_parameters) 131 | prediction = model.decode(output, as_lanes=True) 132 | predictions.extend(prediction) 133 | 134 | for bb in range(len(images)): 135 | img = (images[bb].cpu().permute(1, 2, 0).numpy() * 255).astype(np.uint8) 136 | img, fp, fn = dataloader.dataset.draw_annotation(idx, img=img, pred=prediction[bb]) 137 | cv2.imwrite('/home/LaneATT-main/XXX/'+'pred_'+ str(idx)+ "_"+str(bb)+".jpg", img) 138 | 139 | 140 | #if self.view: 141 | # img = (images[0].cpu().permute(1, 2, 0).numpy() * 255).astype(np.uint8) 142 | # img, fp, fn = dataloader.dataset.draw_annotation(idx, img=img, pred=prediction[0]) 143 | # if self.view == 'mistakes' and fp == 0 and fn == 0: 144 | # continue 145 | #cv2.imshow('pred', img) 146 | #cv2.waitKey(0) 147 | 148 | if save_predictions: 149 | with open('predictions.pkl', 'wb') as handle: 150 | pickle.dump(predictions, handle, protocol=pickle.HIGHEST_PROTOCOL) 151 | self.exp.eval_end_callback(dataloader.dataset.dataset, predictions, epoch) 152 | 153 | def get_train_dataloader(self): 154 | train_dataset = self.cfg.get_dataset('train') 155 | ######DDP 156 | # 给每个rank对应的进程分配训练的样本索引 157 | train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) 158 | # 将样本索引每batch_size个元素组成一个list 159 | train_batch_sampler = torch.utils.data.BatchSampler( 160 | train_sampler, self.cfg['batch_size'], drop_last=True) 161 | 162 | nw = min([os.cpu_count(), self.cfg['batch_size'] if self.cfg['batch_size'] > 1 else 0, 8]) # number of workers 163 | 164 | train_loader = torch.utils.data.DataLoader(dataset=train_dataset, 165 | batch_sampler=train_batch_sampler, 166 | num_workers=nw, 167 | pin_memory=True, 168 | worker_init_fn=self._worker_init_fn_) 169 | 170 | return train_loader, train_sampler 171 | 172 | def get_test_dataloader(self): 173 | test_dataset = self.cfg.get_dataset('test') 174 | test_loader = torch.utils.data.DataLoader(dataset=test_dataset, 175 | batch_size=self.cfg['batch_size'] if not self.view else 1, 176 | shuffle=False, 177 | num_workers=8, 178 | worker_init_fn=self._worker_init_fn_) 179 | return test_loader 180 | 181 | def get_val_dataloader(self): 182 | val_dataset = self.cfg.get_dataset('val') 183 | val_loader = torch.utils.data.DataLoader(dataset=val_dataset, 184 | batch_size=self.cfg['batch_size'], 185 | shuffle=False, 186 | num_workers=8, 187 | worker_init_fn=self._worker_init_fn_) 188 | return val_loader 189 | 190 | @staticmethod 191 | def _worker_init_fn_(_): 192 | torch_seed = torch.initial_seed() 193 | np_seed = torch_seed // 2**32 - 1 194 | random.seed(torch_seed) 195 | np.random.seed(np_seed) 196 | -------------------------------------------------------------------------------- /utils/llamas_metric.py: -------------------------------------------------------------------------------- 1 | """ Evaluation script for the CULane metric on the LLAMAS dataset. 2 | 3 | This script will compute the F1, precision and recall metrics as described in the CULane benchmark. 4 | 5 | The predictions format is the same one used in the CULane benchmark. 6 | In summary, for every annotation file: 7 | labels/a/b/c.json 8 | There should be a prediction file: 9 | predictions/a/b/c.lines.txt 10 | Inside each .lines.txt file each line will contain a sequence of points (x, y) separated by spaces. 11 | For more information, please see https://xingangpan.github.io/projects/CULane.html 12 | 13 | This script uses two methods to compute the IoU: one using an image to draw the lanes (named `discrete` here) and 14 | another one that uses shapes with the shapely library (named `continuous` here). The results achieved with the first 15 | method are very close to the official CULane implementation. Although the second should be a more exact method and is 16 | faster to compute, it deviates more from the official implementation. By default, the method closer to the official 17 | metric is used. 18 | """ 19 | 20 | import os 21 | import argparse 22 | from functools import partial 23 | 24 | import cv2 25 | import numpy as np 26 | from p_tqdm import t_map, p_map 27 | from scipy.interpolate import splprep, splev 28 | from scipy.optimize import linear_sum_assignment 29 | from shapely.geometry import LineString, Polygon 30 | 31 | from . import llamas_utils 32 | 33 | LLAMAS_IMG_RES = (717, 1276) 34 | 35 | 36 | def add_ys(xs): 37 | """For each x in xs, make a tuple with x and its corresponding y.""" 38 | xs = np.array(xs[300:]) 39 | valid = xs >= 0 40 | xs = xs[valid] 41 | assert len(xs) > 1 42 | ys = np.arange(300, 717)[valid] 43 | return list(zip(xs, ys)) 44 | 45 | 46 | def draw_lane(lane, img=None, img_shape=None, width=30): 47 | """Draw a lane (a list of points) on an image by drawing a line with width `width` through each 48 | pair of points i and i+i""" 49 | if img is None: 50 | img = np.zeros(img_shape, dtype=np.uint8) 51 | lane = lane.astype(np.int32) 52 | for p1, p2 in zip(lane[:-1], lane[1:]): 53 | cv2.line(img, tuple(p1), tuple(p2), color=(1,), thickness=width) 54 | return img 55 | 56 | 57 | def discrete_cross_iou(xs, ys, width=30, img_shape=LLAMAS_IMG_RES): 58 | """For each lane in xs, compute its Intersection Over Union (IoU) with each lane in ys by drawing the lanes on 59 | an image""" 60 | xs = [draw_lane(lane, img_shape=img_shape, width=width) > 0 for lane in xs] 61 | ys = [draw_lane(lane, img_shape=img_shape, width=width) > 0 for lane in ys] 62 | 63 | ious = np.zeros((len(xs), len(ys))) 64 | for i, x in enumerate(xs): 65 | for j, y in enumerate(ys): 66 | # IoU by the definition: sum all intersections (binary and) and divide by the sum of the union (binary or) 67 | ious[i, j] = (x & y).sum() / (x | y).sum() 68 | return ious 69 | 70 | 71 | def continuous_cross_iou(xs, ys, width=30): 72 | """For each lane in xs, compute its Intersection Over Union (IoU) with each lane in ys using the area between each 73 | pair of points""" 74 | h, w = IMAGE_HEIGHT, IMAGE_WIDTH 75 | image = Polygon([(0, 0), (0, h - 1), (w - 1, h - 1), (w - 1, 0)]) 76 | xs = [LineString(lane).buffer(distance=width / 2., cap_style=1, join_style=2).intersection(image) for lane in xs] 77 | ys = [LineString(lane).buffer(distance=width / 2., cap_style=1, join_style=2).intersection(image) for lane in ys] 78 | 79 | ious = np.zeros((len(xs), len(ys))) 80 | for i, x in enumerate(xs): 81 | for j, y in enumerate(ys): 82 | ious[i, j] = x.intersection(y).area / x.union(y).area 83 | 84 | return ious 85 | 86 | 87 | def interpolate_lane(points, n=50): 88 | """Spline interpolation of a lane. Used on the predictions""" 89 | x = [x for x, _ in points] 90 | y = [y for _, y in points] 91 | tck, _ = splprep([x, y], s=0, t=n, k=min(3, len(points) - 1)) 92 | 93 | u = np.linspace(0., 1., n) 94 | return np.array(splev(u, tck)).T 95 | 96 | 97 | def culane_metric(pred, anno, width=30, iou_threshold=0.5, unofficial=False, img_shape=LLAMAS_IMG_RES): 98 | """Computes CULane's metric for a single image""" 99 | if len(pred) == 0: 100 | return 0, 0, len(anno) 101 | if len(anno) == 0: 102 | return 0, len(pred), 0 103 | interp_pred = np.array([interpolate_lane(pred_lane, n=50) for pred_lane in pred]) # (4, 50, 2) 104 | anno = np.array([np.array(anno_lane) for anno_lane in anno], dtype=object) 105 | 106 | if unofficial: 107 | ious = continuous_cross_iou(interp_pred, anno, width=width) 108 | else: 109 | ious = discrete_cross_iou(interp_pred, anno, width=width, img_shape=img_shape) 110 | 111 | row_ind, col_ind = linear_sum_assignment(1 - ious) 112 | tp = int((ious[row_ind, col_ind] > iou_threshold).sum()) 113 | fp = len(pred) - tp 114 | fn = len(anno) - tp 115 | return tp, fp, fn 116 | 117 | 118 | def load_prediction(path): 119 | """Loads an image's predictions 120 | Returns a list of lanes, where each lane is a list of points (x,y) 121 | """ 122 | with open(path, 'r') as data_file: 123 | img_data = data_file.readlines() 124 | img_data = [line.split() for line in img_data] 125 | img_data = [list(map(float, lane)) for lane in img_data] 126 | img_data = [[(lane[i], lane[i + 1]) for i in range(0, len(lane), 2)] for lane in img_data] 127 | img_data = [lane for lane in img_data if len(lane) >= 2] 128 | 129 | return img_data 130 | 131 | 132 | def load_prediction_list(label_paths, pred_dir): 133 | return [load_prediction(os.path.join(pred_dir, path.replace('.json', '.lines.txt'))) for path in label_paths] 134 | 135 | 136 | def load_labels(label_dir): 137 | """Loads the annotations and its paths 138 | Each annotation is converted to a list of points (x, y) 139 | """ 140 | label_paths = llamas_utils.get_files_from_folder(label_dir, '.json') 141 | annos = [ 142 | [ 143 | add_ys(xs) for xs in llamas_utils.get_horizontal_values_for_four_lanes(label_path) 144 | if (np.array(xs) >= 0).sum() > 1 145 | ] # lanes annotated with a single point are ignored 146 | for label_path in label_paths 147 | ] 148 | label_paths = [llamas_utils.get_label_base(p) for p in label_paths] 149 | return np.array(annos, dtype=object), np.array(label_paths, dtype=object) 150 | 151 | 152 | def eval_predictions(pred_dir, anno_dir, width=30, unofficial=True, sequential=False): 153 | """Evaluates the predictions in pred_dir and returns CULane's metrics (precision, recall, F1 and its components)""" 154 | print(f'Loading annotation data ({anno_dir})...') 155 | annotations, label_paths = load_labels(anno_dir) 156 | print(f'Loading prediction data ({pred_dir})...') 157 | predictions = load_prediction_list(label_paths, pred_dir) 158 | print('Calculating metric {}...'.format('sequentially' if sequential else 'in parallel')) 159 | if sequential: 160 | results = t_map(partial(culane_metric, width=width, unofficial=unofficial, img_shape=LLAMAS_IMG_RES), 161 | predictions, annotations) 162 | else: 163 | results = p_map(partial(culane_metric, width=width, unofficial=unofficial, img_shape=LLAMAS_IMG_RES), 164 | predictions, annotations) 165 | total_tp = sum(tp for tp, _, _ in results) 166 | total_fp = sum(fp for _, fp, _ in results) 167 | total_fn = sum(fn for _, _, fn in results) 168 | if total_tp == 0: 169 | precision = 0 170 | recall = 0 171 | f1 = 0 172 | else: 173 | precision = float(total_tp) / (total_tp + total_fp) 174 | recall = float(total_tp) / (total_tp + total_fn) 175 | f1 = 2 * precision * recall / (precision + recall) 176 | 177 | return {'TP': total_tp, 'FP': total_fp, 'FN': total_fn, 'Precision': precision, 'Recall': recall, 'F1': f1} 178 | 179 | 180 | def parse_args(): 181 | parser = argparse.ArgumentParser(description="Measure CULane's metric on the LLAMAS dataset") 182 | parser.add_argument("--pred_dir", help="Path to directory containing the predicted lanes", required=True) 183 | parser.add_argument("--anno_dir", help="Path to directory containing the annotated lanes", required=True) 184 | parser.add_argument("--width", type=int, default=30, help="Width of the lane") 185 | parser.add_argument("--sequential", action='store_true', help="Run sequentially instead of in parallel") 186 | parser.add_argument("--unofficial", action='store_true', help="Use a faster but unofficial algorithm") 187 | 188 | return parser.parse_args() 189 | 190 | 191 | def main(): 192 | args = parse_args() 193 | results = eval_predictions(args.pred_dir, 194 | args.anno_dir, 195 | width=args.width, 196 | unofficial=args.unofficial, 197 | sequential=args.sequential) 198 | 199 | header = '=' * 20 + ' Results' + '=' * 20 200 | print(header) 201 | for metric, value in results.items(): 202 | if isinstance(value, float): 203 | print('{}: {:.4f}'.format(metric, value)) 204 | else: 205 | print('{}: {}'.format(metric, value)) 206 | print('=' * len(header)) 207 | 208 | 209 | if __name__ == '__main__': 210 | main() 211 | -------------------------------------------------------------------------------- /lib/datasets/lane_dataset.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | import os 4 | import cv2 5 | import numpy as np 6 | import imgaug.augmenters as iaa 7 | from imgaug.augmenters import Resize 8 | from torchvision.transforms import ToTensor 9 | from torch.utils.data.dataset import Dataset 10 | from scipy.interpolate import InterpolatedUnivariateSpline 11 | from imgaug.augmentables.lines import LineString, LineStringsOnImage 12 | 13 | 14 | from lib.lane import Lane 15 | 16 | from .culane import CULane 17 | from .tusimple import TuSimple 18 | from .llamas import LLAMAS 19 | from .nolabel_dataset import NoLabelDataset 20 | from .ehl_wx import EHLWX 21 | import pdb 22 | 23 | 24 | GT_COLOR = (255, 0, 0) 25 | PRED_HIT_COLOR = (0, 255, 0) 26 | PRED_MISS_COLOR = (0, 0, 255) 27 | IMAGENET_MEAN = np.array([0.485, 0.456, 0.406]) 28 | IMAGENET_STD = np.array([0.229, 0.224, 0.225]) 29 | 30 | 31 | class LaneDataset(Dataset): 32 | def __init__(self, 33 | lane_num_types=None, 34 | S=72, 35 | dataset='tusimple', 36 | augmentations=None, 37 | normalize=False, 38 | img_size=(360, 640), 39 | aug_chance=1., 40 | **kwargs): 41 | super(LaneDataset, self).__init__() 42 | self.use_dataset = dataset 43 | if dataset == 'tusimple': 44 | self.dataset = TuSimple(**kwargs) 45 | elif dataset == 'culane': 46 | self.dataset = CULane(**kwargs) 47 | elif dataset == 'llamas': 48 | self.dataset = LLAMAS(**kwargs) 49 | elif dataset == 'ehl_wx': 50 | self.dataset = EHLWX(**kwargs) 51 | elif dataset == 'nolabel_dataset': 52 | self.dataset = NoLabelDataset(**kwargs) 53 | else: 54 | raise NotImplementedError() 55 | self.lane_types = lane_num_types 56 | self.n_strips = S - 1 57 | self.n_offsets = S 58 | self.normalize = normalize 59 | self.img_h, self.img_w = img_size 60 | self.strip_size = self.img_h / self.n_strips 61 | self.logger = logging.getLogger(__name__) 62 | 63 | # y at each x offset 64 | self.offsets_ys = np.arange(self.img_h, -1, -self.strip_size) 65 | self.transform_annotations() 66 | 67 | if augmentations is not None: 68 | # add augmentations 69 | augmentations = [getattr(iaa, aug['name'])(**aug['parameters']) 70 | for aug in augmentations] # add augmentation 71 | else: 72 | augmentations = [] 73 | 74 | transformations = iaa.Sequential([Resize({'height': self.img_h, 'width': self.img_w})]) 75 | self.to_tensor = ToTensor() 76 | self.transform = iaa.Sequential([iaa.Sometimes(then_list=augmentations, p=aug_chance), transformations]) 77 | self.max_lanes = self.dataset.max_lanes 78 | 79 | @property 80 | def annotations(self): 81 | return self.dataset.annotations 82 | 83 | def transform_annotations(self): 84 | self.logger.info("Transforming annotations to the model's target format...") 85 | self.dataset.annotations = np.array(list(map(self.transform_annotation, self.dataset.annotations))) 86 | self.logger.info('Done.') 87 | 88 | def filter_lane(self, lane): 89 | assert lane[-1][1] <= lane[0][1] 90 | filtered_lane = [] 91 | used = set() 92 | for p in lane: 93 | if p[1] not in used: 94 | filtered_lane.append(p) 95 | used.add(p[1]) 96 | 97 | return filtered_lane 98 | 99 | def transform_annotation(self, anno, img_wh=None): 100 | if img_wh is None: 101 | img_h = self.dataset.get_img_heigth() 102 | img_w = self.dataset.get_img_width() 103 | else: 104 | img_w, img_h = img_wh 105 | 106 | old_lanes = anno['lanes'] 107 | 108 | old_lanes_type = anno['lane_type'] 109 | 110 | # removing lanes with less than 2 points 111 | old_lanes_type = [old_lanes_type[ind] for ind,lane in enumerate(old_lanes) if len(lane) > 1] 112 | old_lanes = filter(lambda x: len(x) > 1, old_lanes) 113 | # sort lane points by Y (bottom to top of the image) 114 | old_lanes = [sorted(lane, key=lambda x: -x[1]) for lane in old_lanes] 115 | # remove points with same Y (keep first occurrence) 116 | old_lanes = [self.filter_lane(lane) for lane in old_lanes] 117 | # normalize the annotation coordinates 118 | old_lanes = [[[x * self.img_w / float(img_w), y * self.img_h / float(img_h)] for x, y in lane] 119 | for lane in old_lanes] 120 | 121 | # create tranformed annotations 122 | # lanes = np.ones((self.dataset.max_lanes, 2 + 1 + 1 + 1 + self.n_offsets), 123 | # dtype=np.float32) * -1e5 # 2 scores, 1 start_y, 1 start_x, 1 length, S+1 coordinates 124 | 125 | ##### create multily type lane tranformed 126 | lanes = np.ones((self.dataset.max_lanes, self.lane_types + 1 + 1 + 1 + self.n_offsets), 127 | dtype=np.float32) * -1e5 # 3 scores, 1 start_y, 1 start_x, 1 length, S+1 coordinates 128 | 129 | # lanes are invalid by default 130 | # lanes[:, 0] = 1 131 | # lanes[:, 1] = 0 132 | # #####多类别type须在这里增加维度 133 | # lanes[:, 2] = 0 134 | lane_type_onehot = np.eye(self.lane_types) 135 | lanes[:, :self.lane_types] = lane_type_onehot[0] 136 | 137 | for lane_idx, lane in enumerate(old_lanes): 138 | try: 139 | xs_outside_image, xs_inside_image = self.sample_lane(lane, self.offsets_ys) 140 | except AssertionError: 141 | continue 142 | if len(xs_inside_image) == 0: 143 | continue 144 | all_xs = np.hstack((xs_outside_image, xs_inside_image)) #水平堆叠 145 | 146 | ########多类别type须在这里增加 维度 类似 onehot 147 | # lanes[lane_idx, 0] = 0 # lane type scores 148 | # lanes[lane_idx, 1] = 0 # lane type 149 | # lanes[lane_idx, 2] = 1 150 | lanes[lane_idx, :self.lane_types] = lane_type_onehot[old_lanes_type[lane_idx]] 151 | 152 | lanes[lane_idx, self.lane_types] = len(xs_outside_image) / self.n_strips # start_y 153 | lanes[lane_idx, self.lane_types+1] = xs_inside_image[0] # start_x 154 | lanes[lane_idx, self.lane_types+2] = len(xs_inside_image) # lane length 155 | lanes[lane_idx, (self.lane_types+3):(self.lane_types+3+len(all_xs))] = all_xs # all y 对应 all x points 156 | 157 | 158 | # lanes[lane_idx, 0] = 0 # lane type scores 159 | # lanes[lane_idx, 1] = 1 # lane type 160 | # lanes[lane_idx, 2] = len(xs_outside_image) / self.n_strips # start_y 161 | # lanes[lane_idx, 3] = xs_inside_image[0] # start_x 162 | # lanes[lane_idx, 4] = len(xs_inside_image) # lane length 163 | # lanes[lane_idx, 5:5 + len(all_xs)] = all_xs # all y 对应 all x points 164 | 165 | #######multly lanetypes 166 | if self.use_dataset == 'culane': 167 | new_anno = {'path': os.path.join(self.dataset.root, anno['org_path']), 'label': lanes, 'lane_type': old_lanes_type, 'old_anno': anno} 168 | else: 169 | new_anno = {'path': anno['path'], 'label': lanes, 'lane_type': old_lanes_type, 'old_anno': anno} 170 | return new_anno 171 | 172 | def sample_lane(self, points, sample_ys): 173 | """ 174 | 175 | :param points: 原始数据集中lane 点集 176 | :param sample_ys: 新定义的lane 点集中 y集合 177 | :return: 178 | """ 179 | # this function expects the points to be sorted 180 | points = np.array(points) 181 | if not np.all(points[1:, 1] < points[:-1, 1]): 182 | raise Exception('Annotaion points have to be sorted') 183 | x, y = points[:, 0], points[:, 1] 184 | 185 | # interpolate points inside domain 186 | assert len(points) > 1 187 | interp = InterpolatedUnivariateSpline(y[::-1], x[::-1], k=min(3, len(points) - 1)) #将原始标注的lane进行拟合 188 | domain_min_y = y.min() 189 | domain_max_y = y.max() 190 | sample_ys_inside_domain = sample_ys[(sample_ys >= domain_min_y) & (sample_ys <= domain_max_y)] # 在原始数据label上转化为新定义的线段中 点 对应的 y 值 191 | assert len(sample_ys_inside_domain) > 0 192 | interp_xs = interp(sample_ys_inside_domain) # 通过拟合 方程 通过 y 求出对应 x 193 | 194 | # extrapolate lane to the bottom of the image with a straight line using the 2 points closest to the bottom 195 | #通过原始数据集中 lane 点集中 起始两点(自下而上) 向下延长 去拟合 新定义的点集 （可解决遮挡线问题） 196 | two_closest_points = points[:2] 197 | extrap = np.polyfit(two_closest_points[:, 1], two_closest_points[:, 0], deg=1) 198 | extrap_ys = sample_ys[sample_ys > domain_max_y] 199 | extrap_xs = np.polyval(extrap, extrap_ys) 200 | all_xs = np.hstack((extrap_xs, interp_xs)) 201 | 202 | # separate between inside and outside points 203 | inside_mask = (all_xs >= 0) & (all_xs < self.img_w) # 去掉超过 x 的边界点 204 | xs_inside_image = all_xs[inside_mask] 205 | xs_outside_image = all_xs[~inside_mask] 206 | 207 | return xs_outside_image, xs_inside_image 208 | 209 | def draw_annotation(self, idx, label=None, pred=None, img=None): 210 | # Get image if not provided 211 | if img is None: 212 | img, label, _ = self.__getitem__(idx) 213 | label = self.label_to_lanes3(label) 214 | img = img.permute(1, 2, 0).numpy() 215 | if self.normalize: 216 | img = img * np.array(IMAGENET_STD) + np.array(IMAGENET_MEAN) 217 | img = (img * 255).astype(np.uint8) 218 | else: 219 | _, label, _ = self.__getitem__(idx) 220 | label = self.label_to_lanes3(label) 221 | img = cv2.resize(img, (self.img_w, self.img_h)) 222 | 223 | img_h, _, _ = img.shape 224 | # Pad image to visualize extrapolated predictions 225 | pad = 0 226 | if pad > 0: 227 | img_pad = np.zeros((self.img_h + 2 * pad, self.img_w + 2 * pad, 3), dtype=np.uint8) 228 | img_pad[pad:-pad, pad:-pad, :] = img 229 | img = img_pad 230 | data = [(None, None, label)] 231 | if pred is not None: 232 | fp, fn, matches, accs = self.dataset.get_metrics(pred, idx) 233 | assert len(matches) == len(pred) 234 | data.append((matches, accs, pred)) 235 | else: 236 | fp = fn = None 237 | for matches, accs, datum in data: 238 | for i, l in enumerate(datum): 239 | if l.metadata['type'] == 0: 240 | continue 241 | if matches is None: 242 | color = GT_COLOR 243 | elif matches[i]: 244 | color = PRED_HIT_COLOR 245 | else: 246 | color = PRED_MISS_COLOR 247 | points = l.points 248 | points[:, 0] *= img.shape[1] 249 | points[:, 1] *= img.shape[0] 250 | points = points.round().astype(int) 251 | points += pad 252 | xs, ys = points[:, 0], points[:, 1] 253 | for curr_p, next_p in zip(points[:-1], points[1:]): 254 | img = cv2.line(img, 255 | tuple(curr_p), 256 | tuple(next_p), 257 | color=color, 258 | thickness=3 if matches is None else 3) 259 | if 'start_x' in l.metadata: 260 | start_x = l.metadata['start_x'] * img.shape[1] 261 | start_y = l.metadata['start_y'] * img.shape[0] 262 | cv2.circle(img, (int(start_x + pad), int(img_h - 1 - start_y + pad)), 263 | radius=5, 264 | color=(0, 0, 255), 265 | thickness=-1) 266 | if len(xs) == 0: 267 | print("Empty pred") 268 | if len(xs) > 0 and accs is not None: 269 | cv2.putText(img, 270 | '{}-{:.0f}'.format(l.metadata['type'], l.metadata['conf'] * 100), 271 | (int(xs[len(xs) // 2] + pad), int(ys[len(xs) // 2] + pad - 50) + (i + 1) * 20), 272 | fontFace=cv2.FONT_HERSHEY_COMPLEX, 273 | fontScale=0.7, 274 | color=(255, 0, 255)) 275 | return img, fp, fn 276 | 277 | def label_to_lanes(self, label): 278 | lanes = [] 279 | for l in label: 280 | if l[1] == 0: 281 | continue 282 | xs = l[6:] / self.img_w 283 | ys = self.offsets_ys / self.img_h 284 | start = int(round(l[3] * self.n_strips)) 285 | length = int(round(l[5])) 286 | xs = xs[start:start + length][::-1] 287 | ys = ys[start:start + length][::-1] 288 | xs = xs.reshape(-1, 1) 289 | ys = ys.reshape(-1, 1) 290 | points = np.hstack((xs, ys)) 291 | 292 | lanes.append(Lane(points=points)) 293 | return lanes 294 | 295 | def lane_to_linestrings(self, lanes): 296 | lines = [] 297 | for lane in lanes: 298 | lines.append(LineString(lane)) 299 | 300 | return lines 301 | 302 | def linestrings_to_lanes(self, lines): 303 | lanes = [] 304 | for line in lines: 305 | lanes.append(line.coords) 306 | 307 | return lanes 308 | 309 | def __getitem__(self, idx): 310 | item = self.dataset[idx] 311 | img_org = cv2.imread(item['path']) 312 | 313 | self.use_dataset = None #### 用于culane 数据增强 314 | 315 | try: 316 | line_strings_org = self.lane_to_linestrings(item['old_anno']['lanes']) 317 | line_strings_org = LineStringsOnImage(line_strings_org, shape=img_org.shape) 318 | except: 319 | print(item['path']) 320 | for i in range(30): 321 | img, line_strings = self.transform(image=img_org.copy(), line_strings=line_strings_org) 322 | line_strings.clip_out_of_image_() 323 | #######multly label 324 | new_anno = {'path': item['path'], 'lanes': self.linestrings_to_lanes(line_strings), 'lane_type': item['old_anno']['lane_type']} 325 | try: 326 | label = self.transform_annotation(new_anno, img_wh=(self.img_w, self.img_h))['label'] 327 | break 328 | except: 329 | if (i + 1) == 30: 330 | self.logger.critical('Transform annotation failed 30 times :(') 331 | exit() 332 | 333 | img = img / 255. 334 | if self.normalize: 335 | img = (img - IMAGENET_MEAN) / IMAGENET_STD 336 | img = self.to_tensor(img.astype(np.float32)) 337 | return (img, label, idx) 338 | 339 | def __len__(self): 340 | return len(self.dataset) 341 | 342 | -------------------------------------------------------------------------------- /utils/llamas_utils.py: -------------------------------------------------------------------------------- 1 | # All following lines (which were slightly modified) were taken from: https://github.com/karstenBehrendt/unsupervised_llamas 2 | # Its license is copied here 3 | 4 | # ##### Begin License ###### 5 | # MIT License 6 | 7 | # Copyright (c) 2019 Karsten Behrendt, Robert Bosch LLC 8 | 9 | # Permission is hereby granted, free of charge, to any person obtaining a copy 10 | # of this software and associated documentation files (the "Software"), to deal 11 | # in the Software without restriction, including without limitation the rights 12 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 13 | # copies of the Software, and to permit persons to whom the Software is 14 | # furnished to do so, subject to the following conditions: 15 | 16 | # The above copyright notice and this permission notice shall be included in all 17 | # copies or substantial portions of the Software. 18 | 19 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 20 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 21 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 22 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 23 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 24 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 25 | # SOFTWARE. 26 | 27 | # ##### End License ###### 28 | 29 | # Start code under the previous license 30 | 31 | import os 32 | import json 33 | import numpy as np 34 | 35 | 36 | def get_files_from_folder(directory, extension=None): 37 | """Get all files within a folder that fit the extension """ 38 | # NOTE Can be replaced by glob for newer python versions 39 | label_files = [] 40 | for root, _, files in os.walk(directory): 41 | for some_file in files: 42 | label_files.append(os.path.abspath(os.path.join(root, some_file))) 43 | if extension is not None: 44 | label_files = list(filter(lambda x: x.endswith(extension), label_files)) 45 | return label_files 46 | 47 | 48 | def get_label_base(label_path): 49 | """ Gets directory independent label path """ 50 | return '/'.join(label_path.split('/')[-2:]) 51 | 52 | 53 | def get_labels(dataset_root, split='test'): 54 | """ Gets label files of specified dataset split """ 55 | label_paths = get_files_from_folder(os.path.join(dataset_root, split), '.json') 56 | return label_paths 57 | 58 | 59 | def _extend_lane(lane, projection_matrix): 60 | """Extends marker closest to the camera 61 | 62 | Adds an extra marker that reaches the end of the image 63 | 64 | Parameters 65 | ---------- 66 | lane : iterable of markers 67 | projection_matrix : 3x3 projection matrix 68 | """ 69 | # Unfortunately, we did not store markers beyond the image plane. That hurts us now 70 | # z is the orthongal distance to the car. It's good enough 71 | 72 | # The markers are automatically detected, mapped, and labeled. There exist faulty ones, 73 | # e.g., horizontal markers which need to be filtered 74 | filtered_markers = filter( 75 | lambda x: (x['pixel_start']['y'] != x['pixel_end']['y'] and x['pixel_start']['x'] != x['pixel_end']['x']), 76 | lane['markers']) 77 | # might be the first marker in the list but not guaranteed 78 | closest_marker = min(filtered_markers, key=lambda x: x['world_start']['z']) 79 | 80 | if closest_marker['world_start']['z'] < 0: # This one likely equals "if False" 81 | return lane 82 | 83 | # World marker extension approximation 84 | x_gradient = (closest_marker['world_end']['x'] - closest_marker['world_start']['x']) /\ 85 | (closest_marker['world_end']['z'] - closest_marker['world_start']['z']) 86 | y_gradient = (closest_marker['world_end']['y'] - closest_marker['world_start']['y']) /\ 87 | (closest_marker['world_end']['z'] - closest_marker['world_start']['z']) 88 | 89 | zero_x = closest_marker['world_start']['x'] - (closest_marker['world_start']['z'] - 1) * x_gradient 90 | zero_y = closest_marker['world_start']['y'] - (closest_marker['world_start']['z'] - 1) * y_gradient 91 | 92 | # Pixel marker extension approximation 93 | pixel_x_gradient = (closest_marker['pixel_end']['x'] - closest_marker['pixel_start']['x']) /\ 94 | (closest_marker['pixel_end']['y'] - closest_marker['pixel_start']['y']) 95 | pixel_y_gradient = (closest_marker['pixel_end']['y'] - closest_marker['pixel_start']['y']) /\ 96 | (closest_marker['pixel_end']['x'] - closest_marker['pixel_start']['x']) 97 | 98 | pixel_zero_x = closest_marker['pixel_start']['x'] + (716 - closest_marker['pixel_start']['y']) * pixel_x_gradient 99 | if pixel_zero_x < 0: 100 | left_y = closest_marker['pixel_start']['y'] - closest_marker['pixel_start']['x'] * pixel_y_gradient 101 | new_pixel_point = (0, left_y) 102 | elif pixel_zero_x > 1276: 103 | right_y = closest_marker['pixel_start']['y'] + (1276 - closest_marker['pixel_start']['x']) * pixel_y_gradient 104 | new_pixel_point = (1276, right_y) 105 | else: 106 | new_pixel_point = (pixel_zero_x, 716) 107 | 108 | new_marker = { 109 | 'lane_marker_id': 'FAKE', 110 | 'world_end': { 111 | 'x': closest_marker['world_start']['x'], 112 | 'y': closest_marker['world_start']['y'], 113 | 'z': closest_marker['world_start']['z'] 114 | }, 115 | 'world_start': { 116 | 'x': zero_x, 117 | 'y': zero_y, 118 | 'z': 1 119 | }, 120 | 'pixel_end': { 121 | 'x': closest_marker['pixel_start']['x'], 122 | 'y': closest_marker['pixel_start']['y'] 123 | }, 124 | 'pixel_start': { 125 | 'x': ir(new_pixel_point[0]), 126 | 'y': ir(new_pixel_point[1]) 127 | } 128 | } 129 | lane['markers'].insert(0, new_marker) 130 | 131 | return lane 132 | 133 | 134 | class SplineCreator(): 135 | """ 136 | For each lane divder 137 | - all lines are projected 138 | - linearly interpolated to limit oscillations 139 | - interpolated by a spline 140 | - subsampled to receive individual pixel values 141 | 142 | The spline creation can be optimized! 143 | - Better spline parameters 144 | - Extend lowest marker to reach bottom of image would also help 145 | - Extending last marker may in some cases be interesting too 146 | Any help is welcome. 147 | 148 | Call create_all_points and get the points in self.sampled_points 149 | It has an x coordinate for each value for each lane 150 | 151 | """ 152 | def __init__(self, json_path): 153 | self.json_path = json_path 154 | self.json_content = read_json(json_path) 155 | self.lanes = self.json_content['lanes'] 156 | self.lane_marker_points = {} 157 | self.sampled_points = {} # <--- the interesting part 158 | self.debug_image = np.zeros((717, 1276, 3), dtype=np.uint8) 159 | 160 | def _sample_points(self, lane, ypp=5, between_markers=True): 161 | """ Markers are given by start and endpoint. This one adds extra points 162 | which need to be considered for the interpolation. Otherwise the spline 163 | could arbitrarily oscillate between start and end of the individual markers 164 | 165 | Parameters 166 | ---------- 167 | lane: polyline, in theory but there are artifacts which lead to inconsistencies 168 | in ordering. There may be parallel lines. The lines may be dashed. It's messy. 169 | ypp: y-pixels per point, e.g. 10 leads to a point every ten pixels 170 | between_markers : bool, interpolates inbetween dashes 171 | 172 | Notes 173 | ----- 174 | Especially, adding points in the lower parts of the image (high y-values) because 175 | the start and end points are too sparse. 176 | Removing upper lane markers that have starting and end points mapped into the same pixel. 177 | """ 178 | 179 | # Collect all x values from all markers along a given line. There may be multiple 180 | # intersecting markers, i.e., multiple entries for some y values 181 | x_values = [[] for i in range(717)] 182 | for marker in lane['markers']: 183 | x_values[marker['pixel_start']['y']].append(marker['pixel_start']['x']) 184 | 185 | height = marker['pixel_start']['y'] - marker['pixel_end']['y'] 186 | if height > 2: 187 | slope = (marker['pixel_end']['x'] - marker['pixel_start']['x']) / height 188 | step_size = (marker['pixel_start']['y'] - marker['pixel_end']['y']) / float(height) 189 | for i in range(height + 1): 190 | x = marker['pixel_start']['x'] + slope * step_size * i 191 | y = marker['pixel_start']['y'] - step_size * i 192 | x_values[ir(y)].append(ir(x)) 193 | 194 | # Calculate average x values for each y value 195 | for y, xs in enumerate(x_values): 196 | if not xs: 197 | x_values[y] = -1 198 | else: 199 | x_values[y] = sum(xs) / float(len(xs)) 200 | 201 | # In the following, we will only interpolate between markers if needed 202 | if not between_markers: 203 | return x_values # TODO ypp 204 | 205 | # # interpolate between markers 206 | current_y = 0 207 | while x_values[current_y] == -1: # skip missing first entries 208 | current_y += 1 209 | 210 | # Also possible using numpy.interp when accounting for beginning and end 211 | next_set_y = 0 212 | try: 213 | while current_y < 717: 214 | if x_values[current_y] != -1: # set. Nothing to be done 215 | current_y += 1 216 | continue 217 | 218 | # Finds target x value for interpolation 219 | while next_set_y <= current_y or x_values[next_set_y] == -1: 220 | next_set_y += 1 221 | if next_set_y >= 717: 222 | raise StopIteration 223 | 224 | x_values[current_y] = x_values[current_y - 1] + (x_values[next_set_y] - x_values[current_y - 1]) /\ 225 | (next_set_y - current_y + 1) 226 | current_y += 1 227 | 228 | except StopIteration: 229 | pass # Done with lane 230 | 231 | return x_values 232 | 233 | def _lane_points_fit(self, lane): 234 | # TODO name and docstring 235 | """ Fits spline in image space for the markers of a single lane (side) 236 | 237 | Parameters 238 | ---------- 239 | lane: dict as specified in label 240 | 241 | Returns 242 | ------- 243 | Pixel level values for curve along the y-axis 244 | 245 | Notes 246 | ----- 247 | This one can be drastically improved. Probably fairly easy as well. 248 | """ 249 | # NOTE all variable names represent image coordinates, interpolation coordinates are swapped! 250 | lane = _extend_lane(lane, self.json_content['projection_matrix']) 251 | sampled_points = self._sample_points(lane, ypp=1) 252 | self.sampled_points[lane['lane_id']] = sampled_points 253 | 254 | return sampled_points 255 | 256 | def create_all_points(self, ): 257 | """ Creates splines for given label """ 258 | for lane in self.lanes: 259 | self._lane_points_fit(lane) 260 | 261 | 262 | def get_horizontal_values_for_four_lanes(json_path): 263 | """ Gets an x value for every y coordinate for l1, l0, r0, r1 264 | 265 | This allows to easily train a direct curve approximation. For each value along 266 | the y-axis, the respective x-values can be compared, e.g. squared distance. 267 | Missing values are filled with -1. Missing values are values missing from the spline. 268 | There is no extrapolation to the image start/end (yet). 269 | But values are interpolated between markers. Space between dashed markers is not missing. 270 | 271 | Parameters 272 | ---------- 273 | json_path: str 274 | path to label-file 275 | 276 | Returns 277 | ------- 278 | List of [l1, l0, r0, r1], each of which represents a list of ints the length of 279 | the number of vertical pixels of the image 280 | 281 | Notes 282 | ----- 283 | The points are currently based on the splines. The splines are interpolated based on the 284 | segmentation values. The spline interpolation has lots of room for improvement, e.g. 285 | the lines could be interpolated in 3D, a better approach to spline interpolation could 286 | be used, there is barely any error checking, sometimes the splines oscillate too much. 287 | This was used for a quick poly-line regression training only. 288 | """ 289 | 290 | sc = SplineCreator(json_path) 291 | sc.create_all_points() 292 | 293 | l1 = sc.sampled_points.get('l1', [-1] * 717) 294 | l0 = sc.sampled_points.get('l0', [-1] * 717) 295 | r0 = sc.sampled_points.get('r0', [-1] * 717) 296 | r1 = sc.sampled_points.get('r1', [-1] * 717) 297 | 298 | lanes = [l1, l0, r0, r1] 299 | return lanes 300 | 301 | 302 | def _filter_lanes_by_size(label, min_height=40): 303 | """ May need some tuning """ 304 | filtered_lanes = [] 305 | for lane in label['lanes']: 306 | lane_start = min([int(marker['pixel_start']['y']) for marker in lane['markers']]) 307 | lane_end = max([int(marker['pixel_start']['y']) for marker in lane['markers']]) 308 | if (lane_end - lane_start) < min_height: 309 | continue 310 | filtered_lanes.append(lane) 311 | label['lanes'] = filtered_lanes 312 | 313 | 314 | def _filter_few_markers(label, min_markers=2): 315 | """Filter lines that consist of only few markers""" 316 | filtered_lanes = [] 317 | for lane in label['lanes']: 318 | if len(lane['markers']) >= min_markers: 319 | filtered_lanes.append(lane) 320 | label['lanes'] = filtered_lanes 321 | 322 | 323 | def _fix_lane_names(label): 324 | """ Given keys ['l3', 'l2', 'l0', 'r0', 'r2'] returns ['l2', 'l1', 'l0', 'r0', 'r1']""" 325 | 326 | # Create mapping 327 | l_counter = 0 328 | r_counter = 0 329 | mapping = {} 330 | lane_ids = [lane['lane_id'] for lane in label['lanes']] 331 | for key in sorted(lane_ids): 332 | if key[0] == 'l': 333 | mapping[key] = 'l' + str(l_counter) 334 | l_counter += 1 335 | if key[0] == 'r': 336 | mapping[key] = 'r' + str(r_counter) 337 | r_counter += 1 338 | for lane in label['lanes']: 339 | lane['lane_id'] = mapping[lane['lane_id']] 340 | 341 | 342 | def read_json(json_path, min_lane_height=20): 343 | """ Reads and cleans label file information by path""" 344 | with open(json_path, 'r') as jf: 345 | label_content = json.load(jf) 346 | 347 | _filter_lanes_by_size(label_content, min_height=min_lane_height) 348 | _filter_few_markers(label_content, min_markers=2) 349 | _fix_lane_names(label_content) 350 | 351 | content = {'projection_matrix': label_content['projection_matrix'], 'lanes': label_content['lanes']} 352 | 353 | for lane in content['lanes']: 354 | for marker in lane['markers']: 355 | for pixel_key in marker['pixel_start'].keys(): 356 | marker['pixel_start'][pixel_key] = int(marker['pixel_start'][pixel_key]) 357 | for pixel_key in marker['pixel_end'].keys(): 358 | marker['pixel_end'][pixel_key] = int(marker['pixel_end'][pixel_key]) 359 | for pixel_key in marker['world_start'].keys(): 360 | marker['world_start'][pixel_key] = float(marker['world_start'][pixel_key]) 361 | for pixel_key in marker['world_end'].keys(): 362 | marker['world_end'][pixel_key] = float(marker['world_end'][pixel_key]) 363 | return content 364 | 365 | 366 | def ir(some_value): 367 | """ Rounds and casts to int 368 | Useful for pixel values that cannot be floats 369 | Parameters 370 | ---------- 371 | some_value : float 372 | numeric value 373 | Returns 374 | -------- 375 | Rounded integer 376 | Raises 377 | ------ 378 | ValueError for non scalar types 379 | """ 380 | return int(round(some_value)) 381 | 382 | 383 | # End code under the previous license 384 | --------------------------------------------------------------------------------