├── LICENSE ├── README.md ├── app.py ├── buildSteps.txt ├── gen_wts.py ├── models ├── __init__.py ├── common.py ├── experimental.py ├── hub │ ├── anchors.yaml │ ├── yolov3-spp.yaml │ ├── yolov3-tiny.yaml │ ├── yolov3.yaml │ ├── yolov5-bifpn.yaml │ ├── yolov5-fpn.yaml │ ├── yolov5-p2.yaml │ ├── yolov5-p34.yaml │ ├── yolov5-p6.yaml │ ├── yolov5-p7.yaml │ ├── yolov5-panet.yaml │ ├── yolov5l6.yaml │ ├── yolov5m6.yaml │ ├── yolov5n6.yaml │ ├── yolov5s-LeakyReLU.yaml │ ├── yolov5s-ghost.yaml │ ├── yolov5s-transformer.yaml │ ├── yolov5s6.yaml │ └── yolov5x6.yaml ├── segment │ ├── yolov5l-seg.yaml │ ├── yolov5m-seg.yaml │ ├── yolov5n-seg.yaml │ ├── yolov5s-seg.yaml │ └── yolov5x-seg.yaml ├── tf.py ├── yolo.py ├── yolov5l.yaml ├── yolov5m.yaml ├── yolov5n.yaml ├── yolov5s.yaml └── yolov5x.yaml ├── setup.txt ├── utils ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── augmentations.cpython-36.pyc │ ├── autoanchor.cpython-36.pyc │ ├── dataloaders.cpython-36.pyc │ ├── downloads.cpython-36.pyc │ ├── general.cpython-36.pyc │ ├── metrics.cpython-36.pyc │ ├── plots.cpython-36.pyc │ └── torch_utils.cpython-36.pyc ├── activations.py ├── augmentations.py ├── autoanchor.py ├── autobatch.py ├── aws │ ├── __init__.py │ ├── mime.sh │ ├── resume.py │ └── userdata.sh ├── callbacks.py ├── dataloaders.py ├── docker │ ├── Dockerfile │ ├── Dockerfile-arm64 │ └── Dockerfile-cpu ├── downloads.py ├── flask_rest_api │ ├── README.md │ ├── example_request.py │ └── restapi.py ├── general.py ├── google_app_engine │ ├── Dockerfile │ ├── additional_requirements.txt │ └── app.yaml ├── loggers │ ├── __init__.py │ ├── clearml │ │ ├── README.md │ │ ├── __init__.py │ │ ├── clearml_utils.py │ │ └── hpo.py │ ├── comet │ │ ├── README.md │ │ ├── __init__.py │ │ ├── comet_utils.py │ │ ├── hpo.py │ │ └── optimizer_config.json │ └── wandb │ │ ├── __init__.py │ │ └── wandb_utils.py ├── loss.py ├── metrics.py ├── plots.py ├── segment │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ └── general.cpython-36.pyc │ ├── augmentations.py │ ├── dataloaders.py │ ├── general.py │ ├── loss.py │ ├── metrics.py │ └── plots.py ├── torch_utils.py └── triton.py ├── videos ├── out.jpg ├── testvideo.mp4 └── testvideo1.mp4 ├── yoloDet.py ├── yolov5 ├── CMakeLists.txt ├── images │ ├── bus.jpg │ └── zidane.jpg ├── plugin │ ├── yololayer.cu │ └── yololayer.h ├── src │ ├── calibrator.cpp │ ├── calibrator.h │ ├── config.h │ ├── cuda_utils.h │ ├── logging.h │ ├── macros.h │ ├── model.cpp │ ├── model.h │ ├── postprocess.cpp │ ├── postprocess.h │ ├── preprocess.cu │ ├── preprocess.h │ ├── types.h │ └── utils.h ├── yolov5_cls.cpp ├── yolov5_det.cpp └── yolov5_seg.cpp ├── yolov5n.pt └── yolov5s.pt /README.md: -------------------------------------------------------------------------------- 1 | Running YoloV5 with TensorRT Engine on Jetson. 2 | ========== 3 | 4 | This repository contains step by step guide to build and convert YoloV5 model into a TensorRT engine on Jetson. This has been tested on Jetson Nano or Jetson Xavier 5 | 6 | Please install Jetpack OS version 4.6 as mentioned by Nvidia and follow below steps. Please follow each steps exactly mentioned in the video links below : 7 | 8 | Build YoloV5 TensorRT Engine on Jetson Nano: https://www.youtube.com/watch?v=ErWC3nBuV6k 9 | 10 | Object Detection YoloV5 TensorRT Engine on Jetson Nano: https://www.youtube.com/watch?v=-Vu65N1NRWw 11 | 12 | Jetson Xavier: 13 | 14 | 15 | 16 | Install Libraries 17 | ============= 18 | Please install below libraries:: 19 | 20 | $ sudo apt-get update 21 | $ sudo apt-get install -y liblapack-dev libblas-dev gfortran libfreetype6-dev libopenblas-base libopenmpi-dev libjpeg-dev zlib1g-dev 22 | $ sudo apt-get install -y python3-pip 23 | 24 | 25 | Install below python packages 26 | ============= 27 | Numpy comes pre installed with Jetpack, so make sure you uninstall it first and then confirm if it's uninstalled or not. Then install below packages: 28 | 29 | $ numpy==1.19.0 30 | $ pandas==0.22.0 31 | $ Pillow==8.4.0 32 | $ PyYAML==3.12 33 | $ scipy==1.5.4 34 | $ psutil 35 | $ tqdm==4.64.1 36 | $ imutils 37 | 38 | Install PyCuda 39 | ============= 40 | We need to first export few paths 41 | 42 | $ export PATH=/usr/local/cuda-10.2/bin${PATH:+:${PATH}} 43 | $ export LD_LIBRARY_PATH=/usr/local/cuda-10.2/lib64:$LD_LIBRARY_PATH 44 | $ python3 -m pip install pycuda --user 45 | 46 | 47 | Install Seaborn 48 | ============= 49 | 50 | $ sudo apt install python3-seaborn 51 | 52 | Install torch & torchvision 53 | ============= 54 | 55 | $ wget https://nvidia.box.com/shared/static/fjtbno0vpo676a25cgvuqc1wty0fkkg6.whl -O torch-1.10.0-cp36-cp36m-linux_aarch64.whl 56 | $ pip3 install torch-1.10.0-cp36-cp36m-linux_aarch64.whl 57 | $ git clone --branch v0.11.1 https://github.com/pytorch/vision torchvision 58 | $ cd torchvision 59 | $ sudo python3 setup.py install 60 | 61 | ### Not required but good library 62 | sudo python3 -m pip install -U jetson-stats==3.1.4 63 | 64 | This marks the installation of all the required libraries. 65 | 66 | ------------------------------------------------------------------------------------------ 67 | 68 | Generate wts file from pt file 69 | ============= 70 | Yolov5s.pt and Yolov5n.pt are already provided in the repo. But if you want you can download any other version of the yolov5 model. Then run below command to convert .pt file into .wts file 71 | 72 | $ cd JetsonYoloV5 73 | $ python3 gen_wts.py -w yolov5s.pt -o yolov5s.wts 74 | 75 | Make 76 | ============= 77 | Create a build directory inside yolov5. Copy and paste generated wts file into build directory and run below commands. If using custom model, make sure to update kNumClas in yolov5/src/config.h 78 | 79 | $ cd yolov5/ 80 | $ mkdir build 81 | $ cd build 82 | $ cp ../../yolov5s.wts . 83 | $ cmake .. 84 | $ make 85 | 86 | Build Engine file 87 | ============= 88 | 89 | $ ./yolov5_det -s yolov5s.wts yolov5s.engine s 90 | 91 | 92 | Testing Engine file 93 | ============= 94 | 95 | $ ./yolov5_det -d yolov5s.engine ../images 96 | 97 | This will do inferencing over images and output will be saved in build directory. 98 | 99 | ----------------------------------------------------------------------------------------- 100 | 101 | Python Object Detection 102 | ============= 103 | Use `app.py` to do inferencing on any video file or camera. 104 | 105 | $ python3 app.py 106 | 107 | If you have custom model, make sure to update categories as per your classes in `yolovDet.py` . -------------------------------------------------------------------------------- /app.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import cv2 3 | import imutils 4 | from yoloDet import YoloTRT 5 | 6 | # use path for library and engine file 7 | model = YoloTRT(library="yolov5/build/libmyplugins.so", engine="yolov5/build/yolov5s.engine", conf=0.5, yolo_ver="v5") 8 | 9 | cap = cv2.VideoCapture("videos/testvideo.mp4") 10 | 11 | while True: 12 | ret, frame = cap.read() 13 | frame = imutils.resize(frame, width=600) 14 | detections, t = model.Inference(frame) 15 | # for obj in detections: 16 | # print(obj['class'], obj['conf'], obj['box']) 17 | # print("FPS: {} sec".format(1/t)) 18 | cv2.imshow("Output", frame) 19 | key = cv2.waitKey(1) 20 | if key == ord('q'): 21 | break 22 | cap.release() 23 | cv2.destroyAllWindows() -------------------------------------------------------------------------------- /buildSteps.txt: -------------------------------------------------------------------------------- 1 | # Generate WTS file 2 | python3 gen_wts.py -w yolov5s.pt -o yolov5s.wts 3 | 4 | # Cmake & Make 5 | # If using custom model, make sure to update kNumClas in yolov5/src/config.h 6 | cd yolov5/ 7 | mkdir build 8 | cd build 9 | cp ../../yolov5s.wts . 10 | cmake .. 11 | make 12 | 13 | # Build engine 14 | ./yolov5_det -s yolov5s.wts yolov5s.engine s 15 | 16 | # Test 17 | ./yolov5_det -d yolov5s.engine ../images 18 | -------------------------------------------------------------------------------- /gen_wts.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import argparse 3 | import os 4 | import struct 5 | import torch 6 | from utils.torch_utils import select_device 7 | 8 | 9 | def parse_args(): 10 | parser = argparse.ArgumentParser(description='Convert .pt file to .wts') 11 | parser.add_argument('-w', '--weights', required=True, 12 | help='Input weights (.pt) file path (required)') 13 | parser.add_argument( 14 | '-o', '--output', help='Output (.wts) file path (optional)') 15 | parser.add_argument( 16 | '-t', '--type', type=str, default='detect', choices=['detect', 'cls', 'seg'], 17 | help='determines the model is detection/classification') 18 | args = parser.parse_args() 19 | if not os.path.isfile(args.weights): 20 | raise SystemExit('Invalid input file') 21 | if not args.output: 22 | args.output = os.path.splitext(args.weights)[0] + '.wts' 23 | elif os.path.isdir(args.output): 24 | args.output = os.path.join( 25 | args.output, 26 | os.path.splitext(os.path.basename(args.weights))[0] + '.wts') 27 | return args.weights, args.output, args.type 28 | 29 | 30 | pt_file, wts_file, m_type = parse_args() 31 | print(f'Generating .wts for {m_type} model') 32 | 33 | # Load model 34 | print(f'Loading {pt_file}') 35 | device = select_device('cpu') 36 | model = torch.load(pt_file, map_location=device) # Load FP32 weights 37 | model = model['ema' if model.get('ema') else 'model'].float() 38 | 39 | if m_type in ['detect', 'seg']: 40 | # update anchor_grid info 41 | anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None] 42 | # model.model[-1].anchor_grid = anchor_grid 43 | delattr(model.model[-1], 'anchor_grid') # model.model[-1] is detect layer 44 | # The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight. 45 | model.model[-1].register_buffer("anchor_grid", anchor_grid) 46 | model.model[-1].register_buffer("strides", model.model[-1].stride) 47 | 48 | model.to(device).eval() 49 | 50 | print(f'Writing into {wts_file}') 51 | with open(wts_file, 'w') as f: 52 | f.write('{}\n'.format(len(model.state_dict().keys()))) 53 | for k, v in model.state_dict().items(): 54 | vr = v.reshape(-1).cpu().numpy() 55 | f.write('{} {} '.format(k, len(vr))) 56 | for vv in vr: 57 | f.write(' ') 58 | f.write(struct.pack('>f', float(vv)).hex()) 59 | f.write('\n') 60 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/models/__init__.py -------------------------------------------------------------------------------- /models/experimental.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Experimental modules 4 | """ 5 | import math 6 | 7 | import numpy as np 8 | import torch 9 | import torch.nn as nn 10 | 11 | from utils.downloads import attempt_download 12 | 13 | 14 | class Sum(nn.Module): 15 | # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 16 | def __init__(self, n, weight=False): # n: number of inputs 17 | super().__init__() 18 | self.weight = weight # apply weights boolean 19 | self.iter = range(n - 1) # iter object 20 | if weight: 21 | self.w = nn.Parameter(-torch.arange(1.0, n) / 2, requires_grad=True) # layer weights 22 | 23 | def forward(self, x): 24 | y = x[0] # no weight 25 | if self.weight: 26 | w = torch.sigmoid(self.w) * 2 27 | for i in self.iter: 28 | y = y + x[i + 1] * w[i] 29 | else: 30 | for i in self.iter: 31 | y = y + x[i + 1] 32 | return y 33 | 34 | 35 | class MixConv2d(nn.Module): 36 | # Mixed Depth-wise Conv https://arxiv.org/abs/1907.09595 37 | def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): # ch_in, ch_out, kernel, stride, ch_strategy 38 | super().__init__() 39 | n = len(k) # number of convolutions 40 | if equal_ch: # equal c_ per group 41 | i = torch.linspace(0, n - 1E-6, c2).floor() # c2 indices 42 | c_ = [(i == g).sum() for g in range(n)] # intermediate channels 43 | else: # equal weight.numel() per group 44 | b = [c2] + [0] * n 45 | a = np.eye(n + 1, n, k=-1) 46 | a -= np.roll(a, 1, axis=1) 47 | a *= np.array(k) ** 2 48 | a[0] = 1 49 | c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b 50 | 51 | self.m = nn.ModuleList([ 52 | nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)]) 53 | self.bn = nn.BatchNorm2d(c2) 54 | self.act = nn.SiLU() 55 | 56 | def forward(self, x): 57 | return self.act(self.bn(torch.cat([m(x) for m in self.m], 1))) 58 | 59 | 60 | class Ensemble(nn.ModuleList): 61 | # Ensemble of models 62 | def __init__(self): 63 | super().__init__() 64 | 65 | def forward(self, x, augment=False, profile=False, visualize=False): 66 | y = [module(x, augment, profile, visualize)[0] for module in self] 67 | # y = torch.stack(y).max(0)[0] # max ensemble 68 | # y = torch.stack(y).mean(0) # mean ensemble 69 | y = torch.cat(y, 1) # nms ensemble 70 | return y, None # inference, train output 71 | 72 | 73 | def attempt_load(weights, device=None, inplace=True, fuse=True): 74 | # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a 75 | from models.yolo import Detect, Model 76 | 77 | model = Ensemble() 78 | for w in weights if isinstance(weights, list) else [weights]: 79 | ckpt = torch.load(attempt_download(w), map_location='cpu') # load 80 | ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float() # FP32 model 81 | 82 | # Model compatibility updates 83 | if not hasattr(ckpt, 'stride'): 84 | ckpt.stride = torch.tensor([32.]) 85 | if hasattr(ckpt, 'names') and isinstance(ckpt.names, (list, tuple)): 86 | ckpt.names = dict(enumerate(ckpt.names)) # convert to dict 87 | 88 | model.append(ckpt.fuse().eval() if fuse and hasattr(ckpt, 'fuse') else ckpt.eval()) # model in eval mode 89 | 90 | # Module compatibility updates 91 | for m in model.modules(): 92 | t = type(m) 93 | if t in (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Model): 94 | m.inplace = inplace # torch 1.7.0 compatibility 95 | if t is Detect and not isinstance(m.anchor_grid, list): 96 | delattr(m, 'anchor_grid') 97 | setattr(m, 'anchor_grid', [torch.zeros(1)] * m.nl) 98 | elif t is nn.Upsample and not hasattr(m, 'recompute_scale_factor'): 99 | m.recompute_scale_factor = None # torch 1.11.0 compatibility 100 | 101 | # Return model 102 | if len(model) == 1: 103 | return model[-1] 104 | 105 | # Return detection ensemble 106 | print(f'Ensemble created with {weights}\n') 107 | for k in 'names', 'nc', 'yaml': 108 | setattr(model, k, getattr(model[0], k)) 109 | model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride # max stride 110 | assert all(model[0].nc == m.nc for m in model), f'Models have different class counts: {[m.nc for m in model]}' 111 | return model 112 | -------------------------------------------------------------------------------- /models/hub/anchors.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Default anchors for COCO data 3 | 4 | 5 | # P5 ------------------------------------------------------------------------------------------------------------------- 6 | # P5-640: 7 | anchors_p5_640: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | 13 | # P6 ------------------------------------------------------------------------------------------------------------------- 14 | # P6-640: thr=0.25: 0.9964 BPR, 5.54 anchors past thr, n=12, img_size=640, metric_all=0.281/0.716-mean/best, past_thr=0.469-mean: 9,11, 21,19, 17,41, 43,32, 39,70, 86,64, 65,131, 134,130, 120,265, 282,180, 247,354, 512,387 15 | anchors_p6_640: 16 | - [9,11, 21,19, 17,41] # P3/8 17 | - [43,32, 39,70, 86,64] # P4/16 18 | - [65,131, 134,130, 120,265] # P5/32 19 | - [282,180, 247,354, 512,387] # P6/64 20 | 21 | # P6-1280: thr=0.25: 0.9950 BPR, 5.55 anchors past thr, n=12, img_size=1280, metric_all=0.281/0.714-mean/best, past_thr=0.468-mean: 19,27, 44,40, 38,94, 96,68, 86,152, 180,137, 140,301, 303,264, 238,542, 436,615, 739,380, 925,792 22 | anchors_p6_1280: 23 | - [19,27, 44,40, 38,94] # P3/8 24 | - [96,68, 86,152, 180,137] # P4/16 25 | - [140,301, 303,264, 238,542] # P5/32 26 | - [436,615, 739,380, 925,792] # P6/64 27 | 28 | # P6-1920: thr=0.25: 0.9950 BPR, 5.55 anchors past thr, n=12, img_size=1920, metric_all=0.281/0.714-mean/best, past_thr=0.468-mean: 28,41, 67,59, 57,141, 144,103, 129,227, 270,205, 209,452, 455,396, 358,812, 653,922, 1109,570, 1387,1187 29 | anchors_p6_1920: 30 | - [28,41, 67,59, 57,141] # P3/8 31 | - [144,103, 129,227, 270,205] # P4/16 32 | - [209,452, 455,396, 358,812] # P5/32 33 | - [653,922, 1109,570, 1387,1187] # P6/64 34 | 35 | 36 | # P7 ------------------------------------------------------------------------------------------------------------------- 37 | # P7-640: thr=0.25: 0.9962 BPR, 6.76 anchors past thr, n=15, img_size=640, metric_all=0.275/0.733-mean/best, past_thr=0.466-mean: 11,11, 13,30, 29,20, 30,46, 61,38, 39,92, 78,80, 146,66, 79,163, 149,150, 321,143, 157,303, 257,402, 359,290, 524,372 38 | anchors_p7_640: 39 | - [11,11, 13,30, 29,20] # P3/8 40 | - [30,46, 61,38, 39,92] # P4/16 41 | - [78,80, 146,66, 79,163] # P5/32 42 | - [149,150, 321,143, 157,303] # P6/64 43 | - [257,402, 359,290, 524,372] # P7/128 44 | 45 | # P7-1280: thr=0.25: 0.9968 BPR, 6.71 anchors past thr, n=15, img_size=1280, metric_all=0.273/0.732-mean/best, past_thr=0.463-mean: 19,22, 54,36, 32,77, 70,83, 138,71, 75,173, 165,159, 148,334, 375,151, 334,317, 251,626, 499,474, 750,326, 534,814, 1079,818 46 | anchors_p7_1280: 47 | - [19,22, 54,36, 32,77] # P3/8 48 | - [70,83, 138,71, 75,173] # P4/16 49 | - [165,159, 148,334, 375,151] # P5/32 50 | - [334,317, 251,626, 499,474] # P6/64 51 | - [750,326, 534,814, 1079,818] # P7/128 52 | 53 | # P7-1920: thr=0.25: 0.9968 BPR, 6.71 anchors past thr, n=15, img_size=1920, metric_all=0.273/0.732-mean/best, past_thr=0.463-mean: 29,34, 81,55, 47,115, 105,124, 207,107, 113,259, 247,238, 222,500, 563,227, 501,476, 376,939, 749,711, 1126,489, 801,1222, 1618,1227 54 | anchors_p7_1920: 55 | - [29,34, 81,55, 47,115] # P3/8 56 | - [105,124, 207,107, 113,259] # P4/16 57 | - [247,238, 222,500, 563,227] # P5/32 58 | - [501,476, 376,939, 749,711] # P6/64 59 | - [1126,489, 801,1222, 1618,1227] # P7/128 60 | -------------------------------------------------------------------------------- /models/hub/yolov3-spp.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # darknet53 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [32, 3, 1]], # 0 16 | [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 17 | [-1, 1, Bottleneck, [64]], 18 | [-1, 1, Conv, [128, 3, 2]], # 3-P2/4 19 | [-1, 2, Bottleneck, [128]], 20 | [-1, 1, Conv, [256, 3, 2]], # 5-P3/8 21 | [-1, 8, Bottleneck, [256]], 22 | [-1, 1, Conv, [512, 3, 2]], # 7-P4/16 23 | [-1, 8, Bottleneck, [512]], 24 | [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32 25 | [-1, 4, Bottleneck, [1024]], # 10 26 | ] 27 | 28 | # YOLOv3-SPP head 29 | head: 30 | [[-1, 1, Bottleneck, [1024, False]], 31 | [-1, 1, SPP, [512, [5, 9, 13]]], 32 | [-1, 1, Conv, [1024, 3, 1]], 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large) 35 | 36 | [-2, 1, Conv, [256, 1, 1]], 37 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 38 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 39 | [-1, 1, Bottleneck, [512, False]], 40 | [-1, 1, Bottleneck, [512, False]], 41 | [-1, 1, Conv, [256, 1, 1]], 42 | [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium) 43 | 44 | [-2, 1, Conv, [128, 1, 1]], 45 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 46 | [[-1, 6], 1, Concat, [1]], # cat backbone P3 47 | [-1, 1, Bottleneck, [256, False]], 48 | [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small) 49 | 50 | [[27, 22, 15], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 51 | ] 52 | -------------------------------------------------------------------------------- /models/hub/yolov3-tiny.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,14, 23,27, 37,58] # P4/16 9 | - [81,82, 135,169, 344,319] # P5/32 10 | 11 | # YOLOv3-tiny backbone 12 | backbone: 13 | # [from, number, module, args] 14 | [[-1, 1, Conv, [16, 3, 1]], # 0 15 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 1-P1/2 16 | [-1, 1, Conv, [32, 3, 1]], 17 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 3-P2/4 18 | [-1, 1, Conv, [64, 3, 1]], 19 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 5-P3/8 20 | [-1, 1, Conv, [128, 3, 1]], 21 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 7-P4/16 22 | [-1, 1, Conv, [256, 3, 1]], 23 | [-1, 1, nn.MaxPool2d, [2, 2, 0]], # 9-P5/32 24 | [-1, 1, Conv, [512, 3, 1]], 25 | [-1, 1, nn.ZeroPad2d, [[0, 1, 0, 1]]], # 11 26 | [-1, 1, nn.MaxPool2d, [2, 1, 0]], # 12 27 | ] 28 | 29 | # YOLOv3-tiny head 30 | head: 31 | [[-1, 1, Conv, [1024, 3, 1]], 32 | [-1, 1, Conv, [256, 1, 1]], 33 | [-1, 1, Conv, [512, 3, 1]], # 15 (P5/32-large) 34 | 35 | [-2, 1, Conv, [128, 1, 1]], 36 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 37 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 38 | [-1, 1, Conv, [256, 3, 1]], # 19 (P4/16-medium) 39 | 40 | [[19, 15], 1, Detect, [nc, anchors]], # Detect(P4, P5) 41 | ] 42 | -------------------------------------------------------------------------------- /models/hub/yolov3.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # darknet53 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [32, 3, 1]], # 0 16 | [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 17 | [-1, 1, Bottleneck, [64]], 18 | [-1, 1, Conv, [128, 3, 2]], # 3-P2/4 19 | [-1, 2, Bottleneck, [128]], 20 | [-1, 1, Conv, [256, 3, 2]], # 5-P3/8 21 | [-1, 8, Bottleneck, [256]], 22 | [-1, 1, Conv, [512, 3, 2]], # 7-P4/16 23 | [-1, 8, Bottleneck, [512]], 24 | [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32 25 | [-1, 4, Bottleneck, [1024]], # 10 26 | ] 27 | 28 | # YOLOv3 head 29 | head: 30 | [[-1, 1, Bottleneck, [1024, False]], 31 | [-1, 1, Conv, [512, 1, 1]], 32 | [-1, 1, Conv, [1024, 3, 1]], 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large) 35 | 36 | [-2, 1, Conv, [256, 1, 1]], 37 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 38 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 39 | [-1, 1, Bottleneck, [512, False]], 40 | [-1, 1, Bottleneck, [512, False]], 41 | [-1, 1, Conv, [256, 1, 1]], 42 | [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium) 43 | 44 | [-2, 1, Conv, [128, 1, 1]], 45 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 46 | [[-1, 6], 1, Concat, [1]], # cat backbone P3 47 | [-1, 1, Bottleneck, [256, False]], 48 | [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small) 49 | 50 | [[27, 22, 15], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 51 | ] 52 | -------------------------------------------------------------------------------- /models/hub/yolov5-bifpn.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 BiFPN head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14, 6], 1, Concat, [1]], # cat P4 <--- BiFPN change 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/hub/yolov5-fpn.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 FPN head 28 | head: 29 | [[-1, 3, C3, [1024, False]], # 10 (P5/32-large) 30 | 31 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 32 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 3, C3, [512, False]], # 14 (P4/16-medium) 35 | 36 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 37 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 38 | [-1, 1, Conv, [256, 1, 1]], 39 | [-1, 3, C3, [256, False]], # 18 (P3/8-small) 40 | 41 | [[18, 14, 10], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 42 | ] 43 | -------------------------------------------------------------------------------- /models/hub/yolov5-p2.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 3 # AutoAnchor evolves 3 anchors per P output layer 8 | 9 | # YOLOv5 v6.0 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 13 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 14 | [-1, 3, C3, [128]], 15 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 16 | [-1, 6, C3, [256]], 17 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 18 | [-1, 9, C3, [512]], 19 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 20 | [-1, 3, C3, [1024]], 21 | [-1, 1, SPPF, [1024, 5]], # 9 22 | ] 23 | 24 | # YOLOv5 v6.0 head with (P2, P3, P4, P5) outputs 25 | head: 26 | [[-1, 1, Conv, [512, 1, 1]], 27 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 28 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 29 | [-1, 3, C3, [512, False]], # 13 30 | 31 | [-1, 1, Conv, [256, 1, 1]], 32 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 33 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 34 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 35 | 36 | [-1, 1, Conv, [128, 1, 1]], 37 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 38 | [[-1, 2], 1, Concat, [1]], # cat backbone P2 39 | [-1, 1, C3, [128, False]], # 21 (P2/4-xsmall) 40 | 41 | [-1, 1, Conv, [128, 3, 2]], 42 | [[-1, 18], 1, Concat, [1]], # cat head P3 43 | [-1, 3, C3, [256, False]], # 24 (P3/8-small) 44 | 45 | [-1, 1, Conv, [256, 3, 2]], 46 | [[-1, 14], 1, Concat, [1]], # cat head P4 47 | [-1, 3, C3, [512, False]], # 27 (P4/16-medium) 48 | 49 | [-1, 1, Conv, [512, 3, 2]], 50 | [[-1, 10], 1, Concat, [1]], # cat head P5 51 | [-1, 3, C3, [1024, False]], # 30 (P5/32-large) 52 | 53 | [[21, 24, 27, 30], 1, Detect, [nc, anchors]], # Detect(P2, P3, P4, P5) 54 | ] 55 | -------------------------------------------------------------------------------- /models/hub/yolov5-p34.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 7 | anchors: 3 # AutoAnchor evolves 3 anchors per P output layer 8 | 9 | # YOLOv5 v6.0 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [ [ -1, 1, Conv, [ 64, 6, 2, 2 ] ], # 0-P1/2 13 | [ -1, 1, Conv, [ 128, 3, 2 ] ], # 1-P2/4 14 | [ -1, 3, C3, [ 128 ] ], 15 | [ -1, 1, Conv, [ 256, 3, 2 ] ], # 3-P3/8 16 | [ -1, 6, C3, [ 256 ] ], 17 | [ -1, 1, Conv, [ 512, 3, 2 ] ], # 5-P4/16 18 | [ -1, 9, C3, [ 512 ] ], 19 | [ -1, 1, Conv, [ 1024, 3, 2 ] ], # 7-P5/32 20 | [ -1, 3, C3, [ 1024 ] ], 21 | [ -1, 1, SPPF, [ 1024, 5 ] ], # 9 22 | ] 23 | 24 | # YOLOv5 v6.0 head with (P3, P4) outputs 25 | head: 26 | [ [ -1, 1, Conv, [ 512, 1, 1 ] ], 27 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 28 | [ [ -1, 6 ], 1, Concat, [ 1 ] ], # cat backbone P4 29 | [ -1, 3, C3, [ 512, False ] ], # 13 30 | 31 | [ -1, 1, Conv, [ 256, 1, 1 ] ], 32 | [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], 33 | [ [ -1, 4 ], 1, Concat, [ 1 ] ], # cat backbone P3 34 | [ -1, 3, C3, [ 256, False ] ], # 17 (P3/8-small) 35 | 36 | [ -1, 1, Conv, [ 256, 3, 2 ] ], 37 | [ [ -1, 14 ], 1, Concat, [ 1 ] ], # cat head P4 38 | [ -1, 3, C3, [ 512, False ] ], # 20 (P4/16-medium) 39 | 40 | [ [ 17, 20 ], 1, Detect, [ nc, anchors ] ], # Detect(P3, P4) 41 | ] 42 | -------------------------------------------------------------------------------- /models/hub/yolov5-p6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 3 # AutoAnchor evolves 3 anchors per P output layer 8 | 9 | # YOLOv5 v6.0 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 13 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 14 | [-1, 3, C3, [128]], 15 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 16 | [-1, 6, C3, [256]], 17 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 18 | [-1, 9, C3, [512]], 19 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 20 | [-1, 3, C3, [768]], 21 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 22 | [-1, 3, C3, [1024]], 23 | [-1, 1, SPPF, [1024, 5]], # 11 24 | ] 25 | 26 | # YOLOv5 v6.0 head with (P3, P4, P5, P6) outputs 27 | head: 28 | [[-1, 1, Conv, [768, 1, 1]], 29 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 30 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 31 | [-1, 3, C3, [768, False]], # 15 32 | 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 35 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 36 | [-1, 3, C3, [512, False]], # 19 37 | 38 | [-1, 1, Conv, [256, 1, 1]], 39 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 40 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 41 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 42 | 43 | [-1, 1, Conv, [256, 3, 2]], 44 | [[-1, 20], 1, Concat, [1]], # cat head P4 45 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 46 | 47 | [-1, 1, Conv, [512, 3, 2]], 48 | [[-1, 16], 1, Concat, [1]], # cat head P5 49 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 50 | 51 | [-1, 1, Conv, [768, 3, 2]], 52 | [[-1, 12], 1, Concat, [1]], # cat head P6 53 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 54 | 55 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 56 | ] 57 | -------------------------------------------------------------------------------- /models/hub/yolov5-p7.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 3 # AutoAnchor evolves 3 anchors per P output layer 8 | 9 | # YOLOv5 v6.0 backbone 10 | backbone: 11 | # [from, number, module, args] 12 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 13 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 14 | [-1, 3, C3, [128]], 15 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 16 | [-1, 6, C3, [256]], 17 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 18 | [-1, 9, C3, [512]], 19 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 20 | [-1, 3, C3, [768]], 21 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 22 | [-1, 3, C3, [1024]], 23 | [-1, 1, Conv, [1280, 3, 2]], # 11-P7/128 24 | [-1, 3, C3, [1280]], 25 | [-1, 1, SPPF, [1280, 5]], # 13 26 | ] 27 | 28 | # YOLOv5 v6.0 head with (P3, P4, P5, P6, P7) outputs 29 | head: 30 | [[-1, 1, Conv, [1024, 1, 1]], 31 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 32 | [[-1, 10], 1, Concat, [1]], # cat backbone P6 33 | [-1, 3, C3, [1024, False]], # 17 34 | 35 | [-1, 1, Conv, [768, 1, 1]], 36 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 37 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 38 | [-1, 3, C3, [768, False]], # 21 39 | 40 | [-1, 1, Conv, [512, 1, 1]], 41 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 42 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 43 | [-1, 3, C3, [512, False]], # 25 44 | 45 | [-1, 1, Conv, [256, 1, 1]], 46 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 47 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 48 | [-1, 3, C3, [256, False]], # 29 (P3/8-small) 49 | 50 | [-1, 1, Conv, [256, 3, 2]], 51 | [[-1, 26], 1, Concat, [1]], # cat head P4 52 | [-1, 3, C3, [512, False]], # 32 (P4/16-medium) 53 | 54 | [-1, 1, Conv, [512, 3, 2]], 55 | [[-1, 22], 1, Concat, [1]], # cat head P5 56 | [-1, 3, C3, [768, False]], # 35 (P5/32-large) 57 | 58 | [-1, 1, Conv, [768, 3, 2]], 59 | [[-1, 18], 1, Concat, [1]], # cat head P6 60 | [-1, 3, C3, [1024, False]], # 38 (P6/64-xlarge) 61 | 62 | [-1, 1, Conv, [1024, 3, 2]], 63 | [[-1, 14], 1, Concat, [1]], # cat head P7 64 | [-1, 3, C3, [1280, False]], # 41 (P7/128-xxlarge) 65 | 66 | [[29, 32, 35, 38, 41], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6, P7) 67 | ] 68 | -------------------------------------------------------------------------------- /models/hub/yolov5-panet.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 PANet head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/hub/yolov5l6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [19,27, 44,40, 38,94] # P3/8 9 | - [96,68, 86,152, 180,137] # P4/16 10 | - [140,301, 303,264, 238,542] # P5/32 11 | - [436,615, 739,380, 925,792] # P6/64 12 | 13 | # YOLOv5 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 17 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 18 | [-1, 3, C3, [128]], 19 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 20 | [-1, 6, C3, [256]], 21 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 22 | [-1, 9, C3, [512]], 23 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 24 | [-1, 3, C3, [768]], 25 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 26 | [-1, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 head 31 | head: 32 | [[-1, 1, Conv, [768, 1, 1]], 33 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 34 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 35 | [-1, 3, C3, [768, False]], # 15 36 | 37 | [-1, 1, Conv, [512, 1, 1]], 38 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 39 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 40 | [-1, 3, C3, [512, False]], # 19 41 | 42 | [-1, 1, Conv, [256, 1, 1]], 43 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 44 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 45 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 46 | 47 | [-1, 1, Conv, [256, 3, 2]], 48 | [[-1, 20], 1, Concat, [1]], # cat head P4 49 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 50 | 51 | [-1, 1, Conv, [512, 3, 2]], 52 | [[-1, 16], 1, Concat, [1]], # cat head P5 53 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 54 | 55 | [-1, 1, Conv, [768, 3, 2]], 56 | [[-1, 12], 1, Concat, [1]], # cat head P6 57 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 58 | 59 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /models/hub/yolov5m6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.67 # model depth multiple 6 | width_multiple: 0.75 # layer channel multiple 7 | anchors: 8 | - [19,27, 44,40, 38,94] # P3/8 9 | - [96,68, 86,152, 180,137] # P4/16 10 | - [140,301, 303,264, 238,542] # P5/32 11 | - [436,615, 739,380, 925,792] # P6/64 12 | 13 | # YOLOv5 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 17 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 18 | [-1, 3, C3, [128]], 19 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 20 | [-1, 6, C3, [256]], 21 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 22 | [-1, 9, C3, [512]], 23 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 24 | [-1, 3, C3, [768]], 25 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 26 | [-1, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 head 31 | head: 32 | [[-1, 1, Conv, [768, 1, 1]], 33 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 34 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 35 | [-1, 3, C3, [768, False]], # 15 36 | 37 | [-1, 1, Conv, [512, 1, 1]], 38 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 39 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 40 | [-1, 3, C3, [512, False]], # 19 41 | 42 | [-1, 1, Conv, [256, 1, 1]], 43 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 44 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 45 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 46 | 47 | [-1, 1, Conv, [256, 3, 2]], 48 | [[-1, 20], 1, Concat, [1]], # cat head P4 49 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 50 | 51 | [-1, 1, Conv, [512, 3, 2]], 52 | [[-1, 16], 1, Concat, [1]], # cat head P5 53 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 54 | 55 | [-1, 1, Conv, [768, 3, 2]], 56 | [[-1, 12], 1, Concat, [1]], # cat head P6 57 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 58 | 59 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /models/hub/yolov5n6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.25 # layer channel multiple 7 | anchors: 8 | - [19,27, 44,40, 38,94] # P3/8 9 | - [96,68, 86,152, 180,137] # P4/16 10 | - [140,301, 303,264, 238,542] # P5/32 11 | - [436,615, 739,380, 925,792] # P6/64 12 | 13 | # YOLOv5 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 17 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 18 | [-1, 3, C3, [128]], 19 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 20 | [-1, 6, C3, [256]], 21 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 22 | [-1, 9, C3, [512]], 23 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 24 | [-1, 3, C3, [768]], 25 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 26 | [-1, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 head 31 | head: 32 | [[-1, 1, Conv, [768, 1, 1]], 33 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 34 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 35 | [-1, 3, C3, [768, False]], # 15 36 | 37 | [-1, 1, Conv, [512, 1, 1]], 38 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 39 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 40 | [-1, 3, C3, [512, False]], # 19 41 | 42 | [-1, 1, Conv, [256, 1, 1]], 43 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 44 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 45 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 46 | 47 | [-1, 1, Conv, [256, 3, 2]], 48 | [[-1, 20], 1, Concat, [1]], # cat head P4 49 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 50 | 51 | [-1, 1, Conv, [512, 3, 2]], 52 | [[-1, 16], 1, Concat, [1]], # cat head P5 53 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 54 | 55 | [-1, 1, Conv, [768, 3, 2]], 56 | [[-1, 12], 1, Concat, [1]], # cat head P6 57 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 58 | 59 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /models/hub/yolov5s-LeakyReLU.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | activation: nn.LeakyReLU(0.1) # <----- Conv() activation used throughout entire YOLOv5 model 6 | depth_multiple: 0.33 # model depth multiple 7 | width_multiple: 0.50 # layer channel multiple 8 | anchors: 9 | - [10,13, 16,30, 33,23] # P3/8 10 | - [30,61, 62,45, 59,119] # P4/16 11 | - [116,90, 156,198, 373,326] # P5/32 12 | 13 | # YOLOv5 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 17 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 18 | [-1, 3, C3, [128]], 19 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 20 | [-1, 6, C3, [256]], 21 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 22 | [-1, 9, C3, [512]], 23 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 24 | [-1, 3, C3, [1024]], 25 | [-1, 1, SPPF, [1024, 5]], # 9 26 | ] 27 | 28 | # YOLOv5 v6.0 head 29 | head: 30 | [[-1, 1, Conv, [512, 1, 1]], 31 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 32 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 33 | [-1, 3, C3, [512, False]], # 13 34 | 35 | [-1, 1, Conv, [256, 1, 1]], 36 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 37 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 38 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 39 | 40 | [-1, 1, Conv, [256, 3, 2]], 41 | [[-1, 14], 1, Concat, [1]], # cat head P4 42 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 43 | 44 | [-1, 1, Conv, [512, 3, 2]], 45 | [[-1, 10], 1, Concat, [1]], # cat head P5 46 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 47 | 48 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 49 | ] 50 | -------------------------------------------------------------------------------- /models/hub/yolov5s-ghost.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, GhostConv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3Ghost, [128]], 18 | [-1, 1, GhostConv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3Ghost, [256]], 20 | [-1, 1, GhostConv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3Ghost, [512]], 22 | [-1, 1, GhostConv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3Ghost, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, GhostConv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3Ghost, [512, False]], # 13 33 | 34 | [-1, 1, GhostConv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3Ghost, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, GhostConv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3Ghost, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, GhostConv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3Ghost, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/hub/yolov5s-transformer.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3TR, [1024]], # 9 <--- C3TR() Transformer module 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/hub/yolov5s6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 7 | anchors: 8 | - [19,27, 44,40, 38,94] # P3/8 9 | - [96,68, 86,152, 180,137] # P4/16 10 | - [140,301, 303,264, 238,542] # P5/32 11 | - [436,615, 739,380, 925,792] # P6/64 12 | 13 | # YOLOv5 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 17 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 18 | [-1, 3, C3, [128]], 19 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 20 | [-1, 6, C3, [256]], 21 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 22 | [-1, 9, C3, [512]], 23 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 24 | [-1, 3, C3, [768]], 25 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 26 | [-1, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 head 31 | head: 32 | [[-1, 1, Conv, [768, 1, 1]], 33 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 34 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 35 | [-1, 3, C3, [768, False]], # 15 36 | 37 | [-1, 1, Conv, [512, 1, 1]], 38 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 39 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 40 | [-1, 3, C3, [512, False]], # 19 41 | 42 | [-1, 1, Conv, [256, 1, 1]], 43 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 44 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 45 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 46 | 47 | [-1, 1, Conv, [256, 3, 2]], 48 | [[-1, 20], 1, Concat, [1]], # cat head P4 49 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 50 | 51 | [-1, 1, Conv, [512, 3, 2]], 52 | [[-1, 16], 1, Concat, [1]], # cat head P5 53 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 54 | 55 | [-1, 1, Conv, [768, 3, 2]], 56 | [[-1, 12], 1, Concat, [1]], # cat head P6 57 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 58 | 59 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /models/hub/yolov5x6.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.33 # model depth multiple 6 | width_multiple: 1.25 # layer channel multiple 7 | anchors: 8 | - [19,27, 44,40, 38,94] # P3/8 9 | - [96,68, 86,152, 180,137] # P4/16 10 | - [140,301, 303,264, 238,542] # P5/32 11 | - [436,615, 739,380, 925,792] # P6/64 12 | 13 | # YOLOv5 v6.0 backbone 14 | backbone: 15 | # [from, number, module, args] 16 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 17 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 18 | [-1, 3, C3, [128]], 19 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 20 | [-1, 6, C3, [256]], 21 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 22 | [-1, 9, C3, [512]], 23 | [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 24 | [-1, 3, C3, [768]], 25 | [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 26 | [-1, 3, C3, [1024]], 27 | [-1, 1, SPPF, [1024, 5]], # 11 28 | ] 29 | 30 | # YOLOv5 v6.0 head 31 | head: 32 | [[-1, 1, Conv, [768, 1, 1]], 33 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 34 | [[-1, 8], 1, Concat, [1]], # cat backbone P5 35 | [-1, 3, C3, [768, False]], # 15 36 | 37 | [-1, 1, Conv, [512, 1, 1]], 38 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 39 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 40 | [-1, 3, C3, [512, False]], # 19 41 | 42 | [-1, 1, Conv, [256, 1, 1]], 43 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 44 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 45 | [-1, 3, C3, [256, False]], # 23 (P3/8-small) 46 | 47 | [-1, 1, Conv, [256, 3, 2]], 48 | [[-1, 20], 1, Concat, [1]], # cat head P4 49 | [-1, 3, C3, [512, False]], # 26 (P4/16-medium) 50 | 51 | [-1, 1, Conv, [512, 3, 2]], 52 | [[-1, 16], 1, Concat, [1]], # cat head P5 53 | [-1, 3, C3, [768, False]], # 29 (P5/32-large) 54 | 55 | [-1, 1, Conv, [768, 3, 2]], 56 | [[-1, 12], 1, Concat, [1]], # cat head P6 57 | [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) 58 | 59 | [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) 60 | ] 61 | -------------------------------------------------------------------------------- /models/segment/yolov5l-seg.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/segment/yolov5m-seg.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.67 # model depth multiple 6 | width_multiple: 0.75 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/segment/yolov5n-seg.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.25 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/segment/yolov5s-seg.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.5 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/segment/yolov5x-seg.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.33 # model depth multiple 6 | width_multiple: 1.25 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Segment, [nc, anchors, 32, 256]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolov5l.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.0 # model depth multiple 6 | width_multiple: 1.0 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolov5m.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.67 # model depth multiple 6 | width_multiple: 0.75 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolov5n.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.25 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolov5s.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 0.33 # model depth multiple 6 | width_multiple: 0.50 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /models/yolov5x.yaml: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # Parameters 4 | nc: 80 # number of classes 5 | depth_multiple: 1.33 # model depth multiple 6 | width_multiple: 1.25 # layer channel multiple 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 v6.0 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, C3, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 6, C3, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, C3, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 3, C3, [1024]], 24 | [-1, 1, SPPF, [1024, 5]], # 9 25 | ] 26 | 27 | # YOLOv5 v6.0 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, C3, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /setup.txt: -------------------------------------------------------------------------------- 1 | # Install libraries 2 | sudo apt-get update 3 | sudo apt-get install -y liblapack-dev libblas-dev gfortran libfreetype6-dev libopenblas-base libopenmpi-dev libjpeg-dev zlib1g-dev 4 | sudo apt-get install -y python3-pip 5 | 6 | # Update Pip 7 | python3 -m pip install --upgrade pip 8 | 9 | # Install below necessary packages 10 | # For numpy, first uninstall the version already installed, and then install numpy==1.19.0 11 | numpy==1.19.0 12 | pandas==0.22.0 13 | Pillow==8.4.0 14 | PyYAML==3.12 15 | scipy==1.5.4 16 | psutil 17 | tqdm==4.64.1 18 | imutils 19 | 20 | # Install Pycuda 21 | export PATH=/usr/local/cuda-10.2/bin${PATH:+:${PATH}} 22 | export LD_LIBRARY_PATH=/usr/local/cuda-10.2/lib64:$LD_LIBRARY_PATH 23 | python3 -m pip install pycuda --user 24 | 25 | # Install Seaborn 26 | sudo apt install python3-seaborn 27 | 28 | # Install torch & torchvision 29 | wget https://nvidia.box.com/shared/static/fjtbno0vpo676a25cgvuqc1wty0fkkg6.whl -O torch-1.10.0-cp36-cp36m-linux_aarch64.whl 30 | pip3 install torch-1.10.0-cp36-cp36m-linux_aarch64.whl 31 | git clone --branch v0.11.1 https://github.com/pytorch/vision torchvision 32 | cd torchvision 33 | sudo python3 setup.py install 34 | 35 | # Not required but good library 36 | jetson-stats==3.1.4 37 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | utils/initialization 4 | """ 5 | 6 | import contextlib 7 | import platform 8 | import threading 9 | 10 | 11 | def emojis(str=''): 12 | # Return platform-dependent emoji-safe version of string 13 | return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str 14 | 15 | 16 | class TryExcept(contextlib.ContextDecorator): 17 | # YOLOv5 TryExcept class. Usage: @TryExcept() decorator or 'with TryExcept():' context manager 18 | def __init__(self, msg=''): 19 | self.msg = msg 20 | 21 | def __enter__(self): 22 | pass 23 | 24 | def __exit__(self, exc_type, value, traceback): 25 | if value: 26 | print(emojis(f"{self.msg}{': ' if self.msg else ''}{value}")) 27 | return True 28 | 29 | 30 | def threaded(func): 31 | # Multi-threads a target function and returns thread. Usage: @threaded decorator 32 | def wrapper(*args, **kwargs): 33 | thread = threading.Thread(target=func, args=args, kwargs=kwargs, daemon=True) 34 | thread.start() 35 | return thread 36 | 37 | return wrapper 38 | 39 | 40 | def join_threads(verbose=False): 41 | # Join all daemon threads, i.e. atexit.register(lambda: join_threads()) 42 | main_thread = threading.current_thread() 43 | for t in threading.enumerate(): 44 | if t is not main_thread: 45 | if verbose: 46 | print(f'Joining thread {t.name}') 47 | t.join() 48 | 49 | 50 | def notebook_init(verbose=True): 51 | # Check system software and hardware 52 | print('Checking setup...') 53 | 54 | import os 55 | import shutil 56 | 57 | from utils.general import check_font, check_requirements, is_colab 58 | from utils.torch_utils import select_device # imports 59 | 60 | check_font() 61 | 62 | import psutil 63 | 64 | if is_colab(): 65 | shutil.rmtree('/content/sample_data', ignore_errors=True) # remove colab /sample_data directory 66 | 67 | # System info 68 | display = None 69 | if verbose: 70 | gb = 1 << 30 # bytes to GiB (1024 ** 3) 71 | ram = psutil.virtual_memory().total 72 | total, used, free = shutil.disk_usage('/') 73 | with contextlib.suppress(Exception): # clear display if ipython is installed 74 | from IPython import display 75 | display.clear_output() 76 | s = f'({os.cpu_count()} CPUs, {ram / gb:.1f} GB RAM, {(total - free) / gb:.1f}/{total / gb:.1f} GB disk)' 77 | else: 78 | s = '' 79 | 80 | select_device(newline=False) 81 | print(emojis(f'Setup complete ✅ {s}')) 82 | return display 83 | -------------------------------------------------------------------------------- /utils/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/augmentations.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/__pycache__/augmentations.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/autoanchor.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/__pycache__/autoanchor.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/dataloaders.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/__pycache__/dataloaders.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/downloads.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/__pycache__/downloads.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/general.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/__pycache__/general.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/metrics.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/__pycache__/metrics.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/plots.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/__pycache__/plots.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/torch_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/__pycache__/torch_utils.cpython-36.pyc -------------------------------------------------------------------------------- /utils/activations.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Activation functions 4 | """ 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | 10 | 11 | class SiLU(nn.Module): 12 | # SiLU activation https://arxiv.org/pdf/1606.08415.pdf 13 | @staticmethod 14 | def forward(x): 15 | return x * torch.sigmoid(x) 16 | 17 | 18 | class Hardswish(nn.Module): 19 | # Hard-SiLU activation 20 | @staticmethod 21 | def forward(x): 22 | # return x * F.hardsigmoid(x) # for TorchScript and CoreML 23 | return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0 # for TorchScript, CoreML and ONNX 24 | 25 | 26 | class Mish(nn.Module): 27 | # Mish activation https://github.com/digantamisra98/Mish 28 | @staticmethod 29 | def forward(x): 30 | return x * F.softplus(x).tanh() 31 | 32 | 33 | class MemoryEfficientMish(nn.Module): 34 | # Mish activation memory-efficient 35 | class F(torch.autograd.Function): 36 | 37 | @staticmethod 38 | def forward(ctx, x): 39 | ctx.save_for_backward(x) 40 | return x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x))) 41 | 42 | @staticmethod 43 | def backward(ctx, grad_output): 44 | x = ctx.saved_tensors[0] 45 | sx = torch.sigmoid(x) 46 | fx = F.softplus(x).tanh() 47 | return grad_output * (fx + x * sx * (1 - fx * fx)) 48 | 49 | def forward(self, x): 50 | return self.F.apply(x) 51 | 52 | 53 | class FReLU(nn.Module): 54 | # FReLU activation https://arxiv.org/abs/2007.11824 55 | def __init__(self, c1, k=3): # ch_in, kernel 56 | super().__init__() 57 | self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1, bias=False) 58 | self.bn = nn.BatchNorm2d(c1) 59 | 60 | def forward(self, x): 61 | return torch.max(x, self.bn(self.conv(x))) 62 | 63 | 64 | class AconC(nn.Module): 65 | r""" ACON activation (activate or not) 66 | AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter 67 | according to "Activate or Not: Learning Customized Activation" . 68 | """ 69 | 70 | def __init__(self, c1): 71 | super().__init__() 72 | self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1)) 73 | self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1)) 74 | self.beta = nn.Parameter(torch.ones(1, c1, 1, 1)) 75 | 76 | def forward(self, x): 77 | dpx = (self.p1 - self.p2) * x 78 | return dpx * torch.sigmoid(self.beta * dpx) + self.p2 * x 79 | 80 | 81 | class MetaAconC(nn.Module): 82 | r""" ACON activation (activate or not) 83 | MetaAconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is generated by a small network 84 | according to "Activate or Not: Learning Customized Activation" . 85 | """ 86 | 87 | def __init__(self, c1, k=1, s=1, r=16): # ch_in, kernel, stride, r 88 | super().__init__() 89 | c2 = max(r, c1 // r) 90 | self.p1 = nn.Parameter(torch.randn(1, c1, 1, 1)) 91 | self.p2 = nn.Parameter(torch.randn(1, c1, 1, 1)) 92 | self.fc1 = nn.Conv2d(c1, c2, k, s, bias=True) 93 | self.fc2 = nn.Conv2d(c2, c1, k, s, bias=True) 94 | # self.bn1 = nn.BatchNorm2d(c2) 95 | # self.bn2 = nn.BatchNorm2d(c1) 96 | 97 | def forward(self, x): 98 | y = x.mean(dim=2, keepdims=True).mean(dim=3, keepdims=True) 99 | # batch-size 1 bug/instabilities https://github.com/ultralytics/yolov5/issues/2891 100 | # beta = torch.sigmoid(self.bn2(self.fc2(self.bn1(self.fc1(y))))) # bug/unstable 101 | beta = torch.sigmoid(self.fc2(self.fc1(y))) # bug patch BN layers removed 102 | dpx = (self.p1 - self.p2) * x 103 | return dpx * torch.sigmoid(beta * dpx) + self.p2 * x 104 | -------------------------------------------------------------------------------- /utils/autoanchor.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | AutoAnchor utils 4 | """ 5 | 6 | import random 7 | 8 | import numpy as np 9 | import torch 10 | import yaml 11 | from tqdm import tqdm 12 | 13 | from utils import TryExcept 14 | from utils.general import LOGGER, TQDM_BAR_FORMAT, colorstr 15 | 16 | PREFIX = colorstr('AutoAnchor: ') 17 | 18 | 19 | def check_anchor_order(m): 20 | # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary 21 | a = m.anchors.prod(-1).mean(-1).view(-1) # mean anchor area per output layer 22 | da = a[-1] - a[0] # delta a 23 | ds = m.stride[-1] - m.stride[0] # delta s 24 | if da and (da.sign() != ds.sign()): # same order 25 | LOGGER.info(f'{PREFIX}Reversing anchor order') 26 | m.anchors[:] = m.anchors.flip(0) 27 | 28 | 29 | @TryExcept(f'{PREFIX}ERROR') 30 | def check_anchors(dataset, model, thr=4.0, imgsz=640): 31 | # Check anchor fit to data, recompute if necessary 32 | m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect() 33 | shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True) 34 | scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale 35 | wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh 36 | 37 | def metric(k): # compute metric 38 | r = wh[:, None] / k[None] 39 | x = torch.min(r, 1 / r).min(2)[0] # ratio metric 40 | best = x.max(1)[0] # best_x 41 | aat = (x > 1 / thr).float().sum(1).mean() # anchors above threshold 42 | bpr = (best > 1 / thr).float().mean() # best possible recall 43 | return bpr, aat 44 | 45 | stride = m.stride.to(m.anchors.device).view(-1, 1, 1) # model strides 46 | anchors = m.anchors.clone() * stride # current anchors 47 | bpr, aat = metric(anchors.cpu().view(-1, 2)) 48 | s = f'\n{PREFIX}{aat:.2f} anchors/target, {bpr:.3f} Best Possible Recall (BPR). ' 49 | if bpr > 0.98: # threshold to recompute 50 | LOGGER.info(f'{s}Current anchors are a good fit to dataset ✅') 51 | else: 52 | LOGGER.info(f'{s}Anchors are a poor fit to dataset ⚠️, attempting to improve...') 53 | na = m.anchors.numel() // 2 # number of anchors 54 | anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False) 55 | new_bpr = metric(anchors)[0] 56 | if new_bpr > bpr: # replace anchors 57 | anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors) 58 | m.anchors[:] = anchors.clone().view_as(m.anchors) 59 | check_anchor_order(m) # must be in pixel-space (not grid-space) 60 | m.anchors /= stride 61 | s = f'{PREFIX}Done ✅ (optional: update model *.yaml to use these anchors in the future)' 62 | else: 63 | s = f'{PREFIX}Done ⚠️ (original anchors better than new anchors, proceeding with original anchors)' 64 | LOGGER.info(s) 65 | 66 | 67 | def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True): 68 | """ Creates kmeans-evolved anchors from training dataset 69 | 70 | Arguments: 71 | dataset: path to data.yaml, or a loaded dataset 72 | n: number of anchors 73 | img_size: image size used for training 74 | thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 75 | gen: generations to evolve anchors using genetic algorithm 76 | verbose: print all results 77 | 78 | Return: 79 | k: kmeans evolved anchors 80 | 81 | Usage: 82 | from utils.autoanchor import *; _ = kmean_anchors() 83 | """ 84 | from scipy.cluster.vq import kmeans 85 | 86 | npr = np.random 87 | thr = 1 / thr 88 | 89 | def metric(k, wh): # compute metrics 90 | r = wh[:, None] / k[None] 91 | x = torch.min(r, 1 / r).min(2)[0] # ratio metric 92 | # x = wh_iou(wh, torch.tensor(k)) # iou metric 93 | return x, x.max(1)[0] # x, best_x 94 | 95 | def anchor_fitness(k): # mutation fitness 96 | _, best = metric(torch.tensor(k, dtype=torch.float32), wh) 97 | return (best * (best > thr).float()).mean() # fitness 98 | 99 | def print_results(k, verbose=True): 100 | k = k[np.argsort(k.prod(1))] # sort small to large 101 | x, best = metric(k, wh0) 102 | bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr 103 | s = f'{PREFIX}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr\n' \ 104 | f'{PREFIX}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, ' \ 105 | f'past_thr={x[x > thr].mean():.3f}-mean: ' 106 | for x in k: 107 | s += '%i,%i, ' % (round(x[0]), round(x[1])) 108 | if verbose: 109 | LOGGER.info(s[:-2]) 110 | return k 111 | 112 | if isinstance(dataset, str): # *.yaml file 113 | with open(dataset, errors='ignore') as f: 114 | data_dict = yaml.safe_load(f) # model dict 115 | from utils.dataloaders import LoadImagesAndLabels 116 | dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True) 117 | 118 | # Get label wh 119 | shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True) 120 | wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh 121 | 122 | # Filter 123 | i = (wh0 < 3.0).any(1).sum() 124 | if i: 125 | LOGGER.info(f'{PREFIX}WARNING ⚠️ Extremely small objects found: {i} of {len(wh0)} labels are <3 pixels in size') 126 | wh = wh0[(wh0 >= 2.0).any(1)].astype(np.float32) # filter > 2 pixels 127 | # wh = wh * (npr.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1 128 | 129 | # Kmeans init 130 | try: 131 | LOGGER.info(f'{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...') 132 | assert n <= len(wh) # apply overdetermined constraint 133 | s = wh.std(0) # sigmas for whitening 134 | k = kmeans(wh / s, n, iter=30)[0] * s # points 135 | assert n == len(k) # kmeans may return fewer points than requested if wh is insufficient or too similar 136 | except Exception: 137 | LOGGER.warning(f'{PREFIX}WARNING ⚠️ switching strategies from kmeans to random init') 138 | k = np.sort(npr.rand(n * 2)).reshape(n, 2) * img_size # random init 139 | wh, wh0 = (torch.tensor(x, dtype=torch.float32) for x in (wh, wh0)) 140 | k = print_results(k, verbose=False) 141 | 142 | # Plot 143 | # k, d = [None] * 20, [None] * 20 144 | # for i in tqdm(range(1, 21)): 145 | # k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance 146 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True) 147 | # ax = ax.ravel() 148 | # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.') 149 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh 150 | # ax[0].hist(wh[wh[:, 0]<100, 0],400) 151 | # ax[1].hist(wh[wh[:, 1]<100, 1],400) 152 | # fig.savefig('wh.png', dpi=200) 153 | 154 | # Evolve 155 | f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma 156 | pbar = tqdm(range(gen), bar_format=TQDM_BAR_FORMAT) # progress bar 157 | for _ in pbar: 158 | v = np.ones(sh) 159 | while (v == 1).all(): # mutate until a change occurs (prevent duplicates) 160 | v = ((npr.random(sh) < mp) * random.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0) 161 | kg = (k.copy() * v).clip(min=2.0) 162 | fg = anchor_fitness(kg) 163 | if fg > f: 164 | f, k = fg, kg.copy() 165 | pbar.desc = f'{PREFIX}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}' 166 | if verbose: 167 | print_results(k, verbose) 168 | 169 | return print_results(k).astype(np.float32) 170 | -------------------------------------------------------------------------------- /utils/autobatch.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Auto-batch utils 4 | """ 5 | 6 | from copy import deepcopy 7 | 8 | import numpy as np 9 | import torch 10 | 11 | from utils.general import LOGGER, colorstr 12 | from utils.torch_utils import profile 13 | 14 | 15 | def check_train_batch_size(model, imgsz=640, amp=True): 16 | # Check YOLOv5 training batch size 17 | with torch.cuda.amp.autocast(amp): 18 | return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size 19 | 20 | 21 | def autobatch(model, imgsz=640, fraction=0.8, batch_size=16): 22 | # Automatically estimate best YOLOv5 batch size to use `fraction` of available CUDA memory 23 | # Usage: 24 | # import torch 25 | # from utils.autobatch import autobatch 26 | # model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False) 27 | # print(autobatch(model)) 28 | 29 | # Check device 30 | prefix = colorstr('AutoBatch: ') 31 | LOGGER.info(f'{prefix}Computing optimal batch size for --imgsz {imgsz}') 32 | device = next(model.parameters()).device # get model device 33 | if device.type == 'cpu': 34 | LOGGER.info(f'{prefix}CUDA not detected, using default CPU batch-size {batch_size}') 35 | return batch_size 36 | if torch.backends.cudnn.benchmark: 37 | LOGGER.info(f'{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}') 38 | return batch_size 39 | 40 | # Inspect CUDA memory 41 | gb = 1 << 30 # bytes to GiB (1024 ** 3) 42 | d = str(device).upper() # 'CUDA:0' 43 | properties = torch.cuda.get_device_properties(device) # device properties 44 | t = properties.total_memory / gb # GiB total 45 | r = torch.cuda.memory_reserved(device) / gb # GiB reserved 46 | a = torch.cuda.memory_allocated(device) / gb # GiB allocated 47 | f = t - (r + a) # GiB free 48 | LOGGER.info(f'{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free') 49 | 50 | # Profile batch sizes 51 | batch_sizes = [1, 2, 4, 8, 16] 52 | try: 53 | img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes] 54 | results = profile(img, model, n=3, device=device) 55 | except Exception as e: 56 | LOGGER.warning(f'{prefix}{e}') 57 | 58 | # Fit a solution 59 | y = [x[2] for x in results if x] # memory [2] 60 | p = np.polyfit(batch_sizes[:len(y)], y, deg=1) # first degree polynomial fit 61 | b = int((f * fraction - p[1]) / p[0]) # y intercept (optimal batch size) 62 | if None in results: # some sizes failed 63 | i = results.index(None) # first fail index 64 | if b >= batch_sizes[i]: # y intercept above failure point 65 | b = batch_sizes[max(i - 1, 0)] # select prior safe point 66 | if b < 1 or b > 1024: # b outside of safe range 67 | b = batch_size 68 | LOGGER.warning(f'{prefix}WARNING ⚠️ CUDA anomaly detected, recommend restart environment and retry command.') 69 | 70 | fraction = (np.polyval(p, b) + r + a) / t # actual fraction predicted 71 | LOGGER.info(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅') 72 | return b 73 | -------------------------------------------------------------------------------- /utils/aws/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/aws/__init__.py -------------------------------------------------------------------------------- /utils/aws/mime.sh: -------------------------------------------------------------------------------- 1 | # AWS EC2 instance startup 'MIME' script https://aws.amazon.com/premiumsupport/knowledge-center/execute-user-data-ec2/ 2 | # This script will run on every instance restart, not only on first start 3 | # --- DO NOT COPY ABOVE COMMENTS WHEN PASTING INTO USERDATA --- 4 | 5 | Content-Type: multipart/mixed; boundary="//" 6 | MIME-Version: 1.0 7 | 8 | --// 9 | Content-Type: text/cloud-config; charset="us-ascii" 10 | MIME-Version: 1.0 11 | Content-Transfer-Encoding: 7bit 12 | Content-Disposition: attachment; filename="cloud-config.txt" 13 | 14 | #cloud-config 15 | cloud_final_modules: 16 | - [scripts-user, always] 17 | 18 | --// 19 | Content-Type: text/x-shellscript; charset="us-ascii" 20 | MIME-Version: 1.0 21 | Content-Transfer-Encoding: 7bit 22 | Content-Disposition: attachment; filename="userdata.txt" 23 | 24 | #!/bin/bash 25 | # --- paste contents of userdata.sh here --- 26 | --// 27 | -------------------------------------------------------------------------------- /utils/aws/resume.py: -------------------------------------------------------------------------------- 1 | # Resume all interrupted trainings in yolov5/ dir including DDP trainings 2 | # Usage: $ python utils/aws/resume.py 3 | 4 | import os 5 | import sys 6 | from pathlib import Path 7 | 8 | import torch 9 | import yaml 10 | 11 | FILE = Path(__file__).resolve() 12 | ROOT = FILE.parents[2] # YOLOv5 root directory 13 | if str(ROOT) not in sys.path: 14 | sys.path.append(str(ROOT)) # add ROOT to PATH 15 | 16 | port = 0 # --master_port 17 | path = Path('').resolve() 18 | for last in path.rglob('*/**/last.pt'): 19 | ckpt = torch.load(last) 20 | if ckpt['optimizer'] is None: 21 | continue 22 | 23 | # Load opt.yaml 24 | with open(last.parent.parent / 'opt.yaml', errors='ignore') as f: 25 | opt = yaml.safe_load(f) 26 | 27 | # Get device count 28 | d = opt['device'].split(',') # devices 29 | nd = len(d) # number of devices 30 | ddp = nd > 1 or (nd == 0 and torch.cuda.device_count() > 1) # distributed data parallel 31 | 32 | if ddp: # multi-GPU 33 | port += 1 34 | cmd = f'python -m torch.distributed.run --nproc_per_node {nd} --master_port {port} train.py --resume {last}' 35 | else: # single-GPU 36 | cmd = f'python train.py --resume {last}' 37 | 38 | cmd += ' > /dev/null 2>&1 &' # redirect output to dev/null and run in daemon thread 39 | print(cmd) 40 | os.system(cmd) 41 | -------------------------------------------------------------------------------- /utils/aws/userdata.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # AWS EC2 instance startup script https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/user-data.html 3 | # This script will run only once on first instance start (for a re-start script see mime.sh) 4 | # /home/ubuntu (ubuntu) or /home/ec2-user (amazon-linux) is working dir 5 | # Use >300 GB SSD 6 | 7 | cd home/ubuntu 8 | if [ ! -d yolov5 ]; then 9 | echo "Running first-time script." # install dependencies, download COCO, pull Docker 10 | git clone https://github.com/ultralytics/yolov5 -b master && sudo chmod -R 777 yolov5 11 | cd yolov5 12 | bash data/scripts/get_coco.sh && echo "COCO done." & 13 | sudo docker pull ultralytics/yolov5:latest && echo "Docker done." & 14 | python -m pip install --upgrade pip && pip install -r requirements.txt && python detect.py && echo "Requirements done." & 15 | wait && echo "All tasks done." # finish background tasks 16 | else 17 | echo "Running re-start script." # resume interrupted runs 18 | i=0 19 | list=$(sudo docker ps -qa) # container list i.e. $'one\ntwo\nthree\nfour' 20 | while IFS= read -r id; do 21 | ((i++)) 22 | echo "restarting container $i: $id" 23 | sudo docker start $id 24 | # sudo docker exec -it $id python train.py --resume # single-GPU 25 | sudo docker exec -d $id python utils/aws/resume.py # multi-scenario 26 | done <<<"$list" 27 | fi 28 | -------------------------------------------------------------------------------- /utils/callbacks.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Callback utils 4 | """ 5 | 6 | import threading 7 | 8 | 9 | class Callbacks: 10 | """" 11 | Handles all registered callbacks for YOLOv5 Hooks 12 | """ 13 | 14 | def __init__(self): 15 | # Define the available callbacks 16 | self._callbacks = { 17 | 'on_pretrain_routine_start': [], 18 | 'on_pretrain_routine_end': [], 19 | 'on_train_start': [], 20 | 'on_train_epoch_start': [], 21 | 'on_train_batch_start': [], 22 | 'optimizer_step': [], 23 | 'on_before_zero_grad': [], 24 | 'on_train_batch_end': [], 25 | 'on_train_epoch_end': [], 26 | 'on_val_start': [], 27 | 'on_val_batch_start': [], 28 | 'on_val_image_end': [], 29 | 'on_val_batch_end': [], 30 | 'on_val_end': [], 31 | 'on_fit_epoch_end': [], # fit = train + val 32 | 'on_model_save': [], 33 | 'on_train_end': [], 34 | 'on_params_update': [], 35 | 'teardown': [],} 36 | self.stop_training = False # set True to interrupt training 37 | 38 | def register_action(self, hook, name='', callback=None): 39 | """ 40 | Register a new action to a callback hook 41 | 42 | Args: 43 | hook: The callback hook name to register the action to 44 | name: The name of the action for later reference 45 | callback: The callback to fire 46 | """ 47 | assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}" 48 | assert callable(callback), f"callback '{callback}' is not callable" 49 | self._callbacks[hook].append({'name': name, 'callback': callback}) 50 | 51 | def get_registered_actions(self, hook=None): 52 | """" 53 | Returns all the registered actions by callback hook 54 | 55 | Args: 56 | hook: The name of the hook to check, defaults to all 57 | """ 58 | return self._callbacks[hook] if hook else self._callbacks 59 | 60 | def run(self, hook, *args, thread=False, **kwargs): 61 | """ 62 | Loop through the registered actions and fire all callbacks on main thread 63 | 64 | Args: 65 | hook: The name of the hook to check, defaults to all 66 | args: Arguments to receive from YOLOv5 67 | thread: (boolean) Run callbacks in daemon thread 68 | kwargs: Keyword Arguments to receive from YOLOv5 69 | """ 70 | 71 | assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}" 72 | for logger in self._callbacks[hook]: 73 | if thread: 74 | threading.Thread(target=logger['callback'], args=args, kwargs=kwargs, daemon=True).start() 75 | else: 76 | logger['callback'](*args, **kwargs) 77 | -------------------------------------------------------------------------------- /utils/docker/Dockerfile: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Builds ultralytics/yolov5:latest image on DockerHub https://hub.docker.com/r/ultralytics/yolov5 3 | # Image is CUDA-optimized for YOLOv5 single/multi-GPU training and inference 4 | 5 | # Start FROM PyTorch image https://hub.docker.com/r/pytorch/pytorch 6 | FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime 7 | 8 | # Downloads to user config dir 9 | ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Arial.Unicode.ttf /root/.config/Ultralytics/ 10 | 11 | # Install linux packages 12 | ENV DEBIAN_FRONTEND noninteractive 13 | RUN apt update 14 | RUN TZ=Etc/UTC apt install -y tzdata 15 | RUN apt install --no-install-recommends -y gcc git zip curl htop libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg 16 | # RUN alias python=python3 17 | 18 | # Security updates 19 | # https://security.snyk.io/vuln/SNYK-UBUNTU1804-OPENSSL-3314796 20 | RUN apt upgrade --no-install-recommends -y openssl 21 | 22 | # Create working directory 23 | RUN rm -rf /usr/src/app && mkdir -p /usr/src/app 24 | WORKDIR /usr/src/app 25 | 26 | # Copy contents 27 | # COPY . /usr/src/app (issues as not a .git directory) 28 | RUN git clone https://github.com/ultralytics/yolov5 /usr/src/app 29 | 30 | # Install pip packages 31 | COPY requirements.txt . 32 | RUN python3 -m pip install --upgrade pip wheel 33 | RUN pip install --no-cache -r requirements.txt albumentations comet gsutil notebook \ 34 | coremltools onnx onnx-simplifier onnxruntime 'openvino-dev>=2022.3' 35 | # tensorflow tensorflowjs \ 36 | 37 | # Set environment variables 38 | ENV OMP_NUM_THREADS=1 39 | 40 | # Cleanup 41 | ENV DEBIAN_FRONTEND teletype 42 | 43 | 44 | # Usage Examples ------------------------------------------------------------------------------------------------------- 45 | 46 | # Build and Push 47 | # t=ultralytics/yolov5:latest && sudo docker build -f utils/docker/Dockerfile -t $t . && sudo docker push $t 48 | 49 | # Pull and Run 50 | # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all $t 51 | 52 | # Pull and Run with local directory access 53 | # t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all -v "$(pwd)"/datasets:/usr/src/datasets $t 54 | 55 | # Kill all 56 | # sudo docker kill $(sudo docker ps -q) 57 | 58 | # Kill all image-based 59 | # sudo docker kill $(sudo docker ps -qa --filter ancestor=ultralytics/yolov5:latest) 60 | 61 | # DockerHub tag update 62 | # t=ultralytics/yolov5:latest tnew=ultralytics/yolov5:v6.2 && sudo docker pull $t && sudo docker tag $t $tnew && sudo docker push $tnew 63 | 64 | # Clean up 65 | # sudo docker system prune -a --volumes 66 | 67 | # Update Ubuntu drivers 68 | # https://www.maketecheasier.com/install-nvidia-drivers-ubuntu/ 69 | 70 | # DDP test 71 | # python -m torch.distributed.run --nproc_per_node 2 --master_port 1 train.py --epochs 3 72 | 73 | # GCP VM from Image 74 | # docker.io/ultralytics/yolov5:latest 75 | -------------------------------------------------------------------------------- /utils/docker/Dockerfile-arm64: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Builds ultralytics/yolov5:latest-arm64 image on DockerHub https://hub.docker.com/r/ultralytics/yolov5 3 | # Image is aarch64-compatible for Apple M1 and other ARM architectures i.e. Jetson Nano and Raspberry Pi 4 | 5 | # Start FROM Ubuntu image https://hub.docker.com/_/ubuntu 6 | FROM arm64v8/ubuntu:rolling 7 | 8 | # Downloads to user config dir 9 | ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Arial.Unicode.ttf /root/.config/Ultralytics/ 10 | 11 | # Install linux packages 12 | ENV DEBIAN_FRONTEND noninteractive 13 | RUN apt update 14 | RUN TZ=Etc/UTC apt install -y tzdata 15 | RUN apt install --no-install-recommends -y python3-pip git zip curl htop gcc libgl1-mesa-glx libglib2.0-0 libpython3-dev 16 | # RUN alias python=python3 17 | 18 | # Install pip packages 19 | COPY requirements.txt . 20 | RUN python3 -m pip install --upgrade pip wheel 21 | RUN pip install --no-cache -r requirements.txt albumentations gsutil notebook \ 22 | coremltools onnx onnxruntime 23 | # tensorflow-aarch64 tensorflowjs \ 24 | 25 | # Create working directory 26 | RUN mkdir -p /usr/src/app 27 | WORKDIR /usr/src/app 28 | 29 | # Copy contents 30 | # COPY . /usr/src/app (issues as not a .git directory) 31 | RUN git clone https://github.com/ultralytics/yolov5 /usr/src/app 32 | ENV DEBIAN_FRONTEND teletype 33 | 34 | 35 | # Usage Examples ------------------------------------------------------------------------------------------------------- 36 | 37 | # Build and Push 38 | # t=ultralytics/yolov5:latest-arm64 && sudo docker build --platform linux/arm64 -f utils/docker/Dockerfile-arm64 -t $t . && sudo docker push $t 39 | 40 | # Pull and Run 41 | # t=ultralytics/yolov5:latest-arm64 && sudo docker pull $t && sudo docker run -it --ipc=host -v "$(pwd)"/datasets:/usr/src/datasets $t 42 | -------------------------------------------------------------------------------- /utils/docker/Dockerfile-cpu: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | # Builds ultralytics/yolov5:latest-cpu image on DockerHub https://hub.docker.com/r/ultralytics/yolov5 3 | # Image is CPU-optimized for ONNX, OpenVINO and PyTorch YOLOv5 deployments 4 | 5 | # Start FROM Ubuntu image https://hub.docker.com/_/ubuntu 6 | FROM ubuntu:rolling 7 | 8 | # Downloads to user config dir 9 | ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Arial.Unicode.ttf /root/.config/Ultralytics/ 10 | 11 | # Install linux packages 12 | ENV DEBIAN_FRONTEND noninteractive 13 | RUN apt update 14 | RUN TZ=Etc/UTC apt install -y tzdata 15 | RUN apt install --no-install-recommends -y python3-pip git zip curl htop libgl1-mesa-glx libglib2.0-0 libpython3-dev gnupg 16 | # RUN alias python=python3 17 | 18 | # Install pip packages 19 | COPY requirements.txt . 20 | RUN python3 -m pip install --upgrade pip wheel 21 | RUN pip install --no-cache -r requirements.txt albumentations gsutil notebook \ 22 | coremltools onnx onnx-simplifier onnxruntime 'openvino-dev>=2022.3' \ 23 | # tensorflow tensorflowjs \ 24 | --extra-index-url https://download.pytorch.org/whl/cpu 25 | 26 | # Create working directory 27 | RUN mkdir -p /usr/src/app 28 | WORKDIR /usr/src/app 29 | 30 | # Copy contents 31 | # COPY . /usr/src/app (issues as not a .git directory) 32 | RUN git clone https://github.com/ultralytics/yolov5 /usr/src/app 33 | ENV DEBIAN_FRONTEND teletype 34 | 35 | 36 | # Usage Examples ------------------------------------------------------------------------------------------------------- 37 | 38 | # Build and Push 39 | # t=ultralytics/yolov5:latest-cpu && sudo docker build -f utils/docker/Dockerfile-cpu -t $t . && sudo docker push $t 40 | 41 | # Pull and Run 42 | # t=ultralytics/yolov5:latest-cpu && sudo docker pull $t && sudo docker run -it --ipc=host -v "$(pwd)"/datasets:/usr/src/datasets $t 43 | -------------------------------------------------------------------------------- /utils/downloads.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Download utils 4 | """ 5 | 6 | import logging 7 | import os 8 | import subprocess 9 | import urllib 10 | from pathlib import Path 11 | 12 | import requests 13 | import torch 14 | 15 | 16 | def is_url(url, check=True): 17 | # Check if string is URL and check if URL exists 18 | try: 19 | url = str(url) 20 | result = urllib.parse.urlparse(url) 21 | assert all([result.scheme, result.netloc]) # check if is url 22 | return (urllib.request.urlopen(url).getcode() == 200) if check else True # check if exists online 23 | except (AssertionError, urllib.request.HTTPError): 24 | return False 25 | 26 | 27 | def gsutil_getsize(url=''): 28 | # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du 29 | output = subprocess.check_output(['gsutil', 'du', url], shell=True, encoding='utf-8') 30 | if output: 31 | return int(output.split()[0]) 32 | return 0 33 | 34 | 35 | def url_getsize(url='https://ultralytics.com/images/bus.jpg'): 36 | # Return downloadable file size in bytes 37 | response = requests.head(url, allow_redirects=True) 38 | return int(response.headers.get('content-length', -1)) 39 | 40 | 41 | def curl_download(url, filename, *, silent: bool = False) -> bool: 42 | """ 43 | Download a file from a url to a filename using curl. 44 | """ 45 | silent_option = 'sS' if silent else '' # silent 46 | proc = subprocess.run([ 47 | 'curl', 48 | '-#', 49 | f'-{silent_option}L', 50 | url, 51 | '--output', 52 | filename, 53 | '--retry', 54 | '9', 55 | '-C', 56 | '-',]) 57 | return proc.returncode == 0 58 | 59 | 60 | def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''): 61 | # Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes 62 | from utils.general import LOGGER 63 | 64 | file = Path(file) 65 | assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}" 66 | try: # url1 67 | LOGGER.info(f'Downloading {url} to {file}...') 68 | torch.hub.download_url_to_file(url, str(file), progress=LOGGER.level <= logging.INFO) 69 | assert file.exists() and file.stat().st_size > min_bytes, assert_msg # check 70 | except Exception as e: # url2 71 | if file.exists(): 72 | file.unlink() # remove partial downloads 73 | LOGGER.info(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...') 74 | # curl download, retry and resume on fail 75 | curl_download(url2 or url, file) 76 | finally: 77 | if not file.exists() or file.stat().st_size < min_bytes: # check 78 | if file.exists(): 79 | file.unlink() # remove partial downloads 80 | LOGGER.info(f'ERROR: {assert_msg}\n{error_msg}') 81 | LOGGER.info('') 82 | 83 | 84 | def attempt_download(file, repo='ultralytics/yolov5', release='v7.0'): 85 | # Attempt file download from GitHub release assets if not found locally. release = 'latest', 'v7.0', etc. 86 | from utils.general import LOGGER 87 | 88 | def github_assets(repository, version='latest'): 89 | # Return GitHub repo tag (i.e. 'v7.0') and assets (i.e. ['yolov5s.pt', 'yolov5m.pt', ...]) 90 | if version != 'latest': 91 | version = f'tags/{version}' # i.e. tags/v7.0 92 | response = requests.get(f'https://api.github.com/repos/{repository}/releases/{version}').json() # github api 93 | return response['tag_name'], [x['name'] for x in response['assets']] # tag, assets 94 | 95 | file = Path(str(file).strip().replace("'", '')) 96 | if not file.exists(): 97 | # URL specified 98 | name = Path(urllib.parse.unquote(str(file))).name # decode '%2F' to '/' etc. 99 | if str(file).startswith(('http:/', 'https:/')): # download 100 | url = str(file).replace(':/', '://') # Pathlib turns :// -> :/ 101 | file = name.split('?')[0] # parse authentication https://url.com/file.txt?auth... 102 | if Path(file).is_file(): 103 | LOGGER.info(f'Found {url} locally at {file}') # file already exists 104 | else: 105 | safe_download(file=file, url=url, min_bytes=1E5) 106 | return file 107 | 108 | # GitHub assets 109 | assets = [f'yolov5{size}{suffix}.pt' for size in 'nsmlx' for suffix in ('', '6', '-cls', '-seg')] # default 110 | try: 111 | tag, assets = github_assets(repo, release) 112 | except Exception: 113 | try: 114 | tag, assets = github_assets(repo) # latest release 115 | except Exception: 116 | try: 117 | tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1] 118 | except Exception: 119 | tag = release 120 | 121 | if name in assets: 122 | file.parent.mkdir(parents=True, exist_ok=True) # make parent dir (if required) 123 | safe_download(file, 124 | url=f'https://github.com/{repo}/releases/download/{tag}/{name}', 125 | min_bytes=1E5, 126 | error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/{tag}') 127 | 128 | return str(file) 129 | -------------------------------------------------------------------------------- /utils/flask_rest_api/README.md: -------------------------------------------------------------------------------- 1 | # Flask REST API 2 | 3 | [REST](https://en.wikipedia.org/wiki/Representational_state_transfer) [API](https://en.wikipedia.org/wiki/API)s are 4 | commonly used to expose Machine Learning (ML) models to other services. This folder contains an example REST API 5 | created using Flask to expose the YOLOv5s model from [PyTorch Hub](https://pytorch.org/hub/ultralytics_yolov5/). 6 | 7 | ## Requirements 8 | 9 | [Flask](https://palletsprojects.com/p/flask/) is required. Install with: 10 | 11 | ```shell 12 | $ pip install Flask 13 | ``` 14 | 15 | ## Run 16 | 17 | After Flask installation run: 18 | 19 | ```shell 20 | $ python3 restapi.py --port 5000 21 | ``` 22 | 23 | Then use [curl](https://curl.se/) to perform a request: 24 | 25 | ```shell 26 | $ curl -X POST -F image=@zidane.jpg 'http://localhost:5000/v1/object-detection/yolov5s' 27 | ``` 28 | 29 | The model inference results are returned as a JSON response: 30 | 31 | ```json 32 | [ 33 | { 34 | "class": 0, 35 | "confidence": 0.8900438547, 36 | "height": 0.9318675399, 37 | "name": "person", 38 | "width": 0.3264600933, 39 | "xcenter": 0.7438579798, 40 | "ycenter": 0.5207948685 41 | }, 42 | { 43 | "class": 0, 44 | "confidence": 0.8440024257, 45 | "height": 0.7155083418, 46 | "name": "person", 47 | "width": 0.6546785235, 48 | "xcenter": 0.427829951, 49 | "ycenter": 0.6334488392 50 | }, 51 | { 52 | "class": 27, 53 | "confidence": 0.3771208823, 54 | "height": 0.3902671337, 55 | "name": "tie", 56 | "width": 0.0696444362, 57 | "xcenter": 0.3675483763, 58 | "ycenter": 0.7991207838 59 | }, 60 | { 61 | "class": 27, 62 | "confidence": 0.3527112305, 63 | "height": 0.1540903747, 64 | "name": "tie", 65 | "width": 0.0336618312, 66 | "xcenter": 0.7814827561, 67 | "ycenter": 0.5065554976 68 | } 69 | ] 70 | ``` 71 | 72 | An example python script to perform inference using [requests](https://docs.python-requests.org/en/master/) is given 73 | in `example_request.py` 74 | -------------------------------------------------------------------------------- /utils/flask_rest_api/example_request.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Perform test request 4 | """ 5 | 6 | import pprint 7 | 8 | import requests 9 | 10 | DETECTION_URL = 'http://localhost:5000/v1/object-detection/yolov5s' 11 | IMAGE = 'zidane.jpg' 12 | 13 | # Read image 14 | with open(IMAGE, 'rb') as f: 15 | image_data = f.read() 16 | 17 | response = requests.post(DETECTION_URL, files={'image': image_data}).json() 18 | 19 | pprint.pprint(response) 20 | -------------------------------------------------------------------------------- /utils/flask_rest_api/restapi.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Run a Flask REST API exposing one or more YOLOv5s models 4 | """ 5 | 6 | import argparse 7 | import io 8 | 9 | import torch 10 | from flask import Flask, request 11 | from PIL import Image 12 | 13 | app = Flask(__name__) 14 | models = {} 15 | 16 | DETECTION_URL = '/v1/object-detection/' 17 | 18 | 19 | @app.route(DETECTION_URL, methods=['POST']) 20 | def predict(model): 21 | if request.method != 'POST': 22 | return 23 | 24 | if request.files.get('image'): 25 | # Method 1 26 | # with request.files["image"] as f: 27 | # im = Image.open(io.BytesIO(f.read())) 28 | 29 | # Method 2 30 | im_file = request.files['image'] 31 | im_bytes = im_file.read() 32 | im = Image.open(io.BytesIO(im_bytes)) 33 | 34 | if model in models: 35 | results = models[model](im, size=640) # reduce size=320 for faster inference 36 | return results.pandas().xyxy[0].to_json(orient='records') 37 | 38 | 39 | if __name__ == '__main__': 40 | parser = argparse.ArgumentParser(description='Flask API exposing YOLOv5 model') 41 | parser.add_argument('--port', default=5000, type=int, help='port number') 42 | parser.add_argument('--model', nargs='+', default=['yolov5s'], help='model(s) to run, i.e. --model yolov5n yolov5s') 43 | opt = parser.parse_args() 44 | 45 | for m in opt.model: 46 | models[m] = torch.hub.load('ultralytics/yolov5', m, force_reload=True, skip_validation=True) 47 | 48 | app.run(host='0.0.0.0', port=opt.port) # debug=True causes Restarting with stat 49 | -------------------------------------------------------------------------------- /utils/google_app_engine/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM gcr.io/google-appengine/python 2 | 3 | # Create a virtualenv for dependencies. This isolates these packages from 4 | # system-level packages. 5 | # Use -p python3 or -p python3.7 to select python version. Default is version 2. 6 | RUN virtualenv /env -p python3 7 | 8 | # Setting these environment variables are the same as running 9 | # source /env/bin/activate. 10 | ENV VIRTUAL_ENV /env 11 | ENV PATH /env/bin:$PATH 12 | 13 | RUN apt-get update && apt-get install -y python-opencv 14 | 15 | # Copy the application's requirements.txt and run pip to install all 16 | # dependencies into the virtualenv. 17 | ADD requirements.txt /app/requirements.txt 18 | RUN pip install -r /app/requirements.txt 19 | 20 | # Add the application source code. 21 | ADD . /app 22 | 23 | # Run a WSGI server to serve the application. gunicorn must be declared as 24 | # a dependency in requirements.txt. 25 | CMD gunicorn -b :$PORT main:app 26 | -------------------------------------------------------------------------------- /utils/google_app_engine/additional_requirements.txt: -------------------------------------------------------------------------------- 1 | # add these requirements in your app on top of the existing ones 2 | pip==21.1 3 | Flask==1.0.2 4 | gunicorn==19.10.0 5 | werkzeug>=2.2.3 # not directly required, pinned by Snyk to avoid a vulnerability 6 | -------------------------------------------------------------------------------- /utils/google_app_engine/app.yaml: -------------------------------------------------------------------------------- 1 | runtime: custom 2 | env: flex 3 | 4 | service: yolov5app 5 | 6 | liveness_check: 7 | initial_delay_sec: 600 8 | 9 | manual_scaling: 10 | instances: 1 11 | resources: 12 | cpu: 1 13 | memory_gb: 4 14 | disk_size_gb: 20 15 | -------------------------------------------------------------------------------- /utils/loggers/clearml/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/loggers/clearml/__init__.py -------------------------------------------------------------------------------- /utils/loggers/clearml/clearml_utils.py: -------------------------------------------------------------------------------- 1 | """Main Logger class for ClearML experiment tracking.""" 2 | import glob 3 | import re 4 | from pathlib import Path 5 | 6 | import numpy as np 7 | import yaml 8 | 9 | from utils.plots import Annotator, colors 10 | 11 | try: 12 | import clearml 13 | from clearml import Dataset, Task 14 | 15 | assert hasattr(clearml, '__version__') # verify package import not local dir 16 | except (ImportError, AssertionError): 17 | clearml = None 18 | 19 | 20 | def construct_dataset(clearml_info_string): 21 | """Load in a clearml dataset and fill the internal data_dict with its contents. 22 | """ 23 | dataset_id = clearml_info_string.replace('clearml://', '') 24 | dataset = Dataset.get(dataset_id=dataset_id) 25 | dataset_root_path = Path(dataset.get_local_copy()) 26 | 27 | # We'll search for the yaml file definition in the dataset 28 | yaml_filenames = list(glob.glob(str(dataset_root_path / '*.yaml')) + glob.glob(str(dataset_root_path / '*.yml'))) 29 | if len(yaml_filenames) > 1: 30 | raise ValueError('More than one yaml file was found in the dataset root, cannot determine which one contains ' 31 | 'the dataset definition this way.') 32 | elif len(yaml_filenames) == 0: 33 | raise ValueError('No yaml definition found in dataset root path, check that there is a correct yaml file ' 34 | 'inside the dataset root path.') 35 | with open(yaml_filenames[0]) as f: 36 | dataset_definition = yaml.safe_load(f) 37 | 38 | assert set(dataset_definition.keys()).issuperset( 39 | {'train', 'test', 'val', 'nc', 'names'} 40 | ), "The right keys were not found in the yaml file, make sure it at least has the following keys: ('train', 'test', 'val', 'nc', 'names')" 41 | 42 | data_dict = dict() 43 | data_dict['train'] = str( 44 | (dataset_root_path / dataset_definition['train']).resolve()) if dataset_definition['train'] else None 45 | data_dict['test'] = str( 46 | (dataset_root_path / dataset_definition['test']).resolve()) if dataset_definition['test'] else None 47 | data_dict['val'] = str( 48 | (dataset_root_path / dataset_definition['val']).resolve()) if dataset_definition['val'] else None 49 | data_dict['nc'] = dataset_definition['nc'] 50 | data_dict['names'] = dataset_definition['names'] 51 | 52 | return data_dict 53 | 54 | 55 | class ClearmlLogger: 56 | """Log training runs, datasets, models, and predictions to ClearML. 57 | 58 | This logger sends information to ClearML at app.clear.ml or to your own hosted server. By default, 59 | this information includes hyperparameters, system configuration and metrics, model metrics, code information and 60 | basic data metrics and analyses. 61 | 62 | By providing additional command line arguments to train.py, datasets, 63 | models and predictions can also be logged. 64 | """ 65 | 66 | def __init__(self, opt, hyp): 67 | """ 68 | - Initialize ClearML Task, this object will capture the experiment 69 | - Upload dataset version to ClearML Data if opt.upload_dataset is True 70 | 71 | arguments: 72 | opt (namespace) -- Commandline arguments for this run 73 | hyp (dict) -- Hyperparameters for this run 74 | 75 | """ 76 | self.current_epoch = 0 77 | # Keep tracked of amount of logged images to enforce a limit 78 | self.current_epoch_logged_images = set() 79 | # Maximum number of images to log to clearML per epoch 80 | self.max_imgs_to_log_per_epoch = 16 81 | # Get the interval of epochs when bounding box images should be logged 82 | self.bbox_interval = opt.bbox_interval 83 | self.clearml = clearml 84 | self.task = None 85 | self.data_dict = None 86 | if self.clearml: 87 | self.task = Task.init( 88 | project_name=opt.project if opt.project != 'runs/train' else 'YOLOv5', 89 | task_name=opt.name if opt.name != 'exp' else 'Training', 90 | tags=['YOLOv5'], 91 | output_uri=True, 92 | reuse_last_task_id=opt.exist_ok, 93 | auto_connect_frameworks={'pytorch': False} 94 | # We disconnect pytorch auto-detection, because we added manual model save points in the code 95 | ) 96 | # ClearML's hooks will already grab all general parameters 97 | # Only the hyperparameters coming from the yaml config file 98 | # will have to be added manually! 99 | self.task.connect(hyp, name='Hyperparameters') 100 | self.task.connect(opt, name='Args') 101 | 102 | # Make sure the code is easily remotely runnable by setting the docker image to use by the remote agent 103 | self.task.set_base_docker('ultralytics/yolov5:latest', 104 | docker_arguments='--ipc=host -e="CLEARML_AGENT_SKIP_PYTHON_ENV_INSTALL=1"', 105 | docker_setup_bash_script='pip install clearml') 106 | 107 | # Get ClearML Dataset Version if requested 108 | if opt.data.startswith('clearml://'): 109 | # data_dict should have the following keys: 110 | # names, nc (number of classes), test, train, val (all three relative paths to ../datasets) 111 | self.data_dict = construct_dataset(opt.data) 112 | # Set data to data_dict because wandb will crash without this information and opt is the best way 113 | # to give it to them 114 | opt.data = self.data_dict 115 | 116 | def log_debug_samples(self, files, title='Debug Samples'): 117 | """ 118 | Log files (images) as debug samples in the ClearML task. 119 | 120 | arguments: 121 | files (List(PosixPath)) a list of file paths in PosixPath format 122 | title (str) A title that groups together images with the same values 123 | """ 124 | for f in files: 125 | if f.exists(): 126 | it = re.search(r'_batch(\d+)', f.name) 127 | iteration = int(it.groups()[0]) if it else 0 128 | self.task.get_logger().report_image(title=title, 129 | series=f.name.replace(it.group(), ''), 130 | local_path=str(f), 131 | iteration=iteration) 132 | 133 | def log_image_with_boxes(self, image_path, boxes, class_names, image, conf_threshold=0.25): 134 | """ 135 | Draw the bounding boxes on a single image and report the result as a ClearML debug sample. 136 | 137 | arguments: 138 | image_path (PosixPath) the path the original image file 139 | boxes (list): list of scaled predictions in the format - [xmin, ymin, xmax, ymax, confidence, class] 140 | class_names (dict): dict containing mapping of class int to class name 141 | image (Tensor): A torch tensor containing the actual image data 142 | """ 143 | if len(self.current_epoch_logged_images) < self.max_imgs_to_log_per_epoch and self.current_epoch >= 0: 144 | # Log every bbox_interval times and deduplicate for any intermittend extra eval runs 145 | if self.current_epoch % self.bbox_interval == 0 and image_path not in self.current_epoch_logged_images: 146 | im = np.ascontiguousarray(np.moveaxis(image.mul(255).clamp(0, 255).byte().cpu().numpy(), 0, 2)) 147 | annotator = Annotator(im=im, pil=True) 148 | for i, (conf, class_nr, box) in enumerate(zip(boxes[:, 4], boxes[:, 5], boxes[:, :4])): 149 | color = colors(i) 150 | 151 | class_name = class_names[int(class_nr)] 152 | confidence_percentage = round(float(conf) * 100, 2) 153 | label = f'{class_name}: {confidence_percentage}%' 154 | 155 | if conf > conf_threshold: 156 | annotator.rectangle(box.cpu().numpy(), outline=color) 157 | annotator.box_label(box.cpu().numpy(), label=label, color=color) 158 | 159 | annotated_image = annotator.result() 160 | self.task.get_logger().report_image(title='Bounding Boxes', 161 | series=image_path.name, 162 | iteration=self.current_epoch, 163 | image=annotated_image) 164 | self.current_epoch_logged_images.add(image_path) 165 | -------------------------------------------------------------------------------- /utils/loggers/clearml/hpo.py: -------------------------------------------------------------------------------- 1 | from clearml import Task 2 | # Connecting ClearML with the current process, 3 | # from here on everything is logged automatically 4 | from clearml.automation import HyperParameterOptimizer, UniformParameterRange 5 | from clearml.automation.optuna import OptimizerOptuna 6 | 7 | task = Task.init(project_name='Hyper-Parameter Optimization', 8 | task_name='YOLOv5', 9 | task_type=Task.TaskTypes.optimizer, 10 | reuse_last_task_id=False) 11 | 12 | # Example use case: 13 | optimizer = HyperParameterOptimizer( 14 | # This is the experiment we want to optimize 15 | base_task_id='', 16 | # here we define the hyper-parameters to optimize 17 | # Notice: The parameter name should exactly match what you see in the UI: / 18 | # For Example, here we see in the base experiment a section Named: "General" 19 | # under it a parameter named "batch_size", this becomes "General/batch_size" 20 | # If you have `argparse` for example, then arguments will appear under the "Args" section, 21 | # and you should instead pass "Args/batch_size" 22 | hyper_parameters=[ 23 | UniformParameterRange('Hyperparameters/lr0', min_value=1e-5, max_value=1e-1), 24 | UniformParameterRange('Hyperparameters/lrf', min_value=0.01, max_value=1.0), 25 | UniformParameterRange('Hyperparameters/momentum', min_value=0.6, max_value=0.98), 26 | UniformParameterRange('Hyperparameters/weight_decay', min_value=0.0, max_value=0.001), 27 | UniformParameterRange('Hyperparameters/warmup_epochs', min_value=0.0, max_value=5.0), 28 | UniformParameterRange('Hyperparameters/warmup_momentum', min_value=0.0, max_value=0.95), 29 | UniformParameterRange('Hyperparameters/warmup_bias_lr', min_value=0.0, max_value=0.2), 30 | UniformParameterRange('Hyperparameters/box', min_value=0.02, max_value=0.2), 31 | UniformParameterRange('Hyperparameters/cls', min_value=0.2, max_value=4.0), 32 | UniformParameterRange('Hyperparameters/cls_pw', min_value=0.5, max_value=2.0), 33 | UniformParameterRange('Hyperparameters/obj', min_value=0.2, max_value=4.0), 34 | UniformParameterRange('Hyperparameters/obj_pw', min_value=0.5, max_value=2.0), 35 | UniformParameterRange('Hyperparameters/iou_t', min_value=0.1, max_value=0.7), 36 | UniformParameterRange('Hyperparameters/anchor_t', min_value=2.0, max_value=8.0), 37 | UniformParameterRange('Hyperparameters/fl_gamma', min_value=0.0, max_value=4.0), 38 | UniformParameterRange('Hyperparameters/hsv_h', min_value=0.0, max_value=0.1), 39 | UniformParameterRange('Hyperparameters/hsv_s', min_value=0.0, max_value=0.9), 40 | UniformParameterRange('Hyperparameters/hsv_v', min_value=0.0, max_value=0.9), 41 | UniformParameterRange('Hyperparameters/degrees', min_value=0.0, max_value=45.0), 42 | UniformParameterRange('Hyperparameters/translate', min_value=0.0, max_value=0.9), 43 | UniformParameterRange('Hyperparameters/scale', min_value=0.0, max_value=0.9), 44 | UniformParameterRange('Hyperparameters/shear', min_value=0.0, max_value=10.0), 45 | UniformParameterRange('Hyperparameters/perspective', min_value=0.0, max_value=0.001), 46 | UniformParameterRange('Hyperparameters/flipud', min_value=0.0, max_value=1.0), 47 | UniformParameterRange('Hyperparameters/fliplr', min_value=0.0, max_value=1.0), 48 | UniformParameterRange('Hyperparameters/mosaic', min_value=0.0, max_value=1.0), 49 | UniformParameterRange('Hyperparameters/mixup', min_value=0.0, max_value=1.0), 50 | UniformParameterRange('Hyperparameters/copy_paste', min_value=0.0, max_value=1.0)], 51 | # this is the objective metric we want to maximize/minimize 52 | objective_metric_title='metrics', 53 | objective_metric_series='mAP_0.5', 54 | # now we decide if we want to maximize it or minimize it (accuracy we maximize) 55 | objective_metric_sign='max', 56 | # let us limit the number of concurrent experiments, 57 | # this in turn will make sure we do dont bombard the scheduler with experiments. 58 | # if we have an auto-scaler connected, this, by proxy, will limit the number of machine 59 | max_number_of_concurrent_tasks=1, 60 | # this is the optimizer class (actually doing the optimization) 61 | # Currently, we can choose from GridSearch, RandomSearch or OptimizerBOHB (Bayesian optimization Hyper-Band) 62 | optimizer_class=OptimizerOptuna, 63 | # If specified only the top K performing Tasks will be kept, the others will be automatically archived 64 | save_top_k_tasks_only=5, # 5, 65 | compute_time_limit=None, 66 | total_max_jobs=20, 67 | min_iteration_per_job=None, 68 | max_iteration_per_job=None, 69 | ) 70 | 71 | # report every 10 seconds, this is way too often, but we are testing here 72 | optimizer.set_report_period(10 / 60) 73 | # You can also use the line below instead to run all the optimizer tasks locally, without using queues or agent 74 | # an_optimizer.start_locally(job_complete_callback=job_complete_callback) 75 | # set the time limit for the optimization process (2 hours) 76 | optimizer.set_time_limit(in_minutes=120.0) 77 | # Start the optimization process in the local environment 78 | optimizer.start_locally() 79 | # wait until process is done (notice we are controlling the optimization process in the background) 80 | optimizer.wait() 81 | # make sure background optimization stopped 82 | optimizer.stop() 83 | 84 | print('We are done, good bye') 85 | -------------------------------------------------------------------------------- /utils/loggers/comet/comet_utils.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | from urllib.parse import urlparse 4 | 5 | try: 6 | import comet_ml 7 | except (ModuleNotFoundError, ImportError): 8 | comet_ml = None 9 | 10 | import yaml 11 | 12 | logger = logging.getLogger(__name__) 13 | 14 | COMET_PREFIX = 'comet://' 15 | COMET_MODEL_NAME = os.getenv('COMET_MODEL_NAME', 'yolov5') 16 | COMET_DEFAULT_CHECKPOINT_FILENAME = os.getenv('COMET_DEFAULT_CHECKPOINT_FILENAME', 'last.pt') 17 | 18 | 19 | def download_model_checkpoint(opt, experiment): 20 | model_dir = f'{opt.project}/{experiment.name}' 21 | os.makedirs(model_dir, exist_ok=True) 22 | 23 | model_name = COMET_MODEL_NAME 24 | model_asset_list = experiment.get_model_asset_list(model_name) 25 | 26 | if len(model_asset_list) == 0: 27 | logger.error(f'COMET ERROR: No checkpoints found for model name : {model_name}') 28 | return 29 | 30 | model_asset_list = sorted( 31 | model_asset_list, 32 | key=lambda x: x['step'], 33 | reverse=True, 34 | ) 35 | logged_checkpoint_map = {asset['fileName']: asset['assetId'] for asset in model_asset_list} 36 | 37 | resource_url = urlparse(opt.weights) 38 | checkpoint_filename = resource_url.query 39 | 40 | if checkpoint_filename: 41 | asset_id = logged_checkpoint_map.get(checkpoint_filename) 42 | else: 43 | asset_id = logged_checkpoint_map.get(COMET_DEFAULT_CHECKPOINT_FILENAME) 44 | checkpoint_filename = COMET_DEFAULT_CHECKPOINT_FILENAME 45 | 46 | if asset_id is None: 47 | logger.error(f'COMET ERROR: Checkpoint {checkpoint_filename} not found in the given Experiment') 48 | return 49 | 50 | try: 51 | logger.info(f'COMET INFO: Downloading checkpoint {checkpoint_filename}') 52 | asset_filename = checkpoint_filename 53 | 54 | model_binary = experiment.get_asset(asset_id, return_type='binary', stream=False) 55 | model_download_path = f'{model_dir}/{asset_filename}' 56 | with open(model_download_path, 'wb') as f: 57 | f.write(model_binary) 58 | 59 | opt.weights = model_download_path 60 | 61 | except Exception as e: 62 | logger.warning('COMET WARNING: Unable to download checkpoint from Comet') 63 | logger.exception(e) 64 | 65 | 66 | def set_opt_parameters(opt, experiment): 67 | """Update the opts Namespace with parameters 68 | from Comet's ExistingExperiment when resuming a run 69 | 70 | Args: 71 | opt (argparse.Namespace): Namespace of command line options 72 | experiment (comet_ml.APIExperiment): Comet API Experiment object 73 | """ 74 | asset_list = experiment.get_asset_list() 75 | resume_string = opt.resume 76 | 77 | for asset in asset_list: 78 | if asset['fileName'] == 'opt.yaml': 79 | asset_id = asset['assetId'] 80 | asset_binary = experiment.get_asset(asset_id, return_type='binary', stream=False) 81 | opt_dict = yaml.safe_load(asset_binary) 82 | for key, value in opt_dict.items(): 83 | setattr(opt, key, value) 84 | opt.resume = resume_string 85 | 86 | # Save hyperparameters to YAML file 87 | # Necessary to pass checks in training script 88 | save_dir = f'{opt.project}/{experiment.name}' 89 | os.makedirs(save_dir, exist_ok=True) 90 | 91 | hyp_yaml_path = f'{save_dir}/hyp.yaml' 92 | with open(hyp_yaml_path, 'w') as f: 93 | yaml.dump(opt.hyp, f) 94 | opt.hyp = hyp_yaml_path 95 | 96 | 97 | def check_comet_weights(opt): 98 | """Downloads model weights from Comet and updates the 99 | weights path to point to saved weights location 100 | 101 | Args: 102 | opt (argparse.Namespace): Command Line arguments passed 103 | to YOLOv5 training script 104 | 105 | Returns: 106 | None/bool: Return True if weights are successfully downloaded 107 | else return None 108 | """ 109 | if comet_ml is None: 110 | return 111 | 112 | if isinstance(opt.weights, str): 113 | if opt.weights.startswith(COMET_PREFIX): 114 | api = comet_ml.API() 115 | resource = urlparse(opt.weights) 116 | experiment_path = f'{resource.netloc}{resource.path}' 117 | experiment = api.get(experiment_path) 118 | download_model_checkpoint(opt, experiment) 119 | return True 120 | 121 | return None 122 | 123 | 124 | def check_comet_resume(opt): 125 | """Restores run parameters to its original state based on the model checkpoint 126 | and logged Experiment parameters. 127 | 128 | Args: 129 | opt (argparse.Namespace): Command Line arguments passed 130 | to YOLOv5 training script 131 | 132 | Returns: 133 | None/bool: Return True if the run is restored successfully 134 | else return None 135 | """ 136 | if comet_ml is None: 137 | return 138 | 139 | if isinstance(opt.resume, str): 140 | if opt.resume.startswith(COMET_PREFIX): 141 | api = comet_ml.API() 142 | resource = urlparse(opt.resume) 143 | experiment_path = f'{resource.netloc}{resource.path}' 144 | experiment = api.get(experiment_path) 145 | set_opt_parameters(opt, experiment) 146 | download_model_checkpoint(opt, experiment) 147 | 148 | return True 149 | 150 | return None 151 | -------------------------------------------------------------------------------- /utils/loggers/comet/hpo.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import logging 4 | import os 5 | import sys 6 | from pathlib import Path 7 | 8 | import comet_ml 9 | 10 | logger = logging.getLogger(__name__) 11 | 12 | FILE = Path(__file__).resolve() 13 | ROOT = FILE.parents[3] # YOLOv5 root directory 14 | if str(ROOT) not in sys.path: 15 | sys.path.append(str(ROOT)) # add ROOT to PATH 16 | 17 | from train import train 18 | from utils.callbacks import Callbacks 19 | from utils.general import increment_path 20 | from utils.torch_utils import select_device 21 | 22 | # Project Configuration 23 | config = comet_ml.config.get_config() 24 | COMET_PROJECT_NAME = config.get_string(os.getenv('COMET_PROJECT_NAME'), 'comet.project_name', default='yolov5') 25 | 26 | 27 | def get_args(known=False): 28 | parser = argparse.ArgumentParser() 29 | parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path') 30 | parser.add_argument('--cfg', type=str, default='', help='model.yaml path') 31 | parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') 32 | parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path') 33 | parser.add_argument('--epochs', type=int, default=300, help='total training epochs') 34 | parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch') 35 | parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)') 36 | parser.add_argument('--rect', action='store_true', help='rectangular training') 37 | parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training') 38 | parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') 39 | parser.add_argument('--noval', action='store_true', help='only validate final epoch') 40 | parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor') 41 | parser.add_argument('--noplots', action='store_true', help='save no plot files') 42 | parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations') 43 | parser.add_argument('--bucket', type=str, default='', help='gsutil bucket') 44 | parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"') 45 | parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training') 46 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 47 | parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%') 48 | parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class') 49 | parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer') 50 | parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode') 51 | parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)') 52 | parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name') 53 | parser.add_argument('--name', default='exp', help='save to project/name') 54 | parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') 55 | parser.add_argument('--quad', action='store_true', help='quad dataloader') 56 | parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler') 57 | parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon') 58 | parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)') 59 | parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2') 60 | parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)') 61 | parser.add_argument('--seed', type=int, default=0, help='Global training seed') 62 | parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify') 63 | 64 | # Weights & Biases arguments 65 | parser.add_argument('--entity', default=None, help='W&B: Entity') 66 | parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='W&B: Upload data, "val" option') 67 | parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval') 68 | parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use') 69 | 70 | # Comet Arguments 71 | parser.add_argument('--comet_optimizer_config', type=str, help='Comet: Path to a Comet Optimizer Config File.') 72 | parser.add_argument('--comet_optimizer_id', type=str, help='Comet: ID of the Comet Optimizer sweep.') 73 | parser.add_argument('--comet_optimizer_objective', type=str, help="Comet: Set to 'minimize' or 'maximize'.") 74 | parser.add_argument('--comet_optimizer_metric', type=str, help='Comet: Metric to Optimize.') 75 | parser.add_argument('--comet_optimizer_workers', 76 | type=int, 77 | default=1, 78 | help='Comet: Number of Parallel Workers to use with the Comet Optimizer.') 79 | 80 | return parser.parse_known_args()[0] if known else parser.parse_args() 81 | 82 | 83 | def run(parameters, opt): 84 | hyp_dict = {k: v for k, v in parameters.items() if k not in ['epochs', 'batch_size']} 85 | 86 | opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok or opt.evolve)) 87 | opt.batch_size = parameters.get('batch_size') 88 | opt.epochs = parameters.get('epochs') 89 | 90 | device = select_device(opt.device, batch_size=opt.batch_size) 91 | train(hyp_dict, opt, device, callbacks=Callbacks()) 92 | 93 | 94 | if __name__ == '__main__': 95 | opt = get_args(known=True) 96 | 97 | opt.weights = str(opt.weights) 98 | opt.cfg = str(opt.cfg) 99 | opt.data = str(opt.data) 100 | opt.project = str(opt.project) 101 | 102 | optimizer_id = os.getenv('COMET_OPTIMIZER_ID') 103 | if optimizer_id is None: 104 | with open(opt.comet_optimizer_config) as f: 105 | optimizer_config = json.load(f) 106 | optimizer = comet_ml.Optimizer(optimizer_config) 107 | else: 108 | optimizer = comet_ml.Optimizer(optimizer_id) 109 | 110 | opt.comet_optimizer_id = optimizer.id 111 | status = optimizer.status() 112 | 113 | opt.comet_optimizer_objective = status['spec']['objective'] 114 | opt.comet_optimizer_metric = status['spec']['metric'] 115 | 116 | logger.info('COMET INFO: Starting Hyperparameter Sweep') 117 | for parameter in optimizer.get_parameters(): 118 | run(parameter['parameters'], opt) 119 | -------------------------------------------------------------------------------- /utils/loggers/comet/optimizer_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "algorithm": "random", 3 | "parameters": { 4 | "anchor_t": { 5 | "type": "discrete", 6 | "values": [ 7 | 2, 8 | 8 9 | ] 10 | }, 11 | "batch_size": { 12 | "type": "discrete", 13 | "values": [ 14 | 16, 15 | 32, 16 | 64 17 | ] 18 | }, 19 | "box": { 20 | "type": "discrete", 21 | "values": [ 22 | 0.02, 23 | 0.2 24 | ] 25 | }, 26 | "cls": { 27 | "type": "discrete", 28 | "values": [ 29 | 0.2 30 | ] 31 | }, 32 | "cls_pw": { 33 | "type": "discrete", 34 | "values": [ 35 | 0.5 36 | ] 37 | }, 38 | "copy_paste": { 39 | "type": "discrete", 40 | "values": [ 41 | 1 42 | ] 43 | }, 44 | "degrees": { 45 | "type": "discrete", 46 | "values": [ 47 | 0, 48 | 45 49 | ] 50 | }, 51 | "epochs": { 52 | "type": "discrete", 53 | "values": [ 54 | 5 55 | ] 56 | }, 57 | "fl_gamma": { 58 | "type": "discrete", 59 | "values": [ 60 | 0 61 | ] 62 | }, 63 | "fliplr": { 64 | "type": "discrete", 65 | "values": [ 66 | 0 67 | ] 68 | }, 69 | "flipud": { 70 | "type": "discrete", 71 | "values": [ 72 | 0 73 | ] 74 | }, 75 | "hsv_h": { 76 | "type": "discrete", 77 | "values": [ 78 | 0 79 | ] 80 | }, 81 | "hsv_s": { 82 | "type": "discrete", 83 | "values": [ 84 | 0 85 | ] 86 | }, 87 | "hsv_v": { 88 | "type": "discrete", 89 | "values": [ 90 | 0 91 | ] 92 | }, 93 | "iou_t": { 94 | "type": "discrete", 95 | "values": [ 96 | 0.7 97 | ] 98 | }, 99 | "lr0": { 100 | "type": "discrete", 101 | "values": [ 102 | 1e-05, 103 | 0.1 104 | ] 105 | }, 106 | "lrf": { 107 | "type": "discrete", 108 | "values": [ 109 | 0.01, 110 | 1 111 | ] 112 | }, 113 | "mixup": { 114 | "type": "discrete", 115 | "values": [ 116 | 1 117 | ] 118 | }, 119 | "momentum": { 120 | "type": "discrete", 121 | "values": [ 122 | 0.6 123 | ] 124 | }, 125 | "mosaic": { 126 | "type": "discrete", 127 | "values": [ 128 | 0 129 | ] 130 | }, 131 | "obj": { 132 | "type": "discrete", 133 | "values": [ 134 | 0.2 135 | ] 136 | }, 137 | "obj_pw": { 138 | "type": "discrete", 139 | "values": [ 140 | 0.5 141 | ] 142 | }, 143 | "optimizer": { 144 | "type": "categorical", 145 | "values": [ 146 | "SGD", 147 | "Adam", 148 | "AdamW" 149 | ] 150 | }, 151 | "perspective": { 152 | "type": "discrete", 153 | "values": [ 154 | 0 155 | ] 156 | }, 157 | "scale": { 158 | "type": "discrete", 159 | "values": [ 160 | 0 161 | ] 162 | }, 163 | "shear": { 164 | "type": "discrete", 165 | "values": [ 166 | 0 167 | ] 168 | }, 169 | "translate": { 170 | "type": "discrete", 171 | "values": [ 172 | 0 173 | ] 174 | }, 175 | "warmup_bias_lr": { 176 | "type": "discrete", 177 | "values": [ 178 | 0, 179 | 0.2 180 | ] 181 | }, 182 | "warmup_epochs": { 183 | "type": "discrete", 184 | "values": [ 185 | 5 186 | ] 187 | }, 188 | "warmup_momentum": { 189 | "type": "discrete", 190 | "values": [ 191 | 0, 192 | 0.95 193 | ] 194 | }, 195 | "weight_decay": { 196 | "type": "discrete", 197 | "values": [ 198 | 0, 199 | 0.001 200 | ] 201 | } 202 | }, 203 | "spec": { 204 | "maxCombo": 0, 205 | "metric": "metrics/mAP_0.5", 206 | "objective": "maximize" 207 | }, 208 | "trials": 1 209 | } 210 | -------------------------------------------------------------------------------- /utils/loggers/wandb/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/loggers/wandb/__init__.py -------------------------------------------------------------------------------- /utils/loggers/wandb/wandb_utils.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | 3 | # WARNING ⚠️ wandb is deprecated and will be removed in future release. 4 | # See supported integrations at https://github.com/ultralytics/yolov5#integrations 5 | 6 | import logging 7 | import os 8 | import sys 9 | from contextlib import contextmanager 10 | from pathlib import Path 11 | 12 | from utils.general import LOGGER, colorstr 13 | 14 | FILE = Path(__file__).resolve() 15 | ROOT = FILE.parents[3] # YOLOv5 root directory 16 | if str(ROOT) not in sys.path: 17 | sys.path.append(str(ROOT)) # add ROOT to PATH 18 | RANK = int(os.getenv('RANK', -1)) 19 | DEPRECATION_WARNING = f"{colorstr('wandb')}: WARNING ⚠️ wandb is deprecated and will be removed in a future release. " \ 20 | f'See supported integrations at https://github.com/ultralytics/yolov5#integrations.' 21 | 22 | try: 23 | import wandb 24 | 25 | assert hasattr(wandb, '__version__') # verify package import not local dir 26 | LOGGER.warning(DEPRECATION_WARNING) 27 | except (ImportError, AssertionError): 28 | wandb = None 29 | 30 | 31 | class WandbLogger(): 32 | """Log training runs, datasets, models, and predictions to Weights & Biases. 33 | 34 | This logger sends information to W&B at wandb.ai. By default, this information 35 | includes hyperparameters, system configuration and metrics, model metrics, 36 | and basic data metrics and analyses. 37 | 38 | By providing additional command line arguments to train.py, datasets, 39 | models and predictions can also be logged. 40 | 41 | For more on how this logger is used, see the Weights & Biases documentation: 42 | https://docs.wandb.com/guides/integrations/yolov5 43 | """ 44 | 45 | def __init__(self, opt, run_id=None, job_type='Training'): 46 | """ 47 | - Initialize WandbLogger instance 48 | - Upload dataset if opt.upload_dataset is True 49 | - Setup training processes if job_type is 'Training' 50 | 51 | arguments: 52 | opt (namespace) -- Commandline arguments for this run 53 | run_id (str) -- Run ID of W&B run to be resumed 54 | job_type (str) -- To set the job_type for this run 55 | 56 | """ 57 | # Pre-training routine -- 58 | self.job_type = job_type 59 | self.wandb, self.wandb_run = wandb, wandb.run if wandb else None 60 | self.val_artifact, self.train_artifact = None, None 61 | self.train_artifact_path, self.val_artifact_path = None, None 62 | self.result_artifact = None 63 | self.val_table, self.result_table = None, None 64 | self.max_imgs_to_log = 16 65 | self.data_dict = None 66 | if self.wandb: 67 | self.wandb_run = wandb.init(config=opt, 68 | resume='allow', 69 | project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem, 70 | entity=opt.entity, 71 | name=opt.name if opt.name != 'exp' else None, 72 | job_type=job_type, 73 | id=run_id, 74 | allow_val_change=True) if not wandb.run else wandb.run 75 | 76 | if self.wandb_run: 77 | if self.job_type == 'Training': 78 | if isinstance(opt.data, dict): 79 | # This means another dataset manager has already processed the dataset info (e.g. ClearML) 80 | # and they will have stored the already processed dict in opt.data 81 | self.data_dict = opt.data 82 | self.setup_training(opt) 83 | 84 | def setup_training(self, opt): 85 | """ 86 | Setup the necessary processes for training YOLO models: 87 | - Attempt to download model checkpoint and dataset artifacts if opt.resume stats with WANDB_ARTIFACT_PREFIX 88 | - Update data_dict, to contain info of previous run if resumed and the paths of dataset artifact if downloaded 89 | - Setup log_dict, initialize bbox_interval 90 | 91 | arguments: 92 | opt (namespace) -- commandline arguments for this run 93 | 94 | """ 95 | self.log_dict, self.current_epoch = {}, 0 96 | self.bbox_interval = opt.bbox_interval 97 | if isinstance(opt.resume, str): 98 | model_dir, _ = self.download_model_artifact(opt) 99 | if model_dir: 100 | self.weights = Path(model_dir) / 'last.pt' 101 | config = self.wandb_run.config 102 | opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp, opt.imgsz = str( 103 | self.weights), config.save_period, config.batch_size, config.bbox_interval, config.epochs, \ 104 | config.hyp, config.imgsz 105 | 106 | if opt.bbox_interval == -1: 107 | self.bbox_interval = opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else 1 108 | if opt.evolve or opt.noplots: 109 | self.bbox_interval = opt.bbox_interval = opt.epochs + 1 # disable bbox_interval 110 | 111 | def log_model(self, path, opt, epoch, fitness_score, best_model=False): 112 | """ 113 | Log the model checkpoint as W&B artifact 114 | 115 | arguments: 116 | path (Path) -- Path of directory containing the checkpoints 117 | opt (namespace) -- Command line arguments for this run 118 | epoch (int) -- Current epoch number 119 | fitness_score (float) -- fitness score for current epoch 120 | best_model (boolean) -- Boolean representing if the current checkpoint is the best yet. 121 | """ 122 | model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', 123 | type='model', 124 | metadata={ 125 | 'original_url': str(path), 126 | 'epochs_trained': epoch + 1, 127 | 'save period': opt.save_period, 128 | 'project': opt.project, 129 | 'total_epochs': opt.epochs, 130 | 'fitness_score': fitness_score}) 131 | model_artifact.add_file(str(path / 'last.pt'), name='last.pt') 132 | wandb.log_artifact(model_artifact, 133 | aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), 'best' if best_model else '']) 134 | LOGGER.info(f'Saving model artifact on epoch {epoch + 1}') 135 | 136 | def val_one_image(self, pred, predn, path, names, im): 137 | pass 138 | 139 | def log(self, log_dict): 140 | """ 141 | save the metrics to the logging dictionary 142 | 143 | arguments: 144 | log_dict (Dict) -- metrics/media to be logged in current step 145 | """ 146 | if self.wandb_run: 147 | for key, value in log_dict.items(): 148 | self.log_dict[key] = value 149 | 150 | def end_epoch(self): 151 | """ 152 | commit the log_dict, model artifacts and Tables to W&B and flush the log_dict. 153 | 154 | arguments: 155 | best_result (boolean): Boolean representing if the result of this evaluation is best or not 156 | """ 157 | if self.wandb_run: 158 | with all_logging_disabled(): 159 | try: 160 | wandb.log(self.log_dict) 161 | except BaseException as e: 162 | LOGGER.info( 163 | f'An error occurred in wandb logger. The training will proceed without interruption. More info\n{e}' 164 | ) 165 | self.wandb_run.finish() 166 | self.wandb_run = None 167 | self.log_dict = {} 168 | 169 | def finish_run(self): 170 | """ 171 | Log metrics if any and finish the current W&B run 172 | """ 173 | if self.wandb_run: 174 | if self.log_dict: 175 | with all_logging_disabled(): 176 | wandb.log(self.log_dict) 177 | wandb.run.finish() 178 | LOGGER.warning(DEPRECATION_WARNING) 179 | 180 | 181 | @contextmanager 182 | def all_logging_disabled(highest_level=logging.CRITICAL): 183 | """ source - https://gist.github.com/simon-weber/7853144 184 | A context manager that will prevent any logging messages triggered during the body from being processed. 185 | :param highest_level: the maximum logging level in use. 186 | This would only need to be changed if a custom level greater than CRITICAL is defined. 187 | """ 188 | previous_level = logging.root.manager.disable 189 | logging.disable(highest_level) 190 | try: 191 | yield 192 | finally: 193 | logging.disable(previous_level) 194 | -------------------------------------------------------------------------------- /utils/segment/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/segment/__init__.py -------------------------------------------------------------------------------- /utils/segment/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/segment/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /utils/segment/__pycache__/general.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/utils/segment/__pycache__/general.cpython-36.pyc -------------------------------------------------------------------------------- /utils/segment/augmentations.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Image augmentation functions 4 | """ 5 | 6 | import math 7 | import random 8 | 9 | import cv2 10 | import numpy as np 11 | 12 | from ..augmentations import box_candidates 13 | from ..general import resample_segments, segment2box 14 | 15 | 16 | def mixup(im, labels, segments, im2, labels2, segments2): 17 | # Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf 18 | r = np.random.beta(32.0, 32.0) # mixup ratio, alpha=beta=32.0 19 | im = (im * r + im2 * (1 - r)).astype(np.uint8) 20 | labels = np.concatenate((labels, labels2), 0) 21 | segments = np.concatenate((segments, segments2), 0) 22 | return im, labels, segments 23 | 24 | 25 | def random_perspective(im, 26 | targets=(), 27 | segments=(), 28 | degrees=10, 29 | translate=.1, 30 | scale=.1, 31 | shear=10, 32 | perspective=0.0, 33 | border=(0, 0)): 34 | # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10)) 35 | # targets = [cls, xyxy] 36 | 37 | height = im.shape[0] + border[0] * 2 # shape(h,w,c) 38 | width = im.shape[1] + border[1] * 2 39 | 40 | # Center 41 | C = np.eye(3) 42 | C[0, 2] = -im.shape[1] / 2 # x translation (pixels) 43 | C[1, 2] = -im.shape[0] / 2 # y translation (pixels) 44 | 45 | # Perspective 46 | P = np.eye(3) 47 | P[2, 0] = random.uniform(-perspective, perspective) # x perspective (about y) 48 | P[2, 1] = random.uniform(-perspective, perspective) # y perspective (about x) 49 | 50 | # Rotation and Scale 51 | R = np.eye(3) 52 | a = random.uniform(-degrees, degrees) 53 | # a += random.choice([-180, -90, 0, 90]) # add 90deg rotations to small rotations 54 | s = random.uniform(1 - scale, 1 + scale) 55 | # s = 2 ** random.uniform(-scale, scale) 56 | R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s) 57 | 58 | # Shear 59 | S = np.eye(3) 60 | S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # x shear (deg) 61 | S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # y shear (deg) 62 | 63 | # Translation 64 | T = np.eye(3) 65 | T[0, 2] = (random.uniform(0.5 - translate, 0.5 + translate) * width) # x translation (pixels) 66 | T[1, 2] = (random.uniform(0.5 - translate, 0.5 + translate) * height) # y translation (pixels) 67 | 68 | # Combined rotation matrix 69 | M = T @ S @ R @ P @ C # order of operations (right to left) is IMPORTANT 70 | if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any(): # image changed 71 | if perspective: 72 | im = cv2.warpPerspective(im, M, dsize=(width, height), borderValue=(114, 114, 114)) 73 | else: # affine 74 | im = cv2.warpAffine(im, M[:2], dsize=(width, height), borderValue=(114, 114, 114)) 75 | 76 | # Visualize 77 | # import matplotlib.pyplot as plt 78 | # ax = plt.subplots(1, 2, figsize=(12, 6))[1].ravel() 79 | # ax[0].imshow(im[:, :, ::-1]) # base 80 | # ax[1].imshow(im2[:, :, ::-1]) # warped 81 | 82 | # Transform label coordinates 83 | n = len(targets) 84 | new_segments = [] 85 | if n: 86 | new = np.zeros((n, 4)) 87 | segments = resample_segments(segments) # upsample 88 | for i, segment in enumerate(segments): 89 | xy = np.ones((len(segment), 3)) 90 | xy[:, :2] = segment 91 | xy = xy @ M.T # transform 92 | xy = (xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]) # perspective rescale or affine 93 | 94 | # clip 95 | new[i] = segment2box(xy, width, height) 96 | new_segments.append(xy) 97 | 98 | # filter candidates 99 | i = box_candidates(box1=targets[:, 1:5].T * s, box2=new.T, area_thr=0.01) 100 | targets = targets[i] 101 | targets[:, 1:5] = new[i] 102 | new_segments = np.array(new_segments)[i] 103 | 104 | return im, targets, new_segments 105 | -------------------------------------------------------------------------------- /utils/segment/general.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import torch 4 | import torch.nn.functional as F 5 | 6 | 7 | def crop_mask(masks, boxes): 8 | """ 9 | "Crop" predicted masks by zeroing out everything not in the predicted bbox. 10 | Vectorized by Chong (thanks Chong). 11 | 12 | Args: 13 | - masks should be a size [n, h, w] tensor of masks 14 | - boxes should be a size [n, 4] tensor of bbox coords in relative point form 15 | """ 16 | 17 | n, h, w = masks.shape 18 | x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1) # x1 shape(1,1,n) 19 | r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :] # rows shape(1,w,1) 20 | c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None] # cols shape(h,1,1) 21 | 22 | return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2)) 23 | 24 | 25 | def process_mask_upsample(protos, masks_in, bboxes, shape): 26 | """ 27 | Crop after upsample. 28 | protos: [mask_dim, mask_h, mask_w] 29 | masks_in: [n, mask_dim], n is number of masks after nms 30 | bboxes: [n, 4], n is number of masks after nms 31 | shape: input_image_size, (h, w) 32 | 33 | return: h, w, n 34 | """ 35 | 36 | c, mh, mw = protos.shape # CHW 37 | masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw) 38 | masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0] # CHW 39 | masks = crop_mask(masks, bboxes) # CHW 40 | return masks.gt_(0.5) 41 | 42 | 43 | def process_mask(protos, masks_in, bboxes, shape, upsample=False): 44 | """ 45 | Crop before upsample. 46 | proto_out: [mask_dim, mask_h, mask_w] 47 | out_masks: [n, mask_dim], n is number of masks after nms 48 | bboxes: [n, 4], n is number of masks after nms 49 | shape:input_image_size, (h, w) 50 | 51 | return: h, w, n 52 | """ 53 | 54 | c, mh, mw = protos.shape # CHW 55 | ih, iw = shape 56 | masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw) # CHW 57 | 58 | downsampled_bboxes = bboxes.clone() 59 | downsampled_bboxes[:, 0] *= mw / iw 60 | downsampled_bboxes[:, 2] *= mw / iw 61 | downsampled_bboxes[:, 3] *= mh / ih 62 | downsampled_bboxes[:, 1] *= mh / ih 63 | 64 | masks = crop_mask(masks, downsampled_bboxes) # CHW 65 | if upsample: 66 | masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0] # CHW 67 | return masks.gt_(0.5) 68 | 69 | 70 | def process_mask_native(protos, masks_in, bboxes, shape): 71 | """ 72 | Crop after upsample. 73 | protos: [mask_dim, mask_h, mask_w] 74 | masks_in: [n, mask_dim], n is number of masks after nms 75 | bboxes: [n, 4], n is number of masks after nms 76 | shape: input_image_size, (h, w) 77 | 78 | return: h, w, n 79 | """ 80 | c, mh, mw = protos.shape # CHW 81 | masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw) 82 | gain = min(mh / shape[0], mw / shape[1]) # gain = old / new 83 | pad = (mw - shape[1] * gain) / 2, (mh - shape[0] * gain) / 2 # wh padding 84 | top, left = int(pad[1]), int(pad[0]) # y, x 85 | bottom, right = int(mh - pad[1]), int(mw - pad[0]) 86 | masks = masks[:, top:bottom, left:right] 87 | 88 | masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0] # CHW 89 | masks = crop_mask(masks, bboxes) # CHW 90 | return masks.gt_(0.5) 91 | 92 | 93 | def scale_image(im1_shape, masks, im0_shape, ratio_pad=None): 94 | """ 95 | img1_shape: model input shape, [h, w] 96 | img0_shape: origin pic shape, [h, w, 3] 97 | masks: [h, w, num] 98 | """ 99 | # Rescale coordinates (xyxy) from im1_shape to im0_shape 100 | if ratio_pad is None: # calculate from im0_shape 101 | gain = min(im1_shape[0] / im0_shape[0], im1_shape[1] / im0_shape[1]) # gain = old / new 102 | pad = (im1_shape[1] - im0_shape[1] * gain) / 2, (im1_shape[0] - im0_shape[0] * gain) / 2 # wh padding 103 | else: 104 | pad = ratio_pad[1] 105 | top, left = int(pad[1]), int(pad[0]) # y, x 106 | bottom, right = int(im1_shape[0] - pad[1]), int(im1_shape[1] - pad[0]) 107 | 108 | if len(masks.shape) < 2: 109 | raise ValueError(f'"len of masks shape" should be 2 or 3, but got {len(masks.shape)}') 110 | masks = masks[top:bottom, left:right] 111 | # masks = masks.permute(2, 0, 1).contiguous() 112 | # masks = F.interpolate(masks[None], im0_shape[:2], mode='bilinear', align_corners=False)[0] 113 | # masks = masks.permute(1, 2, 0).contiguous() 114 | masks = cv2.resize(masks, (im0_shape[1], im0_shape[0])) 115 | 116 | if len(masks.shape) == 2: 117 | masks = masks[:, :, None] 118 | return masks 119 | 120 | 121 | def mask_iou(mask1, mask2, eps=1e-7): 122 | """ 123 | mask1: [N, n] m1 means number of predicted objects 124 | mask2: [M, n] m2 means number of gt objects 125 | Note: n means image_w x image_h 126 | 127 | return: masks iou, [N, M] 128 | """ 129 | intersection = torch.matmul(mask1, mask2.t()).clamp(0) 130 | union = (mask1.sum(1)[:, None] + mask2.sum(1)[None]) - intersection # (area1 + area2) - intersection 131 | return intersection / (union + eps) 132 | 133 | 134 | def masks_iou(mask1, mask2, eps=1e-7): 135 | """ 136 | mask1: [N, n] m1 means number of predicted objects 137 | mask2: [N, n] m2 means number of gt objects 138 | Note: n means image_w x image_h 139 | 140 | return: masks iou, (N, ) 141 | """ 142 | intersection = (mask1 * mask2).sum(1).clamp(0) # (N, ) 143 | union = (mask1.sum(1) + mask2.sum(1))[None] - intersection # (area1 + area2) - intersection 144 | return intersection / (union + eps) 145 | 146 | 147 | def masks2segments(masks, strategy='largest'): 148 | # Convert masks(n,160,160) into segments(n,xy) 149 | segments = [] 150 | for x in masks.int().cpu().numpy().astype('uint8'): 151 | c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0] 152 | if c: 153 | if strategy == 'concat': # concatenate all segments 154 | c = np.concatenate([x.reshape(-1, 2) for x in c]) 155 | elif strategy == 'largest': # select largest segment 156 | c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2) 157 | else: 158 | c = np.zeros((0, 2)) # no segments found 159 | segments.append(c.astype('float32')) 160 | return segments 161 | -------------------------------------------------------------------------------- /utils/segment/metrics.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ 3 | Model validation metrics 4 | """ 5 | 6 | import numpy as np 7 | 8 | from ..metrics import ap_per_class 9 | 10 | 11 | def fitness(x): 12 | # Model fitness as a weighted combination of metrics 13 | w = [0.0, 0.0, 0.1, 0.9, 0.0, 0.0, 0.1, 0.9] 14 | return (x[:, :8] * w).sum(1) 15 | 16 | 17 | def ap_per_class_box_and_mask( 18 | tp_m, 19 | tp_b, 20 | conf, 21 | pred_cls, 22 | target_cls, 23 | plot=False, 24 | save_dir='.', 25 | names=(), 26 | ): 27 | """ 28 | Args: 29 | tp_b: tp of boxes. 30 | tp_m: tp of masks. 31 | other arguments see `func: ap_per_class`. 32 | """ 33 | results_boxes = ap_per_class(tp_b, 34 | conf, 35 | pred_cls, 36 | target_cls, 37 | plot=plot, 38 | save_dir=save_dir, 39 | names=names, 40 | prefix='Box')[2:] 41 | results_masks = ap_per_class(tp_m, 42 | conf, 43 | pred_cls, 44 | target_cls, 45 | plot=plot, 46 | save_dir=save_dir, 47 | names=names, 48 | prefix='Mask')[2:] 49 | 50 | results = { 51 | 'boxes': { 52 | 'p': results_boxes[0], 53 | 'r': results_boxes[1], 54 | 'ap': results_boxes[3], 55 | 'f1': results_boxes[2], 56 | 'ap_class': results_boxes[4]}, 57 | 'masks': { 58 | 'p': results_masks[0], 59 | 'r': results_masks[1], 60 | 'ap': results_masks[3], 61 | 'f1': results_masks[2], 62 | 'ap_class': results_masks[4]}} 63 | return results 64 | 65 | 66 | class Metric: 67 | 68 | def __init__(self) -> None: 69 | self.p = [] # (nc, ) 70 | self.r = [] # (nc, ) 71 | self.f1 = [] # (nc, ) 72 | self.all_ap = [] # (nc, 10) 73 | self.ap_class_index = [] # (nc, ) 74 | 75 | @property 76 | def ap50(self): 77 | """AP@0.5 of all classes. 78 | Return: 79 | (nc, ) or []. 80 | """ 81 | return self.all_ap[:, 0] if len(self.all_ap) else [] 82 | 83 | @property 84 | def ap(self): 85 | """AP@0.5:0.95 86 | Return: 87 | (nc, ) or []. 88 | """ 89 | return self.all_ap.mean(1) if len(self.all_ap) else [] 90 | 91 | @property 92 | def mp(self): 93 | """mean precision of all classes. 94 | Return: 95 | float. 96 | """ 97 | return self.p.mean() if len(self.p) else 0.0 98 | 99 | @property 100 | def mr(self): 101 | """mean recall of all classes. 102 | Return: 103 | float. 104 | """ 105 | return self.r.mean() if len(self.r) else 0.0 106 | 107 | @property 108 | def map50(self): 109 | """Mean AP@0.5 of all classes. 110 | Return: 111 | float. 112 | """ 113 | return self.all_ap[:, 0].mean() if len(self.all_ap) else 0.0 114 | 115 | @property 116 | def map(self): 117 | """Mean AP@0.5:0.95 of all classes. 118 | Return: 119 | float. 120 | """ 121 | return self.all_ap.mean() if len(self.all_ap) else 0.0 122 | 123 | def mean_results(self): 124 | """Mean of results, return mp, mr, map50, map""" 125 | return (self.mp, self.mr, self.map50, self.map) 126 | 127 | def class_result(self, i): 128 | """class-aware result, return p[i], r[i], ap50[i], ap[i]""" 129 | return (self.p[i], self.r[i], self.ap50[i], self.ap[i]) 130 | 131 | def get_maps(self, nc): 132 | maps = np.zeros(nc) + self.map 133 | for i, c in enumerate(self.ap_class_index): 134 | maps[c] = self.ap[i] 135 | return maps 136 | 137 | def update(self, results): 138 | """ 139 | Args: 140 | results: tuple(p, r, ap, f1, ap_class) 141 | """ 142 | p, r, all_ap, f1, ap_class_index = results 143 | self.p = p 144 | self.r = r 145 | self.all_ap = all_ap 146 | self.f1 = f1 147 | self.ap_class_index = ap_class_index 148 | 149 | 150 | class Metrics: 151 | """Metric for boxes and masks.""" 152 | 153 | def __init__(self) -> None: 154 | self.metric_box = Metric() 155 | self.metric_mask = Metric() 156 | 157 | def update(self, results): 158 | """ 159 | Args: 160 | results: Dict{'boxes': Dict{}, 'masks': Dict{}} 161 | """ 162 | self.metric_box.update(list(results['boxes'].values())) 163 | self.metric_mask.update(list(results['masks'].values())) 164 | 165 | def mean_results(self): 166 | return self.metric_box.mean_results() + self.metric_mask.mean_results() 167 | 168 | def class_result(self, i): 169 | return self.metric_box.class_result(i) + self.metric_mask.class_result(i) 170 | 171 | def get_maps(self, nc): 172 | return self.metric_box.get_maps(nc) + self.metric_mask.get_maps(nc) 173 | 174 | @property 175 | def ap_class_index(self): 176 | # boxes and masks have the same ap_class_index 177 | return self.metric_box.ap_class_index 178 | 179 | 180 | KEYS = [ 181 | 'train/box_loss', 182 | 'train/seg_loss', # train loss 183 | 'train/obj_loss', 184 | 'train/cls_loss', 185 | 'metrics/precision(B)', 186 | 'metrics/recall(B)', 187 | 'metrics/mAP_0.5(B)', 188 | 'metrics/mAP_0.5:0.95(B)', # metrics 189 | 'metrics/precision(M)', 190 | 'metrics/recall(M)', 191 | 'metrics/mAP_0.5(M)', 192 | 'metrics/mAP_0.5:0.95(M)', # metrics 193 | 'val/box_loss', 194 | 'val/seg_loss', # val loss 195 | 'val/obj_loss', 196 | 'val/cls_loss', 197 | 'x/lr0', 198 | 'x/lr1', 199 | 'x/lr2',] 200 | 201 | BEST_KEYS = [ 202 | 'best/epoch', 203 | 'best/precision(B)', 204 | 'best/recall(B)', 205 | 'best/mAP_0.5(B)', 206 | 'best/mAP_0.5:0.95(B)', 207 | 'best/precision(M)', 208 | 'best/recall(M)', 209 | 'best/mAP_0.5(M)', 210 | 'best/mAP_0.5:0.95(M)',] 211 | -------------------------------------------------------------------------------- /utils/segment/plots.py: -------------------------------------------------------------------------------- 1 | import contextlib 2 | import math 3 | from pathlib import Path 4 | 5 | import cv2 6 | import matplotlib.pyplot as plt 7 | import numpy as np 8 | import pandas as pd 9 | import torch 10 | 11 | from .. import threaded 12 | from ..general import xywh2xyxy 13 | from ..plots import Annotator, colors 14 | 15 | 16 | @threaded 17 | def plot_images_and_masks(images, targets, masks, paths=None, fname='images.jpg', names=None): 18 | # Plot image grid with labels 19 | if isinstance(images, torch.Tensor): 20 | images = images.cpu().float().numpy() 21 | if isinstance(targets, torch.Tensor): 22 | targets = targets.cpu().numpy() 23 | if isinstance(masks, torch.Tensor): 24 | masks = masks.cpu().numpy().astype(int) 25 | 26 | max_size = 1920 # max image size 27 | max_subplots = 16 # max image subplots, i.e. 4x4 28 | bs, _, h, w = images.shape # batch size, _, height, width 29 | bs = min(bs, max_subplots) # limit plot images 30 | ns = np.ceil(bs ** 0.5) # number of subplots (square) 31 | if np.max(images[0]) <= 1: 32 | images *= 255 # de-normalise (optional) 33 | 34 | # Build Image 35 | mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init 36 | for i, im in enumerate(images): 37 | if i == max_subplots: # if last batch has fewer images than we expect 38 | break 39 | x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin 40 | im = im.transpose(1, 2, 0) 41 | mosaic[y:y + h, x:x + w, :] = im 42 | 43 | # Resize (optional) 44 | scale = max_size / ns / max(h, w) 45 | if scale < 1: 46 | h = math.ceil(scale * h) 47 | w = math.ceil(scale * w) 48 | mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h))) 49 | 50 | # Annotate 51 | fs = int((h + w) * ns * 0.01) # font size 52 | annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names) 53 | for i in range(i + 1): 54 | x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin 55 | annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2) # borders 56 | if paths: 57 | annotator.text((x + 5, y + 5), text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames 58 | if len(targets) > 0: 59 | idx = targets[:, 0] == i 60 | ti = targets[idx] # image targets 61 | 62 | boxes = xywh2xyxy(ti[:, 2:6]).T 63 | classes = ti[:, 1].astype('int') 64 | labels = ti.shape[1] == 6 # labels if no conf column 65 | conf = None if labels else ti[:, 6] # check for confidence presence (label vs pred) 66 | 67 | if boxes.shape[1]: 68 | if boxes.max() <= 1.01: # if normalized with tolerance 0.01 69 | boxes[[0, 2]] *= w # scale to pixels 70 | boxes[[1, 3]] *= h 71 | elif scale < 1: # absolute coords need scale if image scales 72 | boxes *= scale 73 | boxes[[0, 2]] += x 74 | boxes[[1, 3]] += y 75 | for j, box in enumerate(boxes.T.tolist()): 76 | cls = classes[j] 77 | color = colors(cls) 78 | cls = names[cls] if names else cls 79 | if labels or conf[j] > 0.25: # 0.25 conf thresh 80 | label = f'{cls}' if labels else f'{cls} {conf[j]:.1f}' 81 | annotator.box_label(box, label, color=color) 82 | 83 | # Plot masks 84 | if len(masks): 85 | if masks.max() > 1.0: # mean that masks are overlap 86 | image_masks = masks[[i]] # (1, 640, 640) 87 | nl = len(ti) 88 | index = np.arange(nl).reshape(nl, 1, 1) + 1 89 | image_masks = np.repeat(image_masks, nl, axis=0) 90 | image_masks = np.where(image_masks == index, 1.0, 0.0) 91 | else: 92 | image_masks = masks[idx] 93 | 94 | im = np.asarray(annotator.im).copy() 95 | for j, box in enumerate(boxes.T.tolist()): 96 | if labels or conf[j] > 0.25: # 0.25 conf thresh 97 | color = colors(classes[j]) 98 | mh, mw = image_masks[j].shape 99 | if mh != h or mw != w: 100 | mask = image_masks[j].astype(np.uint8) 101 | mask = cv2.resize(mask, (w, h)) 102 | mask = mask.astype(bool) 103 | else: 104 | mask = image_masks[j].astype(bool) 105 | with contextlib.suppress(Exception): 106 | im[y:y + h, x:x + w, :][mask] = im[y:y + h, x:x + w, :][mask] * 0.4 + np.array(color) * 0.6 107 | annotator.fromarray(im) 108 | annotator.im.save(fname) # save 109 | 110 | 111 | def plot_results_with_masks(file='path/to/results.csv', dir='', best=True): 112 | # Plot training results.csv. Usage: from utils.plots import *; plot_results('path/to/results.csv') 113 | save_dir = Path(file).parent if file else Path(dir) 114 | fig, ax = plt.subplots(2, 8, figsize=(18, 6), tight_layout=True) 115 | ax = ax.ravel() 116 | files = list(save_dir.glob('results*.csv')) 117 | assert len(files), f'No results.csv files found in {save_dir.resolve()}, nothing to plot.' 118 | for f in files: 119 | try: 120 | data = pd.read_csv(f) 121 | index = np.argmax(0.9 * data.values[:, 8] + 0.1 * data.values[:, 7] + 0.9 * data.values[:, 12] + 122 | 0.1 * data.values[:, 11]) 123 | s = [x.strip() for x in data.columns] 124 | x = data.values[:, 0] 125 | for i, j in enumerate([1, 2, 3, 4, 5, 6, 9, 10, 13, 14, 15, 16, 7, 8, 11, 12]): 126 | y = data.values[:, j] 127 | # y[y == 0] = np.nan # don't show zero values 128 | ax[i].plot(x, y, marker='.', label=f.stem, linewidth=2, markersize=2) 129 | if best: 130 | # best 131 | ax[i].scatter(index, y[index], color='r', label=f'best:{index}', marker='*', linewidth=3) 132 | ax[i].set_title(s[j] + f'\n{round(y[index], 5)}') 133 | else: 134 | # last 135 | ax[i].scatter(x[-1], y[-1], color='r', label='last', marker='*', linewidth=3) 136 | ax[i].set_title(s[j] + f'\n{round(y[-1], 5)}') 137 | # if j in [8, 9, 10]: # share train and val loss y axes 138 | # ax[i].get_shared_y_axes().join(ax[i], ax[i - 5]) 139 | except Exception as e: 140 | print(f'Warning: Plotting error for {f}: {e}') 141 | ax[1].legend() 142 | fig.savefig(save_dir / 'results.png', dpi=200) 143 | plt.close() 144 | -------------------------------------------------------------------------------- /utils/triton.py: -------------------------------------------------------------------------------- 1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license 2 | """ Utils to interact with the Triton Inference Server 3 | """ 4 | 5 | import typing 6 | from urllib.parse import urlparse 7 | 8 | import torch 9 | 10 | 11 | class TritonRemoteModel: 12 | """ A wrapper over a model served by the Triton Inference Server. It can 13 | be configured to communicate over GRPC or HTTP. It accepts Torch Tensors 14 | as input and returns them as outputs. 15 | """ 16 | 17 | def __init__(self, url: str): 18 | """ 19 | Keyword arguments: 20 | url: Fully qualified address of the Triton server - for e.g. grpc://localhost:8000 21 | """ 22 | 23 | parsed_url = urlparse(url) 24 | if parsed_url.scheme == 'grpc': 25 | from tritonclient.grpc import InferenceServerClient, InferInput 26 | 27 | self.client = InferenceServerClient(parsed_url.netloc) # Triton GRPC client 28 | model_repository = self.client.get_model_repository_index() 29 | self.model_name = model_repository.models[0].name 30 | self.metadata = self.client.get_model_metadata(self.model_name, as_json=True) 31 | 32 | def create_input_placeholders() -> typing.List[InferInput]: 33 | return [ 34 | InferInput(i['name'], [int(s) for s in i['shape']], i['datatype']) for i in self.metadata['inputs']] 35 | 36 | else: 37 | from tritonclient.http import InferenceServerClient, InferInput 38 | 39 | self.client = InferenceServerClient(parsed_url.netloc) # Triton HTTP client 40 | model_repository = self.client.get_model_repository_index() 41 | self.model_name = model_repository[0]['name'] 42 | self.metadata = self.client.get_model_metadata(self.model_name) 43 | 44 | def create_input_placeholders() -> typing.List[InferInput]: 45 | return [ 46 | InferInput(i['name'], [int(s) for s in i['shape']], i['datatype']) for i in self.metadata['inputs']] 47 | 48 | self._create_input_placeholders_fn = create_input_placeholders 49 | 50 | @property 51 | def runtime(self): 52 | """Returns the model runtime""" 53 | return self.metadata.get('backend', self.metadata.get('platform')) 54 | 55 | def __call__(self, *args, **kwargs) -> typing.Union[torch.Tensor, typing.Tuple[torch.Tensor, ...]]: 56 | """ Invokes the model. Parameters can be provided via args or kwargs. 57 | args, if provided, are assumed to match the order of inputs of the model. 58 | kwargs are matched with the model input names. 59 | """ 60 | inputs = self._create_inputs(*args, **kwargs) 61 | response = self.client.infer(model_name=self.model_name, inputs=inputs) 62 | result = [] 63 | for output in self.metadata['outputs']: 64 | tensor = torch.as_tensor(response.as_numpy(output['name'])) 65 | result.append(tensor) 66 | return result[0] if len(result) == 1 else result 67 | 68 | def _create_inputs(self, *args, **kwargs): 69 | args_len, kwargs_len = len(args), len(kwargs) 70 | if not args_len and not kwargs_len: 71 | raise RuntimeError('No inputs provided.') 72 | if args_len and kwargs_len: 73 | raise RuntimeError('Cannot specify args and kwargs at the same time') 74 | 75 | placeholders = self._create_input_placeholders_fn() 76 | if args_len: 77 | if args_len != len(placeholders): 78 | raise RuntimeError(f'Expected {len(placeholders)} inputs, got {args_len}.') 79 | for input, value in zip(placeholders, args): 80 | input.set_data_from_numpy(value.cpu().numpy()) 81 | else: 82 | for input in placeholders: 83 | value = kwargs[input.name] 84 | input.set_data_from_numpy(value.cpu().numpy()) 85 | return placeholders 86 | -------------------------------------------------------------------------------- /videos/out.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/videos/out.jpg -------------------------------------------------------------------------------- /videos/testvideo.mp4: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/videos/testvideo.mp4 -------------------------------------------------------------------------------- /videos/testvideo1.mp4: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/videos/testvideo1.mp4 -------------------------------------------------------------------------------- /yolov5/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.10) 2 | 3 | project(yolov5) 4 | 5 | add_definitions(-std=c++11) 6 | add_definitions(-DAPI_EXPORTS) 7 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 8 | set(CMAKE_CXX_STANDARD 11) 9 | set(CMAKE_BUILD_TYPE Debug) 10 | 11 | # TODO(Call for PR): make cmake compatible with Windows 12 | set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc) 13 | enable_language(CUDA) 14 | 15 | # include and link dirs of cuda and tensorrt, you need adapt them if yours are different 16 | # cuda 17 | include_directories(/usr/local/cuda/include) 18 | link_directories(/usr/local/cuda/lib64) 19 | # tensorrt 20 | # TODO(Call for PR): make TRT path configurable from command line 21 | include_directories(/home/nvidia/TensorRT-8.2.5.1/include/) 22 | link_directories(/home/nvidia/TensorRT-8.2.5.1/lib/) 23 | 24 | include_directories(${PROJECT_SOURCE_DIR}/src/) 25 | include_directories(${PROJECT_SOURCE_DIR}/plugin/) 26 | file(GLOB_RECURSE SRCS ${PROJECT_SOURCE_DIR}/src/*.cpp ${PROJECT_SOURCE_DIR}/src/*.cu) 27 | file(GLOB_RECURSE PLUGIN_SRCS ${PROJECT_SOURCE_DIR}/plugin/*.cu) 28 | 29 | add_library(myplugins SHARED ${PLUGIN_SRCS}) 30 | target_link_libraries(myplugins nvinfer cudart) 31 | 32 | find_package(OpenCV) 33 | include_directories(${OpenCV_INCLUDE_DIRS}) 34 | 35 | add_executable(yolov5_det yolov5_det.cpp ${SRCS}) 36 | target_link_libraries(yolov5_det nvinfer) 37 | target_link_libraries(yolov5_det cudart) 38 | target_link_libraries(yolov5_det myplugins) 39 | target_link_libraries(yolov5_det ${OpenCV_LIBS}) 40 | 41 | add_executable(yolov5_cls yolov5_cls.cpp ${SRCS}) 42 | target_link_libraries(yolov5_cls nvinfer) 43 | target_link_libraries(yolov5_cls cudart) 44 | target_link_libraries(yolov5_cls myplugins) 45 | target_link_libraries(yolov5_cls ${OpenCV_LIBS}) 46 | 47 | add_executable(yolov5_seg yolov5_seg.cpp ${SRCS}) 48 | target_link_libraries(yolov5_seg nvinfer) 49 | target_link_libraries(yolov5_seg cudart) 50 | target_link_libraries(yolov5_seg myplugins) 51 | target_link_libraries(yolov5_seg ${OpenCV_LIBS}) 52 | 53 | -------------------------------------------------------------------------------- /yolov5/images/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/yolov5/images/bus.jpg -------------------------------------------------------------------------------- /yolov5/images/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/yolov5/images/zidane.jpg -------------------------------------------------------------------------------- /yolov5/plugin/yololayer.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "types.h" 4 | #include "macros.h" 5 | 6 | #include 7 | #include 8 | 9 | namespace nvinfer1 { 10 | class API YoloLayerPlugin : public IPluginV2IOExt { 11 | public: 12 | YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, bool is_segmentation, const std::vector& vYoloKernel); 13 | YoloLayerPlugin(const void* data, size_t length); 14 | ~YoloLayerPlugin(); 15 | 16 | int getNbOutputs() const TRT_NOEXCEPT override { return 1; } 17 | 18 | Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT override; 19 | 20 | int initialize() TRT_NOEXCEPT override; 21 | 22 | virtual void terminate() TRT_NOEXCEPT override {}; 23 | 24 | virtual size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override { return 0; } 25 | 26 | virtual int enqueue(int batchSize, const void* const* inputs, void*TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT override; 27 | 28 | virtual size_t getSerializationSize() const TRT_NOEXCEPT override; 29 | 30 | virtual void serialize(void* buffer) const TRT_NOEXCEPT override; 31 | 32 | bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const TRT_NOEXCEPT override { 33 | return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; 34 | } 35 | 36 | const char* getPluginType() const TRT_NOEXCEPT override; 37 | 38 | const char* getPluginVersion() const TRT_NOEXCEPT override; 39 | 40 | void destroy() TRT_NOEXCEPT override; 41 | 42 | IPluginV2IOExt* clone() const TRT_NOEXCEPT override; 43 | 44 | void setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT override; 45 | 46 | const char* getPluginNamespace() const TRT_NOEXCEPT override; 47 | 48 | DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override; 49 | 50 | bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT override; 51 | 52 | bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override; 53 | 54 | void attachToContext( 55 | cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT override; 56 | 57 | void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT override; 58 | 59 | void detachFromContext() TRT_NOEXCEPT override; 60 | 61 | private: 62 | void forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize = 1); 63 | int mThreadCount = 256; 64 | const char* mPluginNamespace; 65 | int mKernelCount; 66 | int mClassCount; 67 | int mYoloV5NetWidth; 68 | int mYoloV5NetHeight; 69 | int mMaxOutObject; 70 | bool is_segmentation_; 71 | std::vector mYoloKernel; 72 | void** mAnchor; 73 | }; 74 | 75 | class API YoloPluginCreator : public IPluginCreator { 76 | public: 77 | YoloPluginCreator(); 78 | 79 | ~YoloPluginCreator() override = default; 80 | 81 | const char* getPluginName() const TRT_NOEXCEPT override; 82 | 83 | const char* getPluginVersion() const TRT_NOEXCEPT override; 84 | 85 | const PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override; 86 | 87 | IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT override; 88 | 89 | IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT override; 90 | 91 | void setPluginNamespace(const char* libNamespace) TRT_NOEXCEPT override { 92 | mNamespace = libNamespace; 93 | } 94 | 95 | const char* getPluginNamespace() const TRT_NOEXCEPT override { 96 | return mNamespace.c_str(); 97 | } 98 | 99 | private: 100 | std::string mNamespace; 101 | static PluginFieldCollection mFC; 102 | static std::vector mPluginAttributes; 103 | }; 104 | REGISTER_TENSORRT_PLUGIN(YoloPluginCreator); 105 | }; 106 | 107 | -------------------------------------------------------------------------------- /yolov5/src/calibrator.cpp: -------------------------------------------------------------------------------- 1 | #include "calibrator.h" 2 | #include "cuda_utils.h" 3 | #include "utils.h" 4 | 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | static cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) { 12 | int w, h, x, y; 13 | float r_w = input_w / (img.cols * 1.0); 14 | float r_h = input_h / (img.rows * 1.0); 15 | if (r_h > r_w) { 16 | w = input_w; 17 | h = r_w * img.rows; 18 | x = 0; 19 | y = (input_h - h) / 2; 20 | } else { 21 | w = r_h * img.cols; 22 | h = input_h; 23 | x = (input_w - w) / 2; 24 | y = 0; 25 | } 26 | cv::Mat re(h, w, CV_8UC3); 27 | cv::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR); 28 | cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(128, 128, 128)); 29 | re.copyTo(out(cv::Rect(x, y, re.cols, re.rows))); 30 | return out; 31 | } 32 | 33 | Int8EntropyCalibrator2::Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache) 34 | : batchsize_(batchsize), 35 | input_w_(input_w), 36 | input_h_(input_h), 37 | img_idx_(0), 38 | img_dir_(img_dir), 39 | calib_table_name_(calib_table_name), 40 | input_blob_name_(input_blob_name), 41 | read_cache_(read_cache) { 42 | input_count_ = 3 * input_w * input_h * batchsize; 43 | CUDA_CHECK(cudaMalloc(&device_input_, input_count_ * sizeof(float))); 44 | read_files_in_dir(img_dir, img_files_); 45 | } 46 | 47 | Int8EntropyCalibrator2::~Int8EntropyCalibrator2() { 48 | CUDA_CHECK(cudaFree(device_input_)); 49 | } 50 | 51 | int Int8EntropyCalibrator2::getBatchSize() const TRT_NOEXCEPT { 52 | return batchsize_; 53 | } 54 | 55 | bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT { 56 | if (img_idx_ + batchsize_ > (int)img_files_.size()) { 57 | return false; 58 | } 59 | 60 | std::vector input_imgs_; 61 | for (int i = img_idx_; i < img_idx_ + batchsize_; i++) { 62 | std::cout << img_files_[i] << " " << i << std::endl; 63 | cv::Mat temp = cv::imread(img_dir_ + img_files_[i]); 64 | if (temp.empty()) { 65 | std::cerr << "Fatal error: image cannot open!" << std::endl; 66 | return false; 67 | } 68 | cv::Mat pr_img = preprocess_img(temp, input_w_, input_h_); 69 | input_imgs_.push_back(pr_img); 70 | } 71 | img_idx_ += batchsize_; 72 | cv::Mat blob = cv::dnn::blobFromImages(input_imgs_, 1.0 / 255.0, cv::Size(input_w_, input_h_), cv::Scalar(0, 0, 0), true, false); 73 | 74 | CUDA_CHECK(cudaMemcpy(device_input_, blob.ptr(0), input_count_ * sizeof(float), cudaMemcpyHostToDevice)); 75 | assert(!strcmp(names[0], input_blob_name_)); 76 | bindings[0] = device_input_; 77 | return true; 78 | } 79 | 80 | const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length) TRT_NOEXCEPT { 81 | std::cout << "reading calib cache: " << calib_table_name_ << std::endl; 82 | calib_cache_.clear(); 83 | std::ifstream input(calib_table_name_, std::ios::binary); 84 | input >> std::noskipws; 85 | if (read_cache_ && input.good()) { 86 | std::copy(std::istream_iterator(input), std::istream_iterator(), std::back_inserter(calib_cache_)); 87 | } 88 | length = calib_cache_.size(); 89 | return length ? calib_cache_.data() : nullptr; 90 | } 91 | 92 | void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT { 93 | std::cout << "writing calib cache: " << calib_table_name_ << " size: " << length << std::endl; 94 | std::ofstream output(calib_table_name_, std::ios::binary); 95 | output.write(reinterpret_cast(cache), length); 96 | } 97 | 98 | -------------------------------------------------------------------------------- /yolov5/src/calibrator.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "macros.h" 4 | #include 5 | #include 6 | 7 | //! \class Int8EntropyCalibrator2 8 | //! 9 | //! \brief Implements Entropy calibrator 2. 10 | //! CalibrationAlgoType is kENTROPY_CALIBRATION_2. 11 | //! 12 | class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 { 13 | public: 14 | Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache = true); 15 | 16 | virtual ~Int8EntropyCalibrator2(); 17 | int getBatchSize() const TRT_NOEXCEPT override; 18 | bool getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT override; 19 | const void* readCalibrationCache(size_t& length) TRT_NOEXCEPT override; 20 | void writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT override; 21 | 22 | private: 23 | int batchsize_; 24 | int input_w_; 25 | int input_h_; 26 | int img_idx_; 27 | std::string img_dir_; 28 | std::vector img_files_; 29 | size_t input_count_; 30 | std::string calib_table_name_; 31 | const char* input_blob_name_; 32 | bool read_cache_; 33 | void* device_input_; 34 | std::vector calib_cache_; 35 | }; 36 | 37 | -------------------------------------------------------------------------------- /yolov5/src/config.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | /* -------------------------------------------------------- 4 | * These configs are related to tensorrt model, if these are changed, 5 | * please re-compile and re-serialize the tensorrt model. 6 | * --------------------------------------------------------*/ 7 | 8 | // For INT8, you need prepare the calibration dataset, please refer to 9 | // https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5#int8-quantization 10 | #define USE_FP16 // set USE_INT8 or USE_FP16 or USE_FP32 11 | 12 | // These are used to define input/output tensor names, 13 | // you can set them to whatever you want. 14 | const static char* kInputTensorName = "data"; 15 | const static char* kOutputTensorName = "prob"; 16 | 17 | // Detection model and Segmentation model' number of classes 18 | constexpr static int kNumClass = 80; 19 | 20 | // Classfication model's number of classes 21 | constexpr static int kClsNumClass = 1000; 22 | 23 | constexpr static int kBatchSize = 1; 24 | 25 | // Yolo's input width and height must by divisible by 32 26 | constexpr static int kInputH = 640; 27 | constexpr static int kInputW = 640; 28 | 29 | // Classfication model's input shape 30 | constexpr static int kClsInputH = 224; 31 | constexpr static int kClsInputW = 224; 32 | 33 | // Maximum number of output bounding boxes from yololayer plugin. 34 | // That is maximum number of output bounding boxes before NMS. 35 | constexpr static int kMaxNumOutputBbox = 1000; 36 | 37 | constexpr static int kNumAnchor = 3; 38 | 39 | // The bboxes whose confidence is lower than kIgnoreThresh will be ignored in yololayer plugin. 40 | constexpr static float kIgnoreThresh = 0.1f; 41 | 42 | /* -------------------------------------------------------- 43 | * These configs are NOT related to tensorrt model, if these are changed, 44 | * please re-compile, but no need to re-serialize the tensorrt model. 45 | * --------------------------------------------------------*/ 46 | 47 | // NMS overlapping thresh and final detection confidence thresh 48 | const static float kNmsThresh = 0.45f; 49 | const static float kConfThresh = 0.5f; 50 | 51 | const static int kGpuId = 0; 52 | 53 | // If your image size is larger than 4096 * 3112, please increase this value 54 | const static int kMaxInputImageSize = 4096 * 3112; 55 | 56 | -------------------------------------------------------------------------------- /yolov5/src/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef TRTX_CUDA_UTILS_H_ 2 | #define TRTX_CUDA_UTILS_H_ 3 | 4 | #include 5 | 6 | #ifndef CUDA_CHECK 7 | #define CUDA_CHECK(callstr)\ 8 | {\ 9 | cudaError_t error_code = callstr;\ 10 | if (error_code != cudaSuccess) {\ 11 | std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__;\ 12 | assert(0);\ 13 | }\ 14 | } 15 | #endif // CUDA_CHECK 16 | 17 | #endif // TRTX_CUDA_UTILS_H_ 18 | 19 | -------------------------------------------------------------------------------- /yolov5/src/macros.h: -------------------------------------------------------------------------------- 1 | #ifndef __MACROS_H 2 | #define __MACROS_H 3 | 4 | #include 5 | 6 | #ifdef API_EXPORTS 7 | #if defined(_MSC_VER) 8 | #define API __declspec(dllexport) 9 | #else 10 | #define API __attribute__((visibility("default"))) 11 | #endif 12 | #else 13 | 14 | #if defined(_MSC_VER) 15 | #define API __declspec(dllimport) 16 | #else 17 | #define API 18 | #endif 19 | #endif // API_EXPORTS 20 | 21 | #if NV_TENSORRT_MAJOR >= 8 22 | #define TRT_NOEXCEPT noexcept 23 | #define TRT_CONST_ENQUEUE const 24 | #else 25 | #define TRT_NOEXCEPT 26 | #define TRT_CONST_ENQUEUE 27 | #endif 28 | 29 | #endif // __MACROS_H 30 | -------------------------------------------------------------------------------- /yolov5/src/model.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | nvinfer1::ICudaEngine* build_det_engine(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, 7 | nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, 8 | float& gd, float& gw, std::string& wts_name); 9 | 10 | nvinfer1::ICudaEngine* build_det_p6_engine(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, 11 | nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, 12 | float& gd, float& gw, std::string& wts_name); 13 | 14 | nvinfer1::ICudaEngine* build_cls_engine(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, float& gd, float& gw, std::string& wts_name); 15 | 16 | nvinfer1::ICudaEngine* build_seg_engine(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, float& gd, float& gw, std::string& wts_name); 17 | -------------------------------------------------------------------------------- /yolov5/src/postprocess.cpp: -------------------------------------------------------------------------------- 1 | #include "postprocess.h" 2 | #include "utils.h" 3 | 4 | cv::Rect get_rect(cv::Mat& img, float bbox[4]) { 5 | float l, r, t, b; 6 | float r_w = kInputW / (img.cols * 1.0); 7 | float r_h = kInputH / (img.rows * 1.0); 8 | if (r_h > r_w) { 9 | l = bbox[0] - bbox[2] / 2.f; 10 | r = bbox[0] + bbox[2] / 2.f; 11 | t = bbox[1] - bbox[3] / 2.f - (kInputH - r_w * img.rows) / 2; 12 | b = bbox[1] + bbox[3] / 2.f - (kInputH - r_w * img.rows) / 2; 13 | l = l / r_w; 14 | r = r / r_w; 15 | t = t / r_w; 16 | b = b / r_w; 17 | } else { 18 | l = bbox[0] - bbox[2] / 2.f - (kInputW - r_h * img.cols) / 2; 19 | r = bbox[0] + bbox[2] / 2.f - (kInputW - r_h * img.cols) / 2; 20 | t = bbox[1] - bbox[3] / 2.f; 21 | b = bbox[1] + bbox[3] / 2.f; 22 | l = l / r_h; 23 | r = r / r_h; 24 | t = t / r_h; 25 | b = b / r_h; 26 | } 27 | return cv::Rect(round(l), round(t), round(r - l), round(b - t)); 28 | } 29 | 30 | static float iou(float lbox[4], float rbox[4]) { 31 | float interBox[] = { 32 | (std::max)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left 33 | (std::min)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right 34 | (std::max)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top 35 | (std::min)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom 36 | }; 37 | 38 | if (interBox[2] > interBox[3] || interBox[0] > interBox[1]) 39 | return 0.0f; 40 | 41 | float interBoxS = (interBox[1] - interBox[0])*(interBox[3] - interBox[2]); 42 | return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS); 43 | } 44 | 45 | static bool cmp(const Detection& a, const Detection& b) { 46 | return a.conf > b.conf; 47 | } 48 | 49 | void nms(std::vector& res, float* output, float conf_thresh, float nms_thresh) { 50 | int det_size = sizeof(Detection) / sizeof(float); 51 | std::map> m; 52 | for (int i = 0; i < output[0] && i < kMaxNumOutputBbox; i++) { 53 | if (output[1 + det_size * i + 4] <= conf_thresh) continue; 54 | Detection det; 55 | memcpy(&det, &output[1 + det_size * i], det_size * sizeof(float)); 56 | if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector()); 57 | m[det.class_id].push_back(det); 58 | } 59 | for (auto it = m.begin(); it != m.end(); it++) { 60 | auto& dets = it->second; 61 | std::sort(dets.begin(), dets.end(), cmp); 62 | for (size_t m = 0; m < dets.size(); ++m) { 63 | auto& item = dets[m]; 64 | res.push_back(item); 65 | for (size_t n = m + 1; n < dets.size(); ++n) { 66 | if (iou(item.bbox, dets[n].bbox) > nms_thresh) { 67 | dets.erase(dets.begin() + n); 68 | --n; 69 | } 70 | } 71 | } 72 | } 73 | } 74 | 75 | void batch_nms(std::vector>& res_batch, float *output, int batch_size, int output_size, float conf_thresh, float nms_thresh) { 76 | res_batch.resize(batch_size); 77 | for (int i = 0; i < batch_size; i++) { 78 | nms(res_batch[i], &output[i * output_size], conf_thresh, nms_thresh); 79 | } 80 | } 81 | 82 | void draw_bbox(std::vector& img_batch, std::vector>& res_batch) { 83 | for (size_t i = 0; i < img_batch.size(); i++) { 84 | auto& res = res_batch[i]; 85 | cv::Mat img = img_batch[i]; 86 | for (size_t j = 0; j < res.size(); j++) { 87 | cv::Rect r = get_rect(img, res[j].bbox); 88 | cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2); 89 | cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2); 90 | } 91 | } 92 | } 93 | 94 | static cv::Rect get_downscale_rect(float bbox[4], float scale) { 95 | float left = bbox[0] - bbox[2] / 2; 96 | float top = bbox[1] - bbox[3] / 2; 97 | float right = bbox[0] + bbox[2] / 2; 98 | float bottom = bbox[1] + bbox[3] / 2; 99 | left /= scale; 100 | top /= scale; 101 | right /= scale; 102 | bottom /= scale; 103 | return cv::Rect(round(left), round(top), round(right - left), round(bottom - top)); 104 | } 105 | 106 | std::vector process_mask(const float* proto, int proto_size, std::vector& dets) { 107 | std::vector masks; 108 | for (size_t i = 0; i < dets.size(); i++) { 109 | cv::Mat mask_mat = cv::Mat::zeros(kInputH / 4, kInputW / 4, CV_32FC1); 110 | auto r = get_downscale_rect(dets[i].bbox, 4); 111 | for (int x = r.x; x < r.x + r.width; x++) { 112 | for (int y = r.y; y < r.y + r.height; y++) { 113 | float e = 0.0f; 114 | for (int j = 0; j < 32; j++) { 115 | e += dets[i].mask[j] * proto[j * proto_size / 32 + y * mask_mat.cols + x]; 116 | } 117 | e = 1.0f / (1.0f + expf(-e)); 118 | mask_mat.at(y, x) = e; 119 | } 120 | } 121 | cv::resize(mask_mat, mask_mat, cv::Size(kInputW, kInputH)); 122 | masks.push_back(mask_mat); 123 | } 124 | return masks; 125 | } 126 | 127 | cv::Mat scale_mask(cv::Mat mask, cv::Mat img) { 128 | int x, y, w, h; 129 | float r_w = kInputW / (img.cols * 1.0); 130 | float r_h = kInputH / (img.rows * 1.0); 131 | if (r_h > r_w) { 132 | w = kInputW; 133 | h = r_w * img.rows; 134 | x = 0; 135 | y = (kInputH - h) / 2; 136 | } else { 137 | w = r_h * img.cols; 138 | h = kInputH; 139 | x = (kInputW - w) / 2; 140 | y = 0; 141 | } 142 | cv::Rect r(x, y, w, h); 143 | cv::Mat res; 144 | cv::resize(mask(r), res, img.size()); 145 | return res; 146 | } 147 | 148 | void draw_mask_bbox(cv::Mat& img, std::vector& dets, std::vector& masks, std::unordered_map& labels_map) { 149 | static std::vector colors = {0xFF3838, 0xFF9D97, 0xFF701F, 0xFFB21D, 0xCFD231, 0x48F90A, 150 | 0x92CC17, 0x3DDB86, 0x1A9334, 0x00D4BB, 0x2C99A8, 0x00C2FF, 151 | 0x344593, 0x6473FF, 0x0018EC, 0x8438FF, 0x520085, 0xCB38FF, 152 | 0xFF95C8, 0xFF37C7}; 153 | for (size_t i = 0; i < dets.size(); i++) { 154 | cv::Mat img_mask = scale_mask(masks[i], img); 155 | auto color = colors[(int)dets[i].class_id % colors.size()]; 156 | auto bgr = cv::Scalar(color & 0xFF, color >> 8 & 0xFF, color >> 16 & 0xFF); 157 | 158 | cv::Rect r = get_rect(img, dets[i].bbox); 159 | for (int x = r.x; x < r.x + r.width; x++) { 160 | for (int y = r.y; y < r.y + r.height; y++) { 161 | float val = img_mask.at(y, x); 162 | if (val <= 0.5) continue; 163 | img.at(y, x)[0] = img.at(y, x)[0] / 2 + bgr[0] / 2; 164 | img.at(y, x)[1] = img.at(y, x)[1] / 2 + bgr[1] / 2; 165 | img.at(y, x)[2] = img.at(y, x)[2] / 2 + bgr[2] / 2; 166 | } 167 | } 168 | 169 | cv::rectangle(img, r, bgr, 2); 170 | 171 | // Get the size of the text 172 | cv::Size textSize = cv::getTextSize(labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::FONT_HERSHEY_PLAIN, 1.2, 2, NULL); 173 | // Set the top left corner of the rectangle 174 | cv::Point topLeft(r.x, r.y - textSize.height); 175 | 176 | // Set the bottom right corner of the rectangle 177 | cv::Point bottomRight(r.x + textSize.width, r.y + textSize.height); 178 | 179 | // Set the thickness of the rectangle lines 180 | int lineThickness = 2; 181 | 182 | // Draw the rectangle on the image 183 | cv::rectangle(img, topLeft, bottomRight, bgr, -1); 184 | 185 | cv::putText(img, labels_map[(int)dets[i].class_id] + " " + to_string_with_precision(dets[i].conf), cv::Point(r.x, r.y + 4), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar::all(0xFF), 2); 186 | 187 | } 188 | } 189 | 190 | -------------------------------------------------------------------------------- /yolov5/src/postprocess.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "types.h" 4 | #include 5 | 6 | cv::Rect get_rect(cv::Mat& img, float bbox[4]); 7 | 8 | void nms(std::vector& res, float *output, float conf_thresh, float nms_thresh = 0.5); 9 | 10 | void batch_nms(std::vector>& batch_res, float *output, int batch_size, int output_size, float conf_thresh, float nms_thresh = 0.5); 11 | 12 | void draw_bbox(std::vector& img_batch, std::vector>& res_batch); 13 | 14 | std::vector process_mask(const float* proto, int proto_size, std::vector& dets); 15 | 16 | void draw_mask_bbox(cv::Mat& img, std::vector& dets, std::vector& masks, std::unordered_map& labels_map); 17 | -------------------------------------------------------------------------------- /yolov5/src/preprocess.cu: -------------------------------------------------------------------------------- 1 | #include "preprocess.h" 2 | #include "cuda_utils.h" 3 | 4 | static uint8_t* img_buffer_host = nullptr; 5 | static uint8_t* img_buffer_device = nullptr; 6 | 7 | struct AffineMatrix { 8 | float value[6]; 9 | }; 10 | 11 | __global__ void warpaffine_kernel( 12 | uint8_t* src, int src_line_size, int src_width, 13 | int src_height, float* dst, int dst_width, 14 | int dst_height, uint8_t const_value_st, 15 | AffineMatrix d2s, int edge) { 16 | int position = blockDim.x * blockIdx.x + threadIdx.x; 17 | if (position >= edge) return; 18 | 19 | float m_x1 = d2s.value[0]; 20 | float m_y1 = d2s.value[1]; 21 | float m_z1 = d2s.value[2]; 22 | float m_x2 = d2s.value[3]; 23 | float m_y2 = d2s.value[4]; 24 | float m_z2 = d2s.value[5]; 25 | 26 | int dx = position % dst_width; 27 | int dy = position / dst_width; 28 | float src_x = m_x1 * dx + m_y1 * dy + m_z1 + 0.5f; 29 | float src_y = m_x2 * dx + m_y2 * dy + m_z2 + 0.5f; 30 | float c0, c1, c2; 31 | 32 | if (src_x <= -1 || src_x >= src_width || src_y <= -1 || src_y >= src_height) { 33 | // out of range 34 | c0 = const_value_st; 35 | c1 = const_value_st; 36 | c2 = const_value_st; 37 | } else { 38 | int y_low = floorf(src_y); 39 | int x_low = floorf(src_x); 40 | int y_high = y_low + 1; 41 | int x_high = x_low + 1; 42 | 43 | uint8_t const_value[] = {const_value_st, const_value_st, const_value_st}; 44 | float ly = src_y - y_low; 45 | float lx = src_x - x_low; 46 | float hy = 1 - ly; 47 | float hx = 1 - lx; 48 | float w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx; 49 | uint8_t* v1 = const_value; 50 | uint8_t* v2 = const_value; 51 | uint8_t* v3 = const_value; 52 | uint8_t* v4 = const_value; 53 | 54 | if (y_low >= 0) { 55 | if (x_low >= 0) 56 | v1 = src + y_low * src_line_size + x_low * 3; 57 | 58 | if (x_high < src_width) 59 | v2 = src + y_low * src_line_size + x_high * 3; 60 | } 61 | 62 | if (y_high < src_height) { 63 | if (x_low >= 0) 64 | v3 = src + y_high * src_line_size + x_low * 3; 65 | 66 | if (x_high < src_width) 67 | v4 = src + y_high * src_line_size + x_high * 3; 68 | } 69 | 70 | c0 = w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0]; 71 | c1 = w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1]; 72 | c2 = w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2]; 73 | } 74 | 75 | // bgr to rgb 76 | float t = c2; 77 | c2 = c0; 78 | c0 = t; 79 | 80 | // normalization 81 | c0 = c0 / 255.0f; 82 | c1 = c1 / 255.0f; 83 | c2 = c2 / 255.0f; 84 | 85 | // rgbrgbrgb to rrrgggbbb 86 | int area = dst_width * dst_height; 87 | float* pdst_c0 = dst + dy * dst_width + dx; 88 | float* pdst_c1 = pdst_c0 + area; 89 | float* pdst_c2 = pdst_c1 + area; 90 | *pdst_c0 = c0; 91 | *pdst_c1 = c1; 92 | *pdst_c2 = c2; 93 | } 94 | 95 | void cuda_preprocess( 96 | uint8_t* src, int src_width, int src_height, 97 | float* dst, int dst_width, int dst_height, 98 | cudaStream_t stream) { 99 | 100 | int img_size = src_width * src_height * 3; 101 | // copy data to pinned memory 102 | memcpy(img_buffer_host, src, img_size); 103 | // copy data to device memory 104 | CUDA_CHECK(cudaMemcpyAsync(img_buffer_device, img_buffer_host, img_size, cudaMemcpyHostToDevice, stream)); 105 | 106 | AffineMatrix s2d, d2s; 107 | float scale = std::min(dst_height / (float)src_height, dst_width / (float)src_width); 108 | 109 | s2d.value[0] = scale; 110 | s2d.value[1] = 0; 111 | s2d.value[2] = -scale * src_width * 0.5 + dst_width * 0.5; 112 | s2d.value[3] = 0; 113 | s2d.value[4] = scale; 114 | s2d.value[5] = -scale * src_height * 0.5 + dst_height * 0.5; 115 | 116 | cv::Mat m2x3_s2d(2, 3, CV_32F, s2d.value); 117 | cv::Mat m2x3_d2s(2, 3, CV_32F, d2s.value); 118 | cv::invertAffineTransform(m2x3_s2d, m2x3_d2s); 119 | 120 | memcpy(d2s.value, m2x3_d2s.ptr(0), sizeof(d2s.value)); 121 | 122 | int jobs = dst_height * dst_width; 123 | int threads = 256; 124 | int blocks = ceil(jobs / (float)threads); 125 | 126 | warpaffine_kernel<<>>( 127 | img_buffer_device, src_width * 3, src_width, 128 | src_height, dst, dst_width, 129 | dst_height, 128, d2s, jobs); 130 | } 131 | 132 | void cuda_batch_preprocess(std::vector& img_batch, 133 | float* dst, int dst_width, int dst_height, 134 | cudaStream_t stream) { 135 | int dst_size = dst_width * dst_height * 3; 136 | for (size_t i = 0; i < img_batch.size(); i++) { 137 | cuda_preprocess(img_batch[i].ptr(), img_batch[i].cols, img_batch[i].rows, &dst[dst_size * i], dst_width, dst_height, stream); 138 | CUDA_CHECK(cudaStreamSynchronize(stream)); 139 | } 140 | } 141 | 142 | void cuda_preprocess_init(int max_image_size) { 143 | // prepare input data in pinned memory 144 | CUDA_CHECK(cudaMallocHost((void**)&img_buffer_host, max_image_size * 3)); 145 | // prepare input data in device memory 146 | CUDA_CHECK(cudaMalloc((void**)&img_buffer_device, max_image_size * 3)); 147 | } 148 | 149 | void cuda_preprocess_destroy() { 150 | CUDA_CHECK(cudaFree(img_buffer_device)); 151 | CUDA_CHECK(cudaFreeHost(img_buffer_host)); 152 | } 153 | 154 | -------------------------------------------------------------------------------- /yolov5/src/preprocess.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | void cuda_preprocess_init(int max_image_size); 8 | void cuda_preprocess_destroy(); 9 | void cuda_preprocess(uint8_t* src, int src_width, int src_height, 10 | float* dst, int dst_width, int dst_height, 11 | cudaStream_t stream); 12 | void cuda_batch_preprocess(std::vector& img_batch, 13 | float* dst, int dst_width, int dst_height, 14 | cudaStream_t stream); 15 | 16 | -------------------------------------------------------------------------------- /yolov5/src/types.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "config.h" 4 | 5 | struct YoloKernel { 6 | int width; 7 | int height; 8 | float anchors[kNumAnchor * 2]; 9 | }; 10 | 11 | struct alignas(float) Detection { 12 | float bbox[4]; // center_x center_y w h 13 | float conf; // bbox_conf * cls_conf 14 | float class_id; 15 | float mask[32]; 16 | }; 17 | 18 | -------------------------------------------------------------------------------- /yolov5/src/utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | static inline int read_files_in_dir(const char* p_dir_name, std::vector& file_names) { 12 | DIR *p_dir = opendir(p_dir_name); 13 | if (p_dir == nullptr) { 14 | return -1; 15 | } 16 | 17 | struct dirent* p_file = nullptr; 18 | while ((p_file = readdir(p_dir)) != nullptr) { 19 | if (strcmp(p_file->d_name, ".") != 0 && 20 | strcmp(p_file->d_name, "..") != 0) { 21 | //std::string cur_file_name(p_dir_name); 22 | //cur_file_name += "/"; 23 | //cur_file_name += p_file->d_name; 24 | std::string cur_file_name(p_file->d_name); 25 | file_names.push_back(cur_file_name); 26 | } 27 | } 28 | 29 | closedir(p_dir); 30 | return 0; 31 | } 32 | 33 | // Function to trim leading and trailing whitespace from a string 34 | static inline std::string trim_leading_whitespace(const std::string& str) { 35 | size_t first = str.find_first_not_of(' '); 36 | if (std::string::npos == first) { 37 | return str; 38 | } 39 | size_t last = str.find_last_not_of(' '); 40 | return str.substr(first, (last - first + 1)); 41 | } 42 | 43 | // Src: https://stackoverflow.com/questions/16605967 44 | static inline std::string to_string_with_precision(const float a_value, const int n = 2) { 45 | std::ostringstream out; 46 | out.precision(n); 47 | out << std::fixed << a_value; 48 | return out.str(); 49 | } 50 | 51 | static inline int read_labels(const std::string labels_filename, std::unordered_map& labels_map) { 52 | 53 | std::ifstream file(labels_filename); 54 | // Read each line of the file 55 | std::string line; 56 | int index = 0; 57 | while (std::getline(file, line)) { 58 | // Strip the line of any leading or trailing whitespace 59 | line = trim_leading_whitespace(line); 60 | 61 | // Add the stripped line to the labels_map, using the loop index as the key 62 | labels_map[index] = line; 63 | index++; 64 | } 65 | // Close the file 66 | file.close(); 67 | 68 | return 0; 69 | } 70 | 71 | -------------------------------------------------------------------------------- /yolov5/yolov5_det.cpp: -------------------------------------------------------------------------------- 1 | #include "cuda_utils.h" 2 | #include "logging.h" 3 | #include "utils.h" 4 | #include "preprocess.h" 5 | #include "postprocess.h" 6 | #include "model.h" 7 | 8 | #include 9 | #include 10 | #include 11 | 12 | using namespace nvinfer1; 13 | 14 | static Logger gLogger; 15 | const static int kOutputSize = kMaxNumOutputBbox * sizeof(Detection) / sizeof(float) + 1; 16 | 17 | bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, bool& is_p6, float& gd, float& gw, std::string& img_dir) { 18 | if (argc < 4) return false; 19 | if (std::string(argv[1]) == "-s" && (argc == 5 || argc == 7)) { 20 | wts = std::string(argv[2]); 21 | engine = std::string(argv[3]); 22 | auto net = std::string(argv[4]); 23 | if (net[0] == 'n') { 24 | gd = 0.33; 25 | gw = 0.25; 26 | } else if (net[0] == 's') { 27 | gd = 0.33; 28 | gw = 0.50; 29 | } else if (net[0] == 'm') { 30 | gd = 0.67; 31 | gw = 0.75; 32 | } else if (net[0] == 'l') { 33 | gd = 1.0; 34 | gw = 1.0; 35 | } else if (net[0] == 'x') { 36 | gd = 1.33; 37 | gw = 1.25; 38 | } else if (net[0] == 'c' && argc == 7) { 39 | gd = atof(argv[5]); 40 | gw = atof(argv[6]); 41 | } else { 42 | return false; 43 | } 44 | if (net.size() == 2 && net[1] == '6') { 45 | is_p6 = true; 46 | } 47 | } else if (std::string(argv[1]) == "-d" && argc == 4) { 48 | engine = std::string(argv[2]); 49 | img_dir = std::string(argv[3]); 50 | } else { 51 | return false; 52 | } 53 | return true; 54 | } 55 | 56 | void prepare_buffers(ICudaEngine* engine, float** gpu_input_buffer, float** gpu_output_buffer, float** cpu_output_buffer) { 57 | assert(engine->getNbBindings() == 2); 58 | // In order to bind the buffers, we need to know the names of the input and output tensors. 59 | // Note that indices are guaranteed to be less than IEngine::getNbBindings() 60 | const int inputIndex = engine->getBindingIndex(kInputTensorName); 61 | const int outputIndex = engine->getBindingIndex(kOutputTensorName); 62 | assert(inputIndex == 0); 63 | assert(outputIndex == 1); 64 | // Create GPU buffers on device 65 | CUDA_CHECK(cudaMalloc((void**)gpu_input_buffer, kBatchSize * 3 * kInputH * kInputW * sizeof(float))); 66 | CUDA_CHECK(cudaMalloc((void**)gpu_output_buffer, kBatchSize * kOutputSize * sizeof(float))); 67 | 68 | *cpu_output_buffer = new float[kBatchSize * kOutputSize]; 69 | } 70 | 71 | void infer(IExecutionContext& context, cudaStream_t& stream, void** gpu_buffers, float* output, int batchsize) { 72 | context.enqueue(batchsize, gpu_buffers, stream, nullptr); 73 | CUDA_CHECK(cudaMemcpyAsync(output, gpu_buffers[1], batchsize * kOutputSize * sizeof(float), cudaMemcpyDeviceToHost, stream)); 74 | cudaStreamSynchronize(stream); 75 | } 76 | 77 | void serialize_engine(unsigned int max_batchsize, bool& is_p6, float& gd, float& gw, std::string& wts_name, std::string& engine_name) { 78 | // Create builder 79 | IBuilder* builder = createInferBuilder(gLogger); 80 | IBuilderConfig* config = builder->createBuilderConfig(); 81 | 82 | // Create model to populate the network, then set the outputs and create an engine 83 | ICudaEngine *engine = nullptr; 84 | if (is_p6) { 85 | engine = build_det_p6_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); 86 | } else { 87 | engine = build_det_engine(max_batchsize, builder, config, DataType::kFLOAT, gd, gw, wts_name); 88 | } 89 | assert(engine != nullptr); 90 | 91 | // Serialize the engine 92 | IHostMemory* serialized_engine = engine->serialize(); 93 | assert(serialized_engine != nullptr); 94 | 95 | // Save engine to file 96 | std::ofstream p(engine_name, std::ios::binary); 97 | if (!p) { 98 | std::cerr << "Could not open plan output file" << std::endl; 99 | assert(false); 100 | } 101 | p.write(reinterpret_cast(serialized_engine->data()), serialized_engine->size()); 102 | 103 | // Close everything down 104 | engine->destroy(); 105 | builder->destroy(); 106 | config->destroy(); 107 | serialized_engine->destroy(); 108 | } 109 | 110 | void deserialize_engine(std::string& engine_name, IRuntime** runtime, ICudaEngine** engine, IExecutionContext** context) { 111 | std::ifstream file(engine_name, std::ios::binary); 112 | if (!file.good()) { 113 | std::cerr << "read " << engine_name << " error!" << std::endl; 114 | assert(false); 115 | } 116 | size_t size = 0; 117 | file.seekg(0, file.end); 118 | size = file.tellg(); 119 | file.seekg(0, file.beg); 120 | char* serialized_engine = new char[size]; 121 | assert(serialized_engine); 122 | file.read(serialized_engine, size); 123 | file.close(); 124 | 125 | *runtime = createInferRuntime(gLogger); 126 | assert(*runtime); 127 | *engine = (*runtime)->deserializeCudaEngine(serialized_engine, size); 128 | assert(*engine); 129 | *context = (*engine)->createExecutionContext(); 130 | assert(*context); 131 | delete[] serialized_engine; 132 | } 133 | 134 | int main(int argc, char** argv) { 135 | cudaSetDevice(kGpuId); 136 | 137 | std::string wts_name = ""; 138 | std::string engine_name = ""; 139 | bool is_p6 = false; 140 | float gd = 0.0f, gw = 0.0f; 141 | std::string img_dir; 142 | 143 | if (!parse_args(argc, argv, wts_name, engine_name, is_p6, gd, gw, img_dir)) { 144 | std::cerr << "arguments not right!" << std::endl; 145 | std::cerr << "./yolov5_det -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file" << std::endl; 146 | std::cerr << "./yolov5_det -d [.engine] ../images // deserialize plan file and run inference" << std::endl; 147 | return -1; 148 | } 149 | 150 | // Create a model using the API directly and serialize it to a file 151 | if (!wts_name.empty()) { 152 | serialize_engine(kBatchSize, is_p6, gd, gw, wts_name, engine_name); 153 | return 0; 154 | } 155 | 156 | // Deserialize the engine from file 157 | IRuntime* runtime = nullptr; 158 | ICudaEngine* engine = nullptr; 159 | IExecutionContext* context = nullptr; 160 | deserialize_engine(engine_name, &runtime, &engine, &context); 161 | cudaStream_t stream; 162 | CUDA_CHECK(cudaStreamCreate(&stream)); 163 | 164 | // Init CUDA preprocessing 165 | cuda_preprocess_init(kMaxInputImageSize); 166 | 167 | // Prepare cpu and gpu buffers 168 | float* gpu_buffers[2]; 169 | float* cpu_output_buffer = nullptr; 170 | prepare_buffers(engine, &gpu_buffers[0], &gpu_buffers[1], &cpu_output_buffer); 171 | 172 | // Read images from directory 173 | std::vector file_names; 174 | if (read_files_in_dir(img_dir.c_str(), file_names) < 0) { 175 | std::cerr << "read_files_in_dir failed." << std::endl; 176 | return -1; 177 | } 178 | 179 | // batch predict 180 | for (size_t i = 0; i < file_names.size(); i += kBatchSize) { 181 | // Get a batch of images 182 | std::vector img_batch; 183 | std::vector img_name_batch; 184 | for (size_t j = i; j < i + kBatchSize && j < file_names.size(); j++) { 185 | cv::Mat img = cv::imread(img_dir + "/" + file_names[j]); 186 | img_batch.push_back(img); 187 | img_name_batch.push_back(file_names[j]); 188 | } 189 | 190 | // Preprocess 191 | cuda_batch_preprocess(img_batch, gpu_buffers[0], kInputW, kInputH, stream); 192 | 193 | // Run inference 194 | auto start = std::chrono::system_clock::now(); 195 | infer(*context, stream, (void**)gpu_buffers, cpu_output_buffer, kBatchSize); 196 | auto end = std::chrono::system_clock::now(); 197 | std::cout << "inference time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; 198 | 199 | // NMS 200 | std::vector> res_batch; 201 | batch_nms(res_batch, cpu_output_buffer, img_batch.size(), kOutputSize, kConfThresh, kNmsThresh); 202 | 203 | // Draw bounding boxes 204 | draw_bbox(img_batch, res_batch); 205 | 206 | // Save images 207 | for (size_t j = 0; j < img_batch.size(); j++) { 208 | cv::imwrite("_" + img_name_batch[j], img_batch[j]); 209 | } 210 | } 211 | 212 | // Release stream and buffers 213 | cudaStreamDestroy(stream); 214 | CUDA_CHECK(cudaFree(gpu_buffers[0])); 215 | CUDA_CHECK(cudaFree(gpu_buffers[1])); 216 | delete[] cpu_output_buffer; 217 | cuda_preprocess_destroy(); 218 | // Destroy the engine 219 | context->destroy(); 220 | engine->destroy(); 221 | runtime->destroy(); 222 | 223 | // Print histogram of the output distribution 224 | // std::cout << "\nOutput:\n\n"; 225 | // for (unsigned int i = 0; i < kOutputSize; i++) { 226 | // std::cout << prob[i] << ", "; 227 | // if (i % 10 == 0) std::cout << std::endl; 228 | // } 229 | // std::cout << std::endl; 230 | 231 | return 0; 232 | } 233 | 234 | -------------------------------------------------------------------------------- /yolov5n.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/yolov5n.pt -------------------------------------------------------------------------------- /yolov5s.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mailrocketsystems/JetsonYolov5/3934892b9e1ffd06a1ef73c2be8222983eeff75d/yolov5s.pt --------------------------------------------------------------------------------