├── .gitattributes ├── .gitignore ├── LICENSE ├── NvidiaXavier.md ├── README.md ├── convertPytorchToONNX.py ├── doc └── img │ └── detected lanes.jpg ├── imageLaneDetection.py ├── input.jpg ├── models └── .gitkeep ├── output.jpg ├── requirements.txt ├── ultrafastLaneDetector ├── __init__.py ├── backbone.py ├── model.py └── ultrafastLaneDetector.py ├── videoLaneDetection.py └── webcamLaneDetection.py /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .nox/ 42 | .coverage 43 | .coverage.* 44 | .cache 45 | nosetests.xml 46 | coverage.xml 47 | *.cover 48 | .hypothesis/ 49 | .pytest_cache/ 50 | 51 | # Translations 52 | *.mo 53 | *.pot 54 | 55 | # Django stuff: 56 | *.log 57 | local_settings.py 58 | db.sqlite3 59 | 60 | # Flask stuff: 61 | instance/ 62 | .webassets-cache 63 | 64 | # Scrapy stuff: 65 | .scrapy 66 | 67 | # Sphinx documentation 68 | docs/_build/ 69 | 70 | # PyBuilder 71 | target/ 72 | 73 | # Jupyter Notebook 74 | .ipynb_checkpoints 75 | 76 | # IPython 77 | profile_default/ 78 | ipython_config.py 79 | 80 | # pyenv 81 | .python-version 82 | 83 | # celery beat schedule file 84 | celerybeat-schedule 85 | 86 | # SageMath parsed files 87 | *.sage.py 88 | 89 | # Environments 90 | .env 91 | .venv 92 | env/ 93 | venv/ 94 | ENV/ 95 | env.bak/ 96 | venv.bak/ 97 | 98 | # Spyder project settings 99 | .spyderproject 100 | .spyproject 101 | 102 | # Rope project settings 103 | .ropeproject 104 | 105 | # mkdocs documentation 106 | /site 107 | 108 | # mypy 109 | .mypy_cache/ 110 | .dmypy.json 111 | dmypy.json 112 | 113 | # Pyre type checker 114 | .pyre/ 115 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Ibai Gorordo 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /NvidiaXavier.md: -------------------------------------------------------------------------------- 1 | # onnx-Ultra-Fast-Lane-Detection-Inference With Nvidia Xavier 2 | Example scripts for the detection of lanes using the [ultra fast lane detection model](https://github.com/cfzd/Ultra-Fast-Lane-Detection) in ONNX. 3 | 4 | ![!Ultra fast lane detection](https://github.com/ibaiGorordo/onnx-Ultra-Fast-Lane-Detection-Inference/blob/main/doc/img/detected%20lanes.jpg) 5 | Source: https://www.flickr.com/photos/32413914@N00/1475776461/ 6 | 7 | # Pytorch inference 8 | For performing the inference in Pytorch, check my other repository **[Ultrafast Lane Detection Inference Pytorch](https://github.com/ibaiGorordo/Ultrafast-Lane-Detection-Inference-Pytorch-)**. 9 | 10 | # Requirements 11 | 12 | * **OpenCV**, **scipy**, **onnx** and **onnxruntime**. **pafy** and **youtube-dl**, **Nvidia Xavier AGX**, **JetPack 4.6** and **Python3.6** are required for youtube video inference. 13 | 14 | # Installation 15 | 16 | ``` 17 | pip3 install opencv-python 18 | pip3 install scipy 19 | Download >> https://nvidia.app.box.com/s/bfs688apyvor4eo8sf3y1oqtnarwafww 20 | Install >> pip3 install onnxruntime_gpu-1.8.0-cp36-cp36m-linux_aarch64.whl 21 | pip3 install scikit-build 22 | if you found not match version try to upgrade the PIP >> sudo -H pip3 install --upgrade pip 23 | 24 | ``` 25 | # ONNX model 26 | The original model was converted to different formats (including .onnx) by [PINTO0309](https://github.com/PINTO0309), download the models from [his repository](https://github.com/PINTO0309/PINTO_model_zoo/tree/main/140_Ultra-Fast-Lane-Detection) and save it into the **[models](https://github.com/ibaiGorordo/TfLite-Ultra-Fast-Lane-Detection-Inference/tree/main/models)** folder. 27 | 28 | ONNX Conversion script: https://github.com/cfzd/Ultra-Fast-Lane-Detection/issues/218 29 | 30 | # Original Pytorch model 31 | The pretrained Pytorch model was taken from the [original repository](https://github.com/cfzd/Ultra-Fast-Lane-Detection). 32 | 33 | # Model info ([link](https://github.com/cfzd/Ultra-Fast-Lane-Detection)) 34 | 35 | * **Input**: RGB image of size 800 x 200 pixels. 36 | * **Output**: Keypoints for a maximum of 4 lanes (left-most lane, left lane, right lane, and right-most lane). 37 | 38 | # Examples 39 | 40 | * **Image inference**: 41 | 42 | ``` 43 | python imageLaneDetection.py 44 | ``` 45 | * **Webcam inference**: 46 | 47 | ``` 48 | python webcamLaneDetection.py 49 | ``` 50 | 51 | * **Video inference**: 52 | 53 | ``` 54 | python videoLaneDetection.py 55 | ``` 56 | 57 | ## Result 58 | - YoutTube: https://www.youtube.com/watch?v=yBmcYDke7Wg 59 | 60 | Original video: https://youtu.be/2CIxM7x-Clc (by Yunfei Guo) 61 | 62 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # onnx-Ultra-Fast-Lane-Detection-Inference 2 | Example scripts for the detection of lanes using the [ultra fast lane detection model](https://github.com/cfzd/Ultra-Fast-Lane-Detection) in ONNX. 3 | 4 | ![!Ultra fast lane detection](https://github.com/ibaiGorordo/onnx-Ultra-Fast-Lane-Detection-Inference/blob/main/doc/img/detected%20lanes.jpg) 5 | Source: https://www.flickr.com/photos/32413914@N00/1475776461/ 6 | 7 | # Pytorch inference 8 | For performing the inference in Pytorch, check my other repository **[Ultrafast Lane Detection Inference Pytorch](https://github.com/ibaiGorordo/Ultrafast-Lane-Detection-Inference-Pytorch-)**. 9 | 10 | # Tested Environment 11 | ## Computer or Laptop 12 | - 13 | 14 | ## Single Board 15 | - Jetson Xavier AGX. JetPack 4.6 16 | # Requirements for Laptop 17 | 18 | * **OpenCV**, **scipy**, **onnx** and **onnxruntime**. **pafy** and **youtube-dl** are required for youtube video inference. 19 | 20 | # Requirements for Nvidia Xavier 21 | 22 | * **OpenCV**, **scipy**, **onnx** and **onnxruntime**. **pafy** and **youtube-dl**, **Nvidia Xavier AGX**, **JetPack 4.6** and **Python3.6** . 23 | 24 | # Installation for Laptop 25 | ``` 26 | pip install -r requirements.txt 27 | pip install pafy youtube-dl 28 | ``` 29 | # Installation for Nvidia Xavier 30 | ``` 31 | pip3 install opencv-python 32 | pip3 install scipy 33 | Download >> https://nvidia.app.box.com/s/bfs688apyvor4eo8sf3y1oqtnarwafww 34 | Install >> pip3 install onnxruntime_gpu-1.8.0-cp36-cp36m-linux_aarch64.whl 35 | pip3 install scikit-build 36 | if you found not match version try to upgrade the PIP >> sudo -H pip3 install --upgrade pip 37 | 38 | ``` 39 | 40 | # ONNX model 41 | The original model was converted to different formats (including .onnx) by [PINTO0309](https://github.com/PINTO0309), download the models from [his repository](https://github.com/PINTO0309/PINTO_model_zoo/tree/main/140_Ultra-Fast-Lane-Detection) and save it into the **[models](https://github.com/ibaiGorordo/TfLite-Ultra-Fast-Lane-Detection-Inference/tree/main/models)** folder. 42 | 43 | ONNX Conversion script: https://github.com/cfzd/Ultra-Fast-Lane-Detection/issues/218 44 | 45 | # Original Pytorch model 46 | The pretrained Pytorch model was taken from the [original repository](https://github.com/cfzd/Ultra-Fast-Lane-Detection). 47 | 48 | # Model info ([link](https://github.com/cfzd/Ultra-Fast-Lane-Detection)) 49 | 50 | * **Input**: RGB image of size 800 x 200 pixels. 51 | * **Output**: Keypoints for a maximum of 4 lanes (left-most lane, left lane, right lane, and right-most lane). 52 | 53 | # Examples 54 | 55 | * **Image inference**: 56 | 57 | ``` 58 | python imageLaneDetection.py 59 | ``` 60 | 61 | * **Webcam inference**: 62 | 63 | ``` 64 | python webcamLaneDetection.py 65 | ``` 66 | 67 | * **Video inference**: 68 | 69 | ``` 70 | python videoLaneDetection.py 71 | ``` 72 | 73 | # [Inference video Example](https://youtu.be/0Owf6gef1Ew) 74 | ![!Ultrafast lane detection on video](https://github.com/ibaiGorordo/Ultrafast-Lane-Detection-Inference-Pytorch-/blob/main/doc/img/laneDetection.gif) 75 | 76 | # Result for Nvidia Xavier 77 | - YoutTube: https://www.youtube.com/watch?v=yBmcYDke7Wg 78 | - Original video: https://youtu.be/2CIxM7x-Clc (by Yunfei Guo) 79 | 80 | -------------------------------------------------------------------------------- /convertPytorchToONNX.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import torch 3 | import scipy.special 4 | import numpy as np 5 | import torchvision 6 | from enum import Enum 7 | import onnx 8 | import onnxruntime 9 | 10 | from ultrafastLaneDetector.model import parsingNet 11 | 12 | class ModelType(Enum): 13 | TUSIMPLE = 0 14 | CULANE = 1 15 | 16 | class ModelConfig(): 17 | 18 | def __init__(self, model_type): 19 | 20 | if model_type == ModelType.TUSIMPLE: 21 | self.init_tusimple_config() 22 | else: 23 | self.init_culane_config() 24 | 25 | def init_tusimple_config(self): 26 | self.img_w = 1280 27 | self.img_h = 720 28 | self.griding_num = 100 29 | self.cls_num_per_lane = 56 30 | 31 | def init_culane_config(self): 32 | self.img_w = 1640 33 | self.img_h = 590 34 | self.griding_num = 200 35 | self.cls_num_per_lane = 18 36 | 37 | def convert_model(model_path, onnx_file_path, model_type=ModelType.TUSIMPLE): 38 | 39 | # Load model configuration based on the model type 40 | cfg = ModelConfig(model_type) 41 | 42 | 43 | # Load the model architecture 44 | net = parsingNet(pretrained = False, backbone='18', cls_dim = (cfg.griding_num+1,cfg.cls_num_per_lane,4), 45 | use_aux=False) # we dont need auxiliary segmentation in testing 46 | 47 | 48 | state_dict = torch.load(model_path, map_location='cpu')['model'] # CPU 49 | 50 | compatible_state_dict = {} 51 | for k, v in state_dict.items(): 52 | if 'module.' in k: 53 | compatible_state_dict[k[7:]] = v 54 | else: 55 | compatible_state_dict[k] = v 56 | 57 | # Load the weights into the model 58 | net.load_state_dict(compatible_state_dict, strict=False) 59 | 60 | img = torch.zeros(1, 3, 288, 800).to('cpu') 61 | torch.onnx.export(net, img, onnx_file_path, verbose=True) 62 | 63 | model = onnx.load(onnx_file_path) 64 | 65 | # Check that the IR is well formed 66 | onnx.checker.check_model(model) 67 | # Print a human readable representation of the graph 68 | print("==============================================================================================") 69 | 70 | if __name__ == '__main__': 71 | 72 | onnx_model_path = "tusimple_18.onnx" 73 | model_path = "models/tusimple_18.pth" 74 | model_type = ModelType.TUSIMPLE 75 | 76 | convert_model(model_path, onnx_model_path, model_type) 77 | 78 | 79 | 80 | -------------------------------------------------------------------------------- /doc/img/detected lanes.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ibaiGorordo/onnx-Ultra-Fast-Lane-Detection-Inference/3d61660eafe516179f1649e70a4cc6b83b800a8b/doc/img/detected lanes.jpg -------------------------------------------------------------------------------- /imageLaneDetection.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | 3 | from ultrafastLaneDetector import UltrafastLaneDetector, ModelType 4 | 5 | model_path = "models/lane.onnx" 6 | model_type = ModelType.TUSIMPLE 7 | 8 | image_path = "input.jpg" 9 | 10 | # Initialize lane detection model 11 | lane_detector = UltrafastLaneDetector(model_path, model_type) 12 | 13 | # Read RGB images 14 | img = cv2.imread(image_path, cv2.IMREAD_COLOR) 15 | 16 | # Detect the lanes 17 | output_img = lane_detector.detect_lanes(img) 18 | 19 | # Draw estimated depth 20 | cv2.namedWindow("Detected lanes", cv2.WINDOW_NORMAL) 21 | cv2.imshow("Detected lanes", output_img) 22 | cv2.waitKey(0) 23 | 24 | cv2.imwrite("output.jpg",output_img) 25 | -------------------------------------------------------------------------------- /input.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ibaiGorordo/onnx-Ultra-Fast-Lane-Detection-Inference/3d61660eafe516179f1649e70a4cc6b83b800a8b/input.jpg -------------------------------------------------------------------------------- /models/.gitkeep: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ibaiGorordo/onnx-Ultra-Fast-Lane-Detection-Inference/3d61660eafe516179f1649e70a4cc6b83b800a8b/models/.gitkeep -------------------------------------------------------------------------------- /output.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ibaiGorordo/onnx-Ultra-Fast-Lane-Detection-Inference/3d61660eafe516179f1649e70a4cc6b83b800a8b/output.jpg -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | opencv-python 2 | scipy 3 | onnx 4 | onnxruntime -------------------------------------------------------------------------------- /ultrafastLaneDetector/__init__.py: -------------------------------------------------------------------------------- 1 | from ultrafastLaneDetector.ultrafastLaneDetector import UltrafastLaneDetector, ModelType -------------------------------------------------------------------------------- /ultrafastLaneDetector/backbone.py: -------------------------------------------------------------------------------- 1 | import torch,pdb 2 | import torchvision 3 | import torch.nn.modules 4 | 5 | class vgg16bn(torch.nn.Module): 6 | def __init__(self,pretrained = False): 7 | super(vgg16bn,self).__init__() 8 | model = list(torchvision.models.vgg16_bn(pretrained=pretrained).features.children()) 9 | model = model[:33]+model[34:43] 10 | self.model = torch.nn.Sequential(*model) 11 | 12 | def forward(self,x): 13 | return self.model(x) 14 | class resnet(torch.nn.Module): 15 | def __init__(self,layers,pretrained = False): 16 | super(resnet,self).__init__() 17 | if layers == '18': 18 | model = torchvision.models.resnet18(pretrained=pretrained) 19 | elif layers == '34': 20 | model = torchvision.models.resnet34(pretrained=pretrained) 21 | elif layers == '50': 22 | model = torchvision.models.resnet50(pretrained=pretrained) 23 | elif layers == '101': 24 | model = torchvision.models.resnet101(pretrained=pretrained) 25 | elif layers == '152': 26 | model = torchvision.models.resnet152(pretrained=pretrained) 27 | elif layers == '50next': 28 | model = torchvision.models.resnext50_32x4d(pretrained=pretrained) 29 | elif layers == '101next': 30 | model = torchvision.models.resnext101_32x8d(pretrained=pretrained) 31 | elif layers == '50wide': 32 | model = torchvision.models.wide_resnet50_2(pretrained=pretrained) 33 | elif layers == '101wide': 34 | model = torchvision.models.wide_resnet101_2(pretrained=pretrained) 35 | else: 36 | raise NotImplementedError 37 | 38 | self.conv1 = model.conv1 39 | self.bn1 = model.bn1 40 | self.relu = model.relu 41 | self.maxpool = model.maxpool 42 | self.layer1 = model.layer1 43 | self.layer2 = model.layer2 44 | self.layer3 = model.layer3 45 | self.layer4 = model.layer4 46 | 47 | def forward(self,x): 48 | x = self.conv1(x) 49 | x = self.bn1(x) 50 | x = self.relu(x) 51 | x = self.maxpool(x) 52 | x = self.layer1(x) 53 | x2 = self.layer2(x) 54 | x3 = self.layer3(x2) 55 | x4 = self.layer4(x3) 56 | return x2,x3,x4 57 | -------------------------------------------------------------------------------- /ultrafastLaneDetector/model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from ultrafastLaneDetector.backbone import resnet 3 | import numpy as np 4 | 5 | class conv_bn_relu(torch.nn.Module): 6 | def __init__(self,in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1,bias=False): 7 | super(conv_bn_relu,self).__init__() 8 | self.conv = torch.nn.Conv2d(in_channels,out_channels, kernel_size, 9 | stride = stride, padding = padding, dilation = dilation,bias = bias) 10 | self.bn = torch.nn.BatchNorm2d(out_channels) 11 | self.relu = torch.nn.ReLU() 12 | 13 | def forward(self,x): 14 | x = self.conv(x) 15 | x = self.bn(x) 16 | x = self.relu(x) 17 | return x 18 | class parsingNet(torch.nn.Module): 19 | def __init__(self, size=(288, 800), pretrained=True, backbone='50', cls_dim=(37, 10, 4), use_aux=False): 20 | super(parsingNet, self).__init__() 21 | 22 | self.size = size 23 | self.w = size[0] 24 | self.h = size[1] 25 | self.cls_dim = cls_dim # (num_gridding, num_cls_per_lane, num_of_lanes) 26 | # num_cls_per_lane is the number of row anchors 27 | self.use_aux = use_aux 28 | self.total_dim = np.prod(cls_dim) 29 | 30 | # input : nchw, 31 | # output: (w+1) * sample_rows * 4 32 | self.model = resnet(backbone, pretrained=pretrained) 33 | 34 | if self.use_aux: 35 | self.aux_header2 = torch.nn.Sequential( 36 | conv_bn_relu(128, 128, kernel_size=3, stride=1, padding=1) if backbone in ['34','18'] else conv_bn_relu(512, 128, kernel_size=3, stride=1, padding=1), 37 | conv_bn_relu(128,128,3,padding=1), 38 | conv_bn_relu(128,128,3,padding=1), 39 | conv_bn_relu(128,128,3,padding=1), 40 | ) 41 | self.aux_header3 = torch.nn.Sequential( 42 | conv_bn_relu(256, 128, kernel_size=3, stride=1, padding=1) if backbone in ['34','18'] else conv_bn_relu(1024, 128, kernel_size=3, stride=1, padding=1), 43 | conv_bn_relu(128,128,3,padding=1), 44 | conv_bn_relu(128,128,3,padding=1), 45 | ) 46 | self.aux_header4 = torch.nn.Sequential( 47 | conv_bn_relu(512, 128, kernel_size=3, stride=1, padding=1) if backbone in ['34','18'] else conv_bn_relu(2048, 128, kernel_size=3, stride=1, padding=1), 48 | conv_bn_relu(128,128,3,padding=1), 49 | ) 50 | self.aux_combine = torch.nn.Sequential( 51 | conv_bn_relu(384, 256, 3,padding=2,dilation=2), 52 | conv_bn_relu(256, 128, 3,padding=2,dilation=2), 53 | conv_bn_relu(128, 128, 3,padding=2,dilation=2), 54 | conv_bn_relu(128, 128, 3,padding=4,dilation=4), 55 | torch.nn.Conv2d(128, cls_dim[-1] + 1,1) 56 | # output : n, num_of_lanes+1, h, w 57 | ) 58 | initialize_weights(self.aux_header2,self.aux_header3,self.aux_header4,self.aux_combine) 59 | 60 | self.cls = torch.nn.Sequential( 61 | torch.nn.Linear(1800, 2048), 62 | torch.nn.ReLU(), 63 | torch.nn.Linear(2048, self.total_dim), 64 | ) 65 | 66 | self.pool = torch.nn.Conv2d(512,8,1) if backbone in ['34','18'] else torch.nn.Conv2d(2048,8,1) 67 | # 1/32,2048 channel 68 | # 288,800 -> 9,40,2048 69 | # (w+1) * sample_rows * 4 70 | # 37 * 10 * 4 71 | #initialize_weights(self.cls) 72 | 73 | def forward(self, x): 74 | # n c h w - > n 2048 sh sw 75 | # -> n 2048 76 | x2,x3,fea = self.model(x) 77 | if self.use_aux: 78 | x2 = self.aux_header2(x2) 79 | x3 = self.aux_header3(x3) 80 | x3 = torch.nn.functional.interpolate(x3,scale_factor = 2,mode='bilinear') 81 | x4 = self.aux_header4(fea) 82 | x4 = torch.nn.functional.interpolate(x4,scale_factor = 4,mode='bilinear') 83 | aux_seg = torch.cat([x2,x3,x4],dim=1) 84 | aux_seg = self.aux_combine(aux_seg) 85 | else: 86 | aux_seg = None 87 | 88 | fea = self.pool(fea).view(-1, 1800) 89 | 90 | group_cls = self.cls(fea).view(-1, *self.cls_dim) 91 | 92 | if self.use_aux: 93 | return group_cls, aux_seg 94 | 95 | return group_cls 96 | 97 | 98 | def initialize_weights(*models): 99 | for model in models: 100 | real_init_weights(model) 101 | def real_init_weights(m): 102 | 103 | if isinstance(m, list): 104 | for mini_m in m: 105 | real_init_weights(mini_m) 106 | else: 107 | if isinstance(m, torch.nn.Conv2d): 108 | torch.nn.init.kaiming_normal_(m.weight, nonlinearity='relu') 109 | if m.bias is not None: 110 | torch.nn.init.constant_(m.bias, 0) 111 | elif isinstance(m, torch.nn.Linear): 112 | m.weight.data.normal_(0.0, std=0.01) 113 | elif isinstance(m, torch.nn.BatchNorm2d): 114 | torch.nn.init.constant_(m.weight, 1) 115 | torch.nn.init.constant_(m.bias, 0) 116 | elif isinstance(m,torch.nn.Module): 117 | for mini_m in m.children(): 118 | real_init_weights(mini_m) 119 | else: 120 | print('unkonwn module', m) -------------------------------------------------------------------------------- /ultrafastLaneDetector/ultrafastLaneDetector.py: -------------------------------------------------------------------------------- 1 | import onnx 2 | import onnxruntime 3 | import scipy.special 4 | from enum import Enum 5 | import cv2 6 | import time 7 | import numpy as np 8 | 9 | lane_colors = [(0,0,255),(0,255,0),(255,0,0),(0,255,255)] 10 | 11 | tusimple_row_anchor = [ 64, 68, 72, 76, 80, 84, 88, 92, 96, 100, 104, 108, 112, 12 | 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164, 13 | 168, 172, 176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216, 14 | 220, 224, 228, 232, 236, 240, 244, 248, 252, 256, 260, 264, 268, 15 | 272, 276, 280, 284] 16 | culane_row_anchor = [121, 131, 141, 150, 160, 170, 180, 189, 199, 209, 219, 228, 238, 248, 258, 267, 277, 287] 17 | 18 | class ModelType(Enum): 19 | TUSIMPLE = 0 20 | CULANE = 1 21 | 22 | class ModelConfig(): 23 | 24 | def __init__(self, model_type): 25 | 26 | if model_type == ModelType.TUSIMPLE: 27 | self.init_tusimple_config() 28 | else: 29 | self.init_culane_config() 30 | 31 | def init_tusimple_config(self): 32 | self.img_w = 1280 33 | self.img_h = 720 34 | self.row_anchor = tusimple_row_anchor 35 | self.griding_num = 100 36 | self.cls_num_per_lane = 56 37 | 38 | def init_culane_config(self): 39 | self.img_w = 1640 40 | self.img_h = 590 41 | self.row_anchor = culane_row_anchor 42 | self.griding_num = 200 43 | self.cls_num_per_lane = 18 44 | 45 | class UltrafastLaneDetector(): 46 | 47 | def __init__(self, model_path, model_type=ModelType.TUSIMPLE): 48 | 49 | self.fps = 0 50 | self.timeLastPrediction = time.time() 51 | self.frameCounter = 0 52 | 53 | # Load model configuration based on the model type 54 | self.cfg = ModelConfig(model_type) 55 | 56 | # Initialize model 57 | self.initialize_model(model_path) 58 | 59 | 60 | def initialize_model(self, model_path): 61 | 62 | self.session = onnxruntime.InferenceSession(model_path) 63 | 64 | # Get model info 65 | self.getModel_input_details() 66 | self.getModel_output_details() 67 | 68 | def detect_lanes(self, image, draw_points=True): 69 | 70 | input_tensor = self.prepare_input(image) 71 | 72 | # Perform inference on the image 73 | output = self.inference(input_tensor) 74 | 75 | # Process output data 76 | self.lanes_points, self.lanes_detected = self.process_output(output, self.cfg) 77 | 78 | # # Draw depth image 79 | visualization_img = self.draw_lanes(image, self.lanes_points, self.lanes_detected, self.cfg, draw_points) 80 | 81 | return visualization_img 82 | 83 | def prepare_input(self, image): 84 | img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 85 | self.img_height, self.img_width, self.img_channels = img.shape 86 | 87 | # Input values should be from -1 to 1 with a size of 288 x 800 pixels 88 | img_input = cv2.resize(img, (self.input_width,self.input_height)).astype(np.float32) 89 | 90 | # Scale input pixel values to -1 to 1 91 | mean=[0.485, 0.456, 0.406] 92 | std=[0.229, 0.224, 0.225] 93 | 94 | img_input = ((img_input/ 255.0 - mean) / std) 95 | img_input = img_input.transpose(2, 0, 1) 96 | img_input = img_input[np.newaxis,:,:,:] 97 | 98 | return img_input.astype(np.float32) 99 | 100 | def inference(self, input_tensor): 101 | input_name = self.session.get_inputs()[0].name 102 | output_name = self.session.get_outputs()[0].name 103 | 104 | output = self.session.run([output_name], {input_name: input_tensor}) 105 | 106 | return output 107 | 108 | 109 | def getModel_input_details(self): 110 | 111 | self.input_shape = self.session.get_inputs()[0].shape 112 | self.channes = self.input_shape[2] 113 | self.input_height = self.input_shape[2] 114 | self.input_width = self.input_shape[3] 115 | 116 | def getModel_output_details(self): 117 | 118 | self.output_shape = self.session.get_outputs()[0].shape 119 | self.num_points = self.output_shape[1] 120 | self.num_anchors = self.output_shape[2] 121 | self.num_lanes = self.output_shape[3] 122 | 123 | @staticmethod 124 | def process_output(output, cfg): 125 | # Parse the output of the model 126 | 127 | processed_output = np.squeeze(output[0]) 128 | print(processed_output.shape) 129 | print(np.min(processed_output), np.max(processed_output)) 130 | print(processed_output.reshape((1,-1))) 131 | processed_output = processed_output[:, ::-1, :] 132 | prob = scipy.special.softmax(processed_output[:-1, :, :], axis=0) 133 | idx = np.arange(cfg.griding_num) + 1 134 | idx = idx.reshape(-1, 1, 1) 135 | loc = np.sum(prob * idx, axis=0) 136 | processed_output = np.argmax(processed_output, axis=0) 137 | loc[processed_output == cfg.griding_num] = 0 138 | processed_output = loc 139 | 140 | 141 | col_sample = np.linspace(0, 800 - 1, cfg.griding_num) 142 | col_sample_w = col_sample[1] - col_sample[0] 143 | 144 | lanes_points = [] 145 | lanes_detected = [] 146 | 147 | max_lanes = processed_output.shape[1] 148 | for lane_num in range(max_lanes): 149 | lane_points = [] 150 | # Check if there are any points detected in the lane 151 | if np.sum(processed_output[:, lane_num] != 0) > 2: 152 | 153 | lanes_detected.append(True) 154 | 155 | # Process each of the points for each lane 156 | for point_num in range(processed_output.shape[0]): 157 | if processed_output[point_num, lane_num] > 0: 158 | lane_point = [int(processed_output[point_num, lane_num] * col_sample_w * cfg.img_w / 800) - 1, int(cfg.img_h * (cfg.row_anchor[cfg.cls_num_per_lane-1-point_num]/288)) - 1 ] 159 | lane_points.append(lane_point) 160 | else: 161 | lanes_detected.append(False) 162 | 163 | lanes_points.append(lane_points) 164 | return np.array(lanes_points), np.array(lanes_detected) 165 | 166 | @staticmethod 167 | def draw_lanes(input_img, lanes_points, lanes_detected, cfg, draw_points=True): 168 | # Write the detected line points in the image 169 | visualization_img = cv2.resize(input_img, (cfg.img_w, cfg.img_h), interpolation = cv2.INTER_AREA) 170 | 171 | # Draw a mask for the current lane 172 | if(lanes_detected[1] and lanes_detected[2]): 173 | lane_segment_img = visualization_img.copy() 174 | 175 | cv2.fillPoly(lane_segment_img, pts = [np.vstack((lanes_points[1],np.flipud(lanes_points[2])))], color =(255,191,0)) 176 | visualization_img = cv2.addWeighted(visualization_img, 0.7, lane_segment_img, 0.3, 0) 177 | 178 | if(draw_points): 179 | for lane_num,lane_points in enumerate(lanes_points): 180 | for lane_point in lane_points: 181 | cv2.circle(visualization_img, (lane_point[0],lane_point[1]), 3, lane_colors[lane_num], -1) 182 | 183 | return visualization_img 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 | 192 | 193 | -------------------------------------------------------------------------------- /videoLaneDetection.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import pafy 3 | from ultrafastLaneDetector import UltrafastLaneDetector, ModelType 4 | 5 | model_path = "models/lane.onnx" 6 | model_type = ModelType.TUSIMPLE 7 | 8 | # Initialize video 9 | # cap = cv2.VideoCapture("video.mp4") 10 | 11 | videoUrl = 'https://youtu.be/2CIxM7x-Clc' 12 | videoPafy = pafy.new(videoUrl) 13 | print(videoPafy.streams) 14 | cap = cv2.VideoCapture(videoPafy.streams[-1].url) 15 | 16 | # Initialize lane detection model 17 | lane_detector = UltrafastLaneDetector(model_path, model_type) 18 | 19 | cv2.namedWindow("Detected lanes", cv2.WINDOW_NORMAL) 20 | 21 | while cap.isOpened(): 22 | try: 23 | # Read frame from the video 24 | ret, frame = cap.read() 25 | except: 26 | continue 27 | 28 | if ret: 29 | 30 | # Detect the lanes 31 | output_img = lane_detector.detect_lanes(frame) 32 | 33 | cv2.imshow("Detected lanes", output_img) 34 | 35 | else: 36 | break 37 | 38 | # Press key q to stop 39 | if cv2.waitKey(1) == ord('q'): 40 | break 41 | 42 | cap.release() 43 | cv2.destroyAllWindows() -------------------------------------------------------------------------------- /webcamLaneDetection.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | 3 | from ultrafastLaneDetector import UltrafastLaneDetector, ModelType 4 | 5 | model_path = "models/lane.onnx" 6 | model_type = ModelType.TUSIMPLE 7 | 8 | # Initialize lane detection model 9 | lane_detector = UltrafastLaneDetector(model_path, model_type) 10 | 11 | # Initialize webcam 12 | cap = cv2.VideoCapture(0) 13 | cv2.namedWindow("Detected lanes", cv2.WINDOW_NORMAL) 14 | 15 | while(True): 16 | ret, frame = cap.read() 17 | 18 | # Detect the lanes 19 | output_img = lane_detector.detect_lanes(frame) 20 | 21 | 22 | cv2.imshow("Detected lanes", output_img) 23 | if cv2.waitKey(1) & 0xFF == ord('q'): 24 | break 25 | 26 | --------------------------------------------------------------------------------