yolov5-qat

data/hyps/hyp.Objects365.yaml

+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Hyperparameters for Objects365 training
+# python train.py --weights yolov5m.pt --data Objects365.yaml --evolve
+# See Hyperparameter Evolution tutorial for details https://github.com/ultralytics/yolov5#tutorials
+
+lr0: 0.00258
+lrf: 0.17
+momentum: 0.779
+weight_decay: 0.00058
+warmup_epochs: 1.33
+warmup_momentum: 0.86
+warmup_bias_lr: 0.0711
+box: 0.0539
+cls: 0.299
+cls_pw: 0.825
+obj: 0.632
+obj_pw: 1.0
+iou_t: 0.2
+anchor_t: 3.44
+anchors: 3.2
+fl_gamma: 0.0
+hsv_h: 0.0188
+hsv_s: 0.704
+hsv_v: 0.36
+degrees: 0.0
+translate: 0.0902
+scale: 0.491
+shear: 0.0
+perspective: 0.0
+flipud: 0.0
+fliplr: 0.5
+mosaic: 1.0
+mixup: 0.0
+copy_paste: 0.0

data/hyps/hyp.VOC.yaml

+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Hyperparameters for VOC training
+# python train.py --batch 128 --weights yolov5m6.pt --data VOC.yaml --epochs 50 --img 512 --hyp hyp.scratch-med.yaml --evolve
+# See Hyperparameter Evolution tutorial for details https://github.com/ultralytics/yolov5#tutorials
+
+# YOLOv5 Hyperparameter Evolution Results
+# Best generation: 467
+# Last generation: 996
+#    metrics/precision,       metrics/recall,      metrics/mAP_0.5, metrics/mAP_0.5:0.95,         val/box_loss,         val/obj_loss,         val/cls_loss
+#              0.87729,              0.85125,              0.91286,              0.72664,            0.0076739,            0.0042529,            0.0013865
+
+lr0: 0.00334
+lrf: 0.15135
+momentum: 0.74832
+weight_decay: 0.00025
+warmup_epochs: 3.3835
+warmup_momentum: 0.59462
+warmup_bias_lr: 0.18657
+box: 0.02
+cls: 0.21638
+cls_pw: 0.5
+obj: 0.51728
+obj_pw: 0.67198
+iou_t: 0.2
+anchor_t: 3.3744
+fl_gamma: 0.0
+hsv_h: 0.01041
+hsv_s: 0.54703
+hsv_v: 0.27739
+degrees: 0.0
+translate: 0.04591
+scale: 0.75544
+shear: 0.0
+perspective: 0.0
+flipud: 0.0
+fliplr: 0.5
+mosaic: 0.85834
+mixup: 0.04266
+copy_paste: 0.0
+anchors: 3.412

data/hyps/hyp.no-augmentation.yaml

+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Hyperparameters when using Albumentations frameworks
+# python train.py --hyp hyp.no-augmentation.yaml
+# See https://github.com/ultralytics/yolov5/pull/3882 for YOLOv5 + Albumentations Usage examples
+
+lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
+lrf: 0.1 # final OneCycleLR learning rate (lr0 * lrf)
+momentum: 0.937 # SGD momentum/Adam beta1
+weight_decay: 0.0005 # optimizer weight decay 5e-4
+warmup_epochs: 3.0 # warmup epochs (fractions ok)
+warmup_momentum: 0.8 # warmup initial momentum
+warmup_bias_lr: 0.1 # warmup initial bias lr
+box: 0.05 # box loss gain
+cls: 0.3 # cls loss gain
+cls_pw: 1.0 # cls BCELoss positive_weight
+obj: 0.7 # obj loss gain (scale with pixels)
+obj_pw: 1.0 # obj BCELoss positive_weight
+iou_t: 0.20 # IoU training threshold
+anchor_t: 4.0 # anchor-multiple threshold
+# anchors: 3  # anchors per output layer (0 to ignore)
+# this parameters are all zero since we want to use albumentation framework
+fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
+hsv_h: 0 # image HSV-Hue augmentation (fraction)
+hsv_s: 0 # image HSV-Saturation augmentation (fraction)
+hsv_v: 0 # image HSV-Value augmentation (fraction)
+degrees: 0.0 # image rotation (+/- deg)
+translate: 0 # image translation (+/- fraction)
+scale: 0 # image scale (+/- gain)
+shear: 0 # image shear (+/- deg)
+perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
+flipud: 0.0 # image flip up-down (probability)
+fliplr: 0.0 # image flip left-right (probability)
+mosaic: 0.0 # image mosaic (probability)
+mixup: 0.0 # image mixup (probability)
+copy_paste: 0.0 # segment copy-paste (probability)

data/hyps/hyp.scratch-high.yaml

+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Hyperparameters for high-augmentation COCO training from scratch
+# python train.py --batch 32 --cfg yolov5m6.yaml --weights '' --data coco.yaml --img 1280 --epochs 300
+# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
+
+lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
+lrf: 0.1 # final OneCycleLR learning rate (lr0 * lrf)
+momentum: 0.937 # SGD momentum/Adam beta1
+weight_decay: 0.0005 # optimizer weight decay 5e-4
+warmup_epochs: 3.0 # warmup epochs (fractions ok)
+warmup_momentum: 0.8 # warmup initial momentum
+warmup_bias_lr: 0.1 # warmup initial bias lr
+box: 0.05 # box loss gain
+cls: 0.3 # cls loss gain
+cls_pw: 1.0 # cls BCELoss positive_weight
+obj: 0.7 # obj loss gain (scale with pixels)
+obj_pw: 1.0 # obj BCELoss positive_weight
+iou_t: 0.20 # IoU training threshold
+anchor_t: 4.0 # anchor-multiple threshold
+# anchors: 3  # anchors per output layer (0 to ignore)
+fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
+hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
+hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
+hsv_v: 0.4 # image HSV-Value augmentation (fraction)
+degrees: 0.0 # image rotation (+/- deg)
+translate: 0.1 # image translation (+/- fraction)
+scale: 0.9 # image scale (+/- gain)
+shear: 0.0 # image shear (+/- deg)
+perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
+flipud: 0.0 # image flip up-down (probability)
+fliplr: 0.5 # image flip left-right (probability)
+mosaic: 1.0 # image mosaic (probability)
+mixup: 0.1 # image mixup (probability)
+copy_paste: 0.1 # segment copy-paste (probability)

data/hyps/hyp.scratch-low.yaml

+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Hyperparameters for low-augmentation COCO training from scratch
+# python train.py --batch 64 --cfg yolov5n6.yaml --weights '' --data coco.yaml --img 640 --epochs 300 --linear
+# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
+
+lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
+lrf: 0.01 # final OneCycleLR learning rate (lr0 * lrf)
+momentum: 0.937 # SGD momentum/Adam beta1
+weight_decay: 0.0005 # optimizer weight decay 5e-4
+warmup_epochs: 3.0 # warmup epochs (fractions ok)
+warmup_momentum: 0.8 # warmup initial momentum
+warmup_bias_lr: 0.1 # warmup initial bias lr
+box: 0.05 # box loss gain
+cls: 0.5 # cls loss gain
+cls_pw: 1.0 # cls BCELoss positive_weight
+obj: 1.0 # obj loss gain (scale with pixels)
+obj_pw: 1.0 # obj BCELoss positive_weight
+iou_t: 0.20 # IoU training threshold
+anchor_t: 4.0 # anchor-multiple threshold
+# anchors: 3  # anchors per output layer (0 to ignore)
+fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
+hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
+hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
+hsv_v: 0.4 # image HSV-Value augmentation (fraction)
+degrees: 0.0 # image rotation (+/- deg)
+translate: 0.1 # image translation (+/- fraction)
+scale: 0.5 # image scale (+/- gain)
+shear: 0.0 # image shear (+/- deg)
+perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
+flipud: 0.0 # image flip up-down (probability)
+fliplr: 0.5 # image flip left-right (probability)
+mosaic: 1.0 # image mosaic (probability)
+mixup: 0.0 # image mixup (probability)
+copy_paste: 0.0 # segment copy-paste (probability)

data/hyps/hyp.scratch-med.yaml

+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Hyperparameters for medium-augmentation COCO training from scratch
+# python train.py --batch 32 --cfg yolov5m6.yaml --weights '' --data coco.yaml --img 1280 --epochs 300
+# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
+
+lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
+lrf: 0.1 # final OneCycleLR learning rate (lr0 * lrf)
+momentum: 0.937 # SGD momentum/Adam beta1
+weight_decay: 0.0005 # optimizer weight decay 5e-4
+warmup_epochs: 3.0 # warmup epochs (fractions ok)
+warmup_momentum: 0.8 # warmup initial momentum
+warmup_bias_lr: 0.1 # warmup initial bias lr
+box: 0.05 # box loss gain
+cls: 0.3 # cls loss gain
+cls_pw: 1.0 # cls BCELoss positive_weight
+obj: 0.7 # obj loss gain (scale with pixels)
+obj_pw: 1.0 # obj BCELoss positive_weight
+iou_t: 0.20 # IoU training threshold
+anchor_t: 4.0 # anchor-multiple threshold
+# anchors: 3  # anchors per output layer (0 to ignore)
+fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
+hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
+hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
+hsv_v: 0.4 # image HSV-Value augmentation (fraction)
+degrees: 0.0 # image rotation (+/- deg)
+translate: 0.1 # image translation (+/- fraction)
+scale: 0.9 # image scale (+/- gain)
+shear: 0.0 # image shear (+/- deg)
+perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
+flipud: 0.0 # image flip up-down (probability)
+fliplr: 0.5 # image flip left-right (probability)
+mosaic: 1.0 # image mosaic (probability)
+mixup: 0.1 # image mixup (probability)
+copy_paste: 0.0 # segment copy-paste (probability)

data/scripts/download_weights.sh

+#!/bin/bash
+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Download latest models from https://github.com/ultralytics/yolov5/releases
+# Example usage: bash data/scripts/download_weights.sh
+# parent
+# └── yolov5
+#     ├── yolov5s.pt  ← downloads here
+#     ├── yolov5m.pt
+#     └── ...
+
+python - <<EOF
+from utils.downloads import attempt_download
+
+p5 = list('nsmlx')  # P5 models
+p6 = [f'{x}6' for x in p5]  # P6 models
+cls = [f'{x}-cls' for x in p5]  # classification models
+seg = [f'{x}-seg' for x in p5]  # classification models
+
+for x in p5 + p6 + cls + seg:
+    attempt_download(f'weights/yolov5{x}.pt')
+
+EOF

data/scripts/get_coco.sh

+#!/bin/bash
+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Download COCO 2017 dataset http://cocodataset.org
+# Example usage: bash data/scripts/get_coco.sh
+# parent
+# ├── yolov5
+# └── datasets
+#     └── coco  ← downloads here
+
+# Arguments (optional) Usage: bash data/scripts/get_coco.sh --train --val --test --segments
+if [ "$#" -gt 0 ]; then
+  for opt in "$@"; do
+    case "${opt}" in
+    --train) train=true ;;
+    --val) val=true ;;
+    --test) test=true ;;
+    --segments) segments=true ;;
+    esac
+  done
+else
+  train=true
+  val=true
+  test=false
+  segments=false
+fi
+
+# Download/unzip labels
+d='../datasets' # unzip directory
+url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
+if [ "$segments" == "true" ]; then
+  f='coco2017labels-segments.zip' # 168 MB
+else
+  f='coco2017labels.zip' # 46 MB
+fi
+echo 'Downloading' $url$f ' ...'
+curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
+
+# Download/unzip images
+d='../datasets/coco/images' # unzip directory
+url=http://images.cocodataset.org/zips/
+if [ "$train" == "true" ]; then
+  f='train2017.zip' # 19G, 118k images
+  echo 'Downloading' $url$f '...'
+  curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
+fi
+if [ "$val" == "true" ]; then
+  f='val2017.zip' # 1G, 5k images
+  echo 'Downloading' $url$f '...'
+  curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
+fi
+if [ "$test" == "true" ]; then
+  f='test2017.zip' # 7G, 41k images (optional)
+  echo 'Downloading' $url$f '...'
+  curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
+fi
+wait # finish background tasks

data/scripts/get_coco128.sh

+#!/bin/bash
+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Download COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017)
+# Example usage: bash data/scripts/get_coco128.sh
+# parent
+# ├── yolov5
+# └── datasets
+#     └── coco128  ← downloads here
+
+# Download/unzip images and labels
+d='../datasets' # unzip directory
+url=https://github.com/ultralytics/yolov5/releases/download/v1.0/
+f='coco128.zip' # or 'coco128-segments.zip', 68 MB
+echo 'Downloading' $url$f ' ...'
+curl -L $url$f -o $f -# && unzip -q $f -d $d && rm $f &
+
+wait # finish background tasks

data/scripts/get_imagenet.sh

+#!/bin/bash
+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Download ILSVRC2012 ImageNet dataset https://image-net.org
+# Example usage: bash data/scripts/get_imagenet.sh
+# parent
+# ├── yolov5
+# └── datasets
+#     └── imagenet  ← downloads here
+
+# Arguments (optional) Usage: bash data/scripts/get_imagenet.sh --train --val
+if [ "$#" -gt 0 ]; then
+  for opt in "$@"; do
+    case "${opt}" in
+    --train) train=true ;;
+    --val) val=true ;;
+    esac
+  done
+else
+  train=true
+  val=true
+fi
+
+# Make dir
+d='../datasets/imagenet' # unzip directory
+mkdir -p $d && cd $d
+
+# Download/unzip train
+if [ "$train" == "true" ]; then
+  wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_train.tar # download 138G, 1281167 images
+  mkdir train && mv ILSVRC2012_img_train.tar train/ && cd train
+  tar -xf ILSVRC2012_img_train.tar && rm -f ILSVRC2012_img_train.tar
+  find . -name "*.tar" | while read NAME; do
+    mkdir -p "${NAME%.tar}"
+    tar -xf "${NAME}" -C "${NAME%.tar}"
+    rm -f "${NAME}"
+  done
+  cd ..
+fi
+
+# Download/unzip val
+if [ "$val" == "true" ]; then
+  wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar # download 6.3G, 50000 images
+  mkdir val && mv ILSVRC2012_img_val.tar val/ && cd val && tar -xf ILSVRC2012_img_val.tar
+  wget -qO- https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh | bash # move into subdirs
+fi
+
+# Delete corrupted image (optional: PNG under JPEG name that may cause dataloaders to fail)
+# rm train/n04266014/n04266014_10835.JPEG
+
+# TFRecords (optional)
+# wget https://raw.githubusercontent.com/tensorflow/models/master/research/slim/datasets/imagenet_lsvrc_2015_synsets.txt

data/scripts/get_imagenet10.sh

+#!/bin/bash
+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Download ILSVRC2012 ImageNet dataset https://image-net.org
+# Example usage: bash data/scripts/get_imagenet.sh
+# parent
+# ├── yolov5
+# └── datasets
+#     └── imagenet  ← downloads here
+
+# Arguments (optional) Usage: bash data/scripts/get_imagenet.sh --train --val
+if [ "$#" -gt 0 ]; then
+  for opt in "$@"; do
+    case "${opt}" in
+    --train) train=true ;;
+    --val) val=true ;;
+    esac
+  done
+else
+  train=true
+  val=true
+fi
+
+# Make dir
+d='../datasets/imagenet10' # unzip directory
+mkdir -p $d && cd $d
+
+# Download/unzip train
+wget https://github.com/ultralytics/yolov5/releases/download/v1.0/imagenet10.zip
+unzip imagenet10.zip && rm imagenet10.zip

data/scripts/get_imagenet100.sh

+#!/bin/bash
+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Download ILSVRC2012 ImageNet dataset https://image-net.org
+# Example usage: bash data/scripts/get_imagenet.sh
+# parent
+# ├── yolov5
+# └── datasets
+#     └── imagenet  ← downloads here
+
+# Arguments (optional) Usage: bash data/scripts/get_imagenet.sh --train --val
+if [ "$#" -gt 0 ]; then
+  for opt in "$@"; do
+    case "${opt}" in
+    --train) train=true ;;
+    --val) val=true ;;
+    esac
+  done
+else
+  train=true
+  val=true
+fi
+
+# Make dir
+d='../datasets/imagenet100' # unzip directory
+mkdir -p $d && cd $d
+
+# Download/unzip train
+wget https://github.com/ultralytics/yolov5/releases/download/v1.0/imagenet100.zip
+unzip imagenet100.zip && rm imagenet100.zip

data/scripts/get_imagenet1000.sh

+#!/bin/bash
+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# Download ILSVRC2012 ImageNet dataset https://image-net.org
+# Example usage: bash data/scripts/get_imagenet.sh
+# parent
+# ├── yolov5
+# └── datasets
+#     └── imagenet  ← downloads here
+
+# Arguments (optional) Usage: bash data/scripts/get_imagenet.sh --train --val
+if [ "$#" -gt 0 ]; then
+  for opt in "$@"; do
+    case "${opt}" in
+    --train) train=true ;;
+    --val) val=true ;;
+    esac
+  done
+else
+  train=true
+  val=true
+fi
+
+# Make dir
+d='../datasets/imagenet1000' # unzip directory
+mkdir -p $d && cd $d
+
+# Download/unzip train
+wget https://github.com/ultralytics/yolov5/releases/download/v1.0/imagenet1000.zip
+unzip imagenet1000.zip && rm imagenet1000.zip

data/xView.yaml

+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+# DIUx xView 2018 Challenge https://challenge.xviewdataset.org by U.S. National Geospatial-Intelligence Agency (NGA)
+# --------  DOWNLOAD DATA MANUALLY and jar xf val_images.zip to 'datasets/xView' before running train command!  --------
+# Example usage: python train.py --data xView.yaml
+# parent
+# ├── yolov5
+# └── datasets
+#     └── xView  ← downloads here (20.7 GB)
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/xView # dataset root dir
+train: images/autosplit_train.txt # train images (relative to 'path') 90% of 847 train images
+val: images/autosplit_val.txt # train images (relative to 'path') 10% of 847 train images
+
+# Classes
+names:
+  0: Fixed-wing Aircraft
+  1: Small Aircraft
+  2: Cargo Plane
+  3: Helicopter
+  4: Passenger Vehicle
+  5: Small Car
+  6: Bus
+  7: Pickup Truck
+  8: Utility Truck
+  9: Truck
+  10: Cargo Truck
+  11: Truck w/Box
+  12: Truck Tractor
+  13: Trailer
+  14: Truck w/Flatbed
+  15: Truck w/Liquid
+  16: Crane Truck
+  17: Railway Vehicle
+  18: Passenger Car
+  19: Cargo Car
+  20: Flat Car
+  21: Tank car
+  22: Locomotive
+  23: Maritime Vessel
+  24: Motorboat
+  25: Sailboat
+  26: Tugboat
+  27: Barge
+  28: Fishing Vessel
+  29: Ferry
+  30: Yacht
+  31: Container Ship
+  32: Oil Tanker
+  33: Engineering Vehicle
+  34: Tower crane
+  35: Container Crane
+  36: Reach Stacker
+  37: Straddle Carrier
+  38: Mobile Crane
+  39: Dump Truck
+  40: Haul Truck
+  41: Scraper/Tractor
+  42: Front loader/Bulldozer
+  43: Excavator
+  44: Cement Mixer
+  45: Ground Grader
+  46: Hut/Tent
+  47: Shed
+  48: Building
+  49: Aircraft Hangar
+  50: Damaged Building
+  51: Facility
+  52: Construction Site
+  53: Vehicle Lot
+  54: Helipad
+  55: Storage Tank
+  56: Shipping container lot
+  57: Shipping Container
+  58: Pylon
+  59: Tower
+
+# Download script/URL (optional) ---------------------------------------------------------------------------------------
+download: |
+  import json
+  import os
+  from pathlib import Path
+
+  import numpy as np
+  from PIL import Image
+  from tqdm import tqdm
+
+  from utils.dataloaders import autosplit
+  from utils.general import download, xyxy2xywhn
+
+
+  def convert_labels(fname=Path('xView/xView_train.geojson')):
+      # Convert xView geoJSON labels to YOLO format
+      path = fname.parent
+      with open(fname) as f:
+          print(f'Loading {fname}...')
+          data = json.load(f)
+
+      # Make dirs
+      labels = Path(path / 'labels' / 'train')
+      os.system(f'rm -rf {labels}')
+      labels.mkdir(parents=True, exist_ok=True)
+
+      # xView classes 11-94 to 0-59
+      xview_class2index = [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 0, 1, 2, -1, 3, -1, 4, 5, 6, 7, 8, -1, 9, 10, 11,
+                           12, 13, 14, 15, -1, -1, 16, 17, 18, 19, 20, 21, 22, -1, 23, 24, 25, -1, 26, 27, -1, 28, -1,
+                           29, 30, 31, 32, 33, 34, 35, 36, 37, -1, 38, 39, 40, 41, 42, 43, 44, 45, -1, -1, -1, -1, 46,
+                           47, 48, 49, -1, 50, 51, -1, 52, -1, -1, -1, 53, 54, -1, 55, -1, -1, 56, -1, 57, -1, 58, 59]
+
+      shapes = {}
+      for feature in tqdm(data['features'], desc=f'Converting {fname}'):
+          p = feature['properties']
+          if p['bounds_imcoords']:
+              id = p['image_id']
+              file = path / 'train_images' / id
+              if file.exists():  # 1395.tif missing
+                  try:
+                      box = np.array([int(num) for num in p['bounds_imcoords'].split(",")])
+                      assert box.shape[0] == 4, f'incorrect box shape {box.shape[0]}'
+                      cls = p['type_id']
+                      cls = xview_class2index[int(cls)]  # xView class to 0-60
+                      assert 59 >= cls >= 0, f'incorrect class index {cls}'
+
+                      # Write YOLO label
+                      if id not in shapes:
+                          shapes[id] = Image.open(file).size
+                      box = xyxy2xywhn(box[None].astype(np.float), w=shapes[id][0], h=shapes[id][1], clip=True)
+                      with open((labels / id).with_suffix('.txt'), 'a') as f:
+                          f.write(f"{cls} {' '.join(f'{x:.6f}' for x in box[0])}\n")  # write label.txt
+                  except Exception as e:
+                      print(f'WARNING: skipping one label for {file}: {e}')
+
+
+  # Download manually from https://challenge.xviewdataset.org
+  dir = Path(yaml['path'])  # dataset root dir
+  # urls = ['https://d307kc0mrhucc3.cloudfront.net/train_labels.zip',  # train labels
+  #         'https://d307kc0mrhucc3.cloudfront.net/train_images.zip',  # 15G, 847 train images
+  #         'https://d307kc0mrhucc3.cloudfront.net/val_images.zip']  # 5G, 282 val images (no labels)
+  # download(urls, dir=dir, delete=False)
+
+  # Convert labels
+  convert_labels(dir / 'xView_train.geojson')
+
+  # Move images
+  images = Path(dir / 'images')
+  images.mkdir(parents=True, exist_ok=True)
+  Path(dir / 'train_images').rename(dir / 'images' / 'train')
+  Path(dir / 'val_images').rename(dir / 'images' / 'val')
+
+  # Split
+  autosplit(dir / 'images' / 'train')

detect.py

+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+"""
+Run YOLOv5 detection inference on images, videos, directories, globs, YouTube, webcam, streams, etc.
+
+Usage - sources:
+    $ python detect.py --weights yolov5s.pt --source 0                               # webcam
+                                                     img.jpg                         # image
+                                                     vid.mp4                         # video
+                                                     screen                          # screenshot
+                                                     path/                           # directory
+                                                     list.txt                        # list of images
+                                                     list.streams                    # list of streams
+                                                     'path/*.jpg'                    # glob
+                                                     'https://youtu.be/LNwODJXcvt4'  # YouTube
+                                                     'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
+
+Usage - formats:
+    $ python detect.py --weights yolov5s.pt                 # PyTorch
+                                 yolov5s.torchscript        # TorchScript
+                                 yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
+                                 yolov5s_openvino_model     # OpenVINO
+                                 yolov5s.engine             # TensorRT
+                                 yolov5s.mlmodel            # CoreML (macOS-only)
+                                 yolov5s_saved_model        # TensorFlow SavedModel
+                                 yolov5s.pb                 # TensorFlow GraphDef
+                                 yolov5s.tflite             # TensorFlow Lite
+                                 yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
+                                 yolov5s_paddle_model       # PaddlePaddle
+"""
+
+import argparse
+import csv
+import os
+import platform
+import sys
+from pathlib import Path
+
+import torch
+
+FILE = Path(__file__).resolve()
+ROOT = FILE.parents[0]  # YOLOv5 root directory
+if str(ROOT) not in sys.path:
+    sys.path.append(str(ROOT))  # add ROOT to PATH
+ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
+
+from ultralytics.utils.plotting import Annotator, colors, save_one_box
+
+from models.common import DetectMultiBackend
+from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
+from utils.general import (
+    LOGGER,
+    Profile,
+    check_file,
+    check_img_size,
+    check_imshow,
+    check_requirements,
+    colorstr,
+    cv2,
+    increment_path,
+    non_max_suppression,
+    print_args,
+    scale_boxes,
+    strip_optimizer,
+    xyxy2xywh,
+)
+from utils.torch_utils import select_device, smart_inference_mode
+
+
+@smart_inference_mode()
+def run(
+    weights=ROOT / "yolov5s.pt",  # model path or triton URL
+    source=ROOT / "data/images",  # file/dir/URL/glob/screen/0(webcam)
+    data=ROOT / "data/coco128.yaml",  # dataset.yaml path
+    imgsz=(640, 640),  # inference size (height, width)
+    conf_thres=0.25,  # confidence threshold
+    iou_thres=0.45,  # NMS IOU threshold
+    max_det=1000,  # maximum detections per image
+    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    view_img=False,  # show results
+    save_txt=False,  # save results to *.txt
+    save_csv=False,  # save results in CSV format
+    save_conf=False,  # save confidences in --save-txt labels
+    save_crop=False,  # save cropped prediction boxes
+    nosave=False,  # do not save images/videos
+    classes=None,  # filter by class: --class 0, or --class 0 2 3
+    agnostic_nms=False,  # class-agnostic NMS
+    augment=False,  # augmented inference
+    visualize=False,  # visualize features
+    update=False,  # update all models
+    project=ROOT / "runs/detect",  # save results to project/name
+    name="exp",  # save results to project/name
+    exist_ok=False,  # existing project/name ok, do not increment
+    line_thickness=3,  # bounding box thickness (pixels)
+    hide_labels=False,  # hide labels
+    hide_conf=False,  # hide confidences
+    half=False,  # use FP16 half-precision inference
+    dnn=False,  # use OpenCV DNN for ONNX inference
+    vid_stride=1,  # video frame-rate stride
+):
+    source = str(source)
+    save_img = not nosave and not source.endswith(".txt")  # save inference images
+    is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
+    is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
+    webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
+    screenshot = source.lower().startswith("screen")
+    if is_url and is_file:
+        source = check_file(source)  # download
+
+    # Directories
+    save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
+    (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
+
+    # Load model
+    device = select_device(device)
+    model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
+    stride, names, pt = model.stride, model.names, model.pt
+    imgsz = check_img_size(imgsz, s=stride)  # check image size
+
+    # Dataloader
+    bs = 1  # batch_size
+    if webcam:
+        view_img = check_imshow(warn=True)
+        dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
+        bs = len(dataset)
+    elif screenshot:
+        dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
+    else:
+        dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
+    vid_path, vid_writer = [None] * bs, [None] * bs
+
+    # Run inference
+    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
+    seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))
+    for path, im, im0s, vid_cap, s in dataset:
+        with dt[0]:
+            im = torch.from_numpy(im).to(model.device)
+            im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
+            im /= 255  # 0 - 255 to 0.0 - 1.0
+            if len(im.shape) == 3:
+                im = im[None]  # expand for batch dim
+            if model.xml and im.shape[0] > 1:
+                ims = torch.chunk(im, im.shape[0], 0)
+
+        # Inference
+        with dt[1]:
+            visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
+            if model.xml and im.shape[0] > 1:
+                pred = None
+                for image in ims:
+                    if pred is None:
+                        pred = model(image, augment=augment, visualize=visualize).unsqueeze(0)
+                    else:
+                        pred = torch.cat((pred, model(image, augment=augment, visualize=visualize).unsqueeze(0)), dim=0)
+                pred = [pred, None]
+            else:
+                pred = model(im, augment=augment, visualize=visualize)
+        # NMS
+        with dt[2]:
+            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
+
+        # Second-stage classifier (optional)
+        # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
+
+        # Define the path for the CSV file
+        csv_path = save_dir / "predictions.csv"
+
+        # Create or append to the CSV file
+        def write_to_csv(image_name, prediction, confidence):
+            """Writes prediction data for an image to a CSV file, appending if the file exists."""
+            data = {"Image Name": image_name, "Prediction": prediction, "Confidence": confidence}
+            with open(csv_path, mode="a", newline="") as f:
+                writer = csv.DictWriter(f, fieldnames=data.keys())
+                if not csv_path.is_file():
+                    writer.writeheader()
+                writer.writerow(data)
+
+        # Process predictions
+        for i, det in enumerate(pred):  # per image
+            seen += 1
+            if webcam:  # batch_size >= 1
+                p, im0, frame = path[i], im0s[i].copy(), dataset.count
+                s += f"{i}: "
+            else:
+                p, im0, frame = path, im0s.copy(), getattr(dataset, "frame", 0)
+
+            p = Path(p)  # to Path
+            save_path = str(save_dir / p.name)  # im.jpg
+            txt_path = str(save_dir / "labels" / p.stem) + ("" if dataset.mode == "image" else f"_{frame}")  # im.txt
+            s += "%gx%g " % im.shape[2:]  # print string
+            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
+            imc = im0.copy() if save_crop else im0  # for save_crop
+            annotator = Annotator(im0, line_width=line_thickness, example=str(names))
+            if len(det):
+                # Rescale boxes from img_size to im0 size
+                det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
+
+                # Print results
+                for c in det[:, 5].unique():
+                    n = (det[:, 5] == c).sum()  # detections per class
+                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
+
+                # Write results
+                for *xyxy, conf, cls in reversed(det):
+                    c = int(cls)  # integer class
+                    label = names[c] if hide_conf else f"{names[c]}"
+                    confidence = float(conf)
+                    confidence_str = f"{confidence:.2f}"
+
+                    if save_csv:
+                        write_to_csv(p.name, label, confidence_str)
+
+                    if save_txt:  # Write to file
+                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
+                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
+                        with open(f"{txt_path}.txt", "a") as f:
+                            f.write(("%g " * len(line)).rstrip() % line + "\n")
+
+                    if save_img or save_crop or view_img:  # Add bbox to image
+                        c = int(cls)  # integer class
+                        label = None if hide_labels else (names[c] if hide_conf else f"{names[c]} {conf:.2f}")
+                        annotator.box_label(xyxy, label, color=colors(c, True))
+                    if save_crop:
+                        save_one_box(xyxy, imc, file=save_dir / "crops" / names[c] / f"{p.stem}.jpg", BGR=True)
+
+            # Stream results
+            im0 = annotator.result()
+            if view_img:
+                if platform.system() == "Linux" and p not in windows:
+                    windows.append(p)
+                    cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
+                    cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
+                cv2.imshow(str(p), im0)
+                cv2.waitKey(1)  # 1 millisecond
+
+            # Save results (image with detections)
+            if save_img:
+                if dataset.mode == "image":
+                    cv2.imwrite(save_path, im0)
+                else:  # 'video' or 'stream'
+                    if vid_path[i] != save_path:  # new video
+                        vid_path[i] = save_path
+                        if isinstance(vid_writer[i], cv2.VideoWriter):
+                            vid_writer[i].release()  # release previous video writer
+                        if vid_cap:  # video
+                            fps = vid_cap.get(cv2.CAP_PROP_FPS)
+                            w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+                            h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+                        else:  # stream
+                            fps, w, h = 30, im0.shape[1], im0.shape[0]
+                        save_path = str(Path(save_path).with_suffix(".mp4"))  # force *.mp4 suffix on results videos
+                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
+                    vid_writer[i].write(im0)
+
+        # Print time (inference-only)
+        LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
+
+    # Print results
+    t = tuple(x.t / seen * 1e3 for x in dt)  # speeds per image
+    LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}" % t)
+    if save_txt or save_img:
+        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ""
+        LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
+    if update:
+        strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)
+
+
+def parse_opt():
+    """Parses command-line arguments for YOLOv5 detection, setting inference options and model configurations."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model path or triton URL")
+    parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
+    parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="(optional) dataset.yaml path")
+    parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640], help="inference size h,w")
+    parser.add_argument("--conf-thres", type=float, default=0.25, help="confidence threshold")
+    parser.add_argument("--iou-thres", type=float, default=0.45, help="NMS IoU threshold")
+    parser.add_argument("--max-det", type=int, default=1000, help="maximum detections per image")
+    parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
+    parser.add_argument("--view-img", action="store_true", help="show results")
+    parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
+    parser.add_argument("--save-csv", action="store_true", help="save results in CSV format")
+    parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels")
+    parser.add_argument("--save-crop", action="store_true", help="save cropped prediction boxes")
+    parser.add_argument("--nosave", action="store_true", help="do not save images/videos")
+    parser.add_argument("--classes", nargs="+", type=int, help="filter by class: --classes 0, or --classes 0 2 3")
+    parser.add_argument("--agnostic-nms", action="store_true", help="class-agnostic NMS")
+    parser.add_argument("--augment", action="store_true", help="augmented inference")
+    parser.add_argument("--visualize", action="store_true", help="visualize features")
+    parser.add_argument("--update", action="store_true", help="update all models")
+    parser.add_argument("--project", default=ROOT / "runs/detect", help="save results to project/name")
+    parser.add_argument("--name", default="exp", help="save results to project/name")
+    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
+    parser.add_argument("--line-thickness", default=3, type=int, help="bounding box thickness (pixels)")
+    parser.add_argument("--hide-labels", default=False, action="store_true", help="hide labels")
+    parser.add_argument("--hide-conf", default=False, action="store_true", help="hide confidences")
+    parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
+    parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
+    parser.add_argument("--vid-stride", type=int, default=1, help="video frame-rate stride")
+    opt = parser.parse_args()
+    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
+    print_args(vars(opt))
+    return opt
+
+
+def main(opt):
+    """Executes YOLOv5 model inference with given options, checking requirements before running the model."""
+    check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop"))
+    run(**vars(opt))
+
+
+if __name__ == "__main__":
+    opt = parse_opt()
+    main(opt)

evaluate.py

+import sys
+import os
+
+sys.path.insert(0, os.path.abspath("."))
+pydir = os.path.dirname(__file__)
+
+import warnings
+import argparse
+from pathlib import Path
+
+import cv2
+import torch
+import torch.nn as nn
+
+import numpy as np
+import val
+from models.common import Conv
+from tqdm import tqdm
+
+# Disable all warning
+warnings.filterwarnings("ignore")
+from models.yolo import DetectionModel
+
+from utils.general import non_max_suppression, xywhn2xyxy, xywh2xyxy, scale_boxes, xyxy2xywhn
+
+from utils.metrics import ConfusionMatrix, ap_per_class
+import time
+import onnxruntime
+
+from trt_utils.trt import TrtModel
+import pycuda.driver as cuda
+
+from utils.dataloaders import LoadImagesAndLabels
+from torch.utils.data import DataLoader
+
+
+names = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 
+         10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 
+         20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee', 
+         30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle', 
+         40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 
+         51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table', 
+         61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 
+         71: 'sink', 72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'}
+
+
+def draw_boxes(model,
+               image_path,
+               mtype):
+    device = torch.device("cuda")
+    image = cv2.imread(image_path)
+    w, h, c = image.shape
+    
+    image_input = cv2.resize(image, (640, 640))
+    image_input = image_input[:, :, ::-1].transpose(2, 1, 0).astype(np.float32)[np.newaxis, ...] / 255.
+    
+    if mtype == 'ori' or mtype == 'qat':
+        image_input = torch.from_numpy(image_input).to(device)
+    
+    pred = model(image_input)
+    
+    if mtype == "trt":
+        pred = pred[-1].reshape(1, -1, 85)
+    
+    if mtype == "trt" or mtype == "onnx":
+        preds = non_max_suppression(torch.from_numpy(pred).to(torch.device("cuda")), conf_thres=0.1, iou_thres=0.65, max_det=1000, agnostic=False)
+    else:
+        preds = non_max_suppression(pred, conf_thres=0.1, iou_thres=0.65, max_det=1000, agnostic=False)
+
+    for bboxes in preds:
+        for bbox in bboxes:
+            bbox = xyxy2xywhn(bbox)
+            bbox = xywhn2xyxy(bbox, w, h)
+            bbox = bbox.cpu().numpy()
+            y1,x1,y2,x2,conf,cid = bbox
+            cv2.rectangle(image, (int(x1),int(y1)), (int(x2),int(y2)), (0, 255, 0), 2)
+            label = f'Class: {names[int(cid)]}, Confidence: {conf:.2f}'
+            cv2.putText(image, label, (int(x1), int(y1) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1, cv2.LINE_AA)
+    
+    cv2.imwrite(f"{mtype}.png", image)
+    
+    del model
+    torch.cuda.empty_cache()
+    
+
+def compute_metric(model,
+                   dataloader,
+                   mtype):
+    
+    device = torch.device("cuda")
+    iouv = torch.linspace(0.5, 0.95, 10, device=device)
+    niou = iouv.numel()
+    stats = []
+    seen = 0
+    total_model_time = 0.
+    
+    for i in range(2):
+        desc = "warmup"
+
+        if i == 1:
+            desc = "val"
+
+        progress_bar = tqdm(dataloader, total=len(dataloader), desc=desc)
+        
+        for data, targets, paths, shapes in progress_bar:
+            data = data.float() / 255.
+            
+            nb, _, height, width = data.shape
+            
+            if mtype == 'ori' or mtype == 'qat':
+                data, targets = data.to(device), targets.to(device)
+            elif mtype == 'onnx' or mtype == 'trt':
+                data, targets = data.numpy(), targets.to(device)
+            else:
+                raise NotImplemented
+            
+            start_time = time.time()
+            pred = model(data)
+            end_time = time.time()
+            
+            if i == 0:
+                continue
+            
+            total_model_time += end_time - start_time
+            
+            if mtype == 'trt':
+                pred = torch.from_numpy(pred[-1].reshape(1, -1, 85)).to(device)
+            elif mtype == "onnx":
+                pred = torch.from_numpy(pred).to(device)
+            else:
+                pass
+            
+            preds = non_max_suppression(pred, conf_thres=0.001, iou_thres=0.6, max_det=300, multi_label=True, agnostic=False)
+            targets[:, 2:] *= torch.tensor((width, height, width, height), device=device)
+            
+            for si, pred in enumerate(preds):
+                seen += 1
+                labels = targets[targets[:, 0] == si, 1:]
+                nl, npr = labels.shape[0], pred.shape[0]  # number of labels, predictions
+                path, shape = Path(paths[si]), shapes[si][0]
+                correct = torch.zeros(npr, niou, dtype=torch.bool, device=device)  # init
+
+                if npr == 0:
+                    if nl:
+                        stats.append((correct, *torch.zeros((2, 0), device=device), labels[:, 0]))
+                    continue
+
+                # Predictions
+                predn = pred.clone()
+                scale_boxes(data[si].shape[1:], predn[:, :4], shape, shapes[si][1])  # native-space pred
+
+                # Evaluate
+                if nl:
+                    tbox = xywh2xyxy(labels[:, 1:5])  # target boxes
+                    scale_boxes(data[si].shape[1:], tbox, shape, shapes[si][1])  # native-space labels
+                    labelsn = torch.cat((labels[:, 0:1], tbox), 1).to(device)  # native-space labels
+                    correct = val.process_batch(predn, labelsn, iouv)
+                stats.append((correct, pred[:, 4], pred[:, 5], labels[:, 0]))  # (correct, conf, pcls, tcls)
+                
+    # Compute metrics
+    stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)]  # to numpy
+    if len(stats) and stats[0].any():
+        tp, fp, p, r, f1, ap, ap_class = ap_per_class(*stats, names=names)
+        ap50, ap = ap[:, 0], ap.mean(1)  # AP@0.5, AP@0.5:0.95
+        mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
+    
+    s = ("%22s" + "%11s" * 6) % ("Class", "Images", "P", "R", "mAP50", "mAP50-95", "Time")
+    pf = "%22s" + "%11i" * 1 + "%11.3g" * 5  # print format
+    print(s)
+    print(pf % ("all", seen, mp, mr, map50, map, total_model_time / seen))
+
+class ONNX:
+    
+    def __init__(self, 
+                 onnx_path, 
+                 device):
+        
+        sess_options = onnxruntime.SessionOptions()
+
+        if onnxruntime.get_device() == "GPU":
+            providers = ['CUDAExecutionProvider']
+        else:
+            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        
+        sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+        
+        self.session = onnxruntime.InferenceSession(onnx_path, sess_options, providers=providers, provider_options=[{"device_id": 0}]*len(providers))
+
+        self.input_name = self.get_input_name()
+        self.output_name = self.get_output_name()
+    
+    def get_input_name(self):
+        input_name = []
+        for node in self.session.get_inputs():
+            input_name.append(node.name)
+        return input_name
+    
+    def get_output_name(self):
+        output_name = []
+        for node in self.session.get_outputs():
+            output_name.append(node.name)
+        return output_name
+
+    def get_input_feed(self, image):
+        input_feed = {}
+        for name in self.input_name:
+            input_feed[name] = image
+        
+        return input_feed
+
+    def inference(self, img):
+        input_feed = self.get_input_feed(img)
+        pred = self.session.run(None, input_feed)[0]
+        
+        return pred
+
+
+class TorchModel:
+    
+    def __init__(self, 
+                 weight,
+                 device):
+        self.device = device
+        self.model = self.load_model(weight)
+    
+    def load_model(self, weight):
+        # def load_yolov5_model(model: str, weight, device) -> Model:
+        if 'yolov5l' in weight:
+            cfg = "models/yolov5l.yaml"
+        elif 'yolov5m' in weight:
+            cfg = "models/yolov5m.yaml"
+        elif 'yolov5n' in weight:
+            cfg = "models/yolov5n.yaml"
+        elif 'yolov5s' in weight:
+            cfg = "models/yolov5s.yaml"
+        elif "yolov5x" in weight:
+            cfg = "models/yolov5x.yaml"
+        else:
+            raise NotImplementedError("Only support yolov5[l, m, n, s, x]")
+        
+        model = DetectionModel(cfg=cfg).to(self.device)
+        weight = torch.load(weight, map_location="cpu")["model"].state_dict()
+        model.load_state_dict(weight,strict=False)
+        for m in model.modules():
+            if type(m) is nn.Upsample:
+                m.recompute_scale_factor = None  # torch 1.11.0 compatibility
+            elif type(m) is Conv:
+                m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
+                
+        model.float()
+        model.eval()
+
+        with torch.no_grad():
+            model.fuse()
+        return model
+
+    @torch.no_grad()
+    def inference(self, img):
+        return self.model(img)
+        
+
+def main_eval(args):
+    dataloader =  DataLoader(LoadImagesAndLabels(args.data_list,
+                                                 img_size=640, augment=False, rect=False, cache_images=True, single_cls=False),
+                             shuffle=False, batch_size=1, collate_fn=LoadImagesAndLabels.collate_fn)
+    
+    if args.mtype == "ori":
+    # # ================ 原始模型 =======================        
+        device = torch.device(f"cuda")
+        model = TorchModel(args.weight, device).inference
+    
+    elif args.mtype=='qat':
+        # =============== qat 模型 ==========================
+        device = torch.device(f"cuda")
+        model = torch.load(args.weight, map_location="cpu")['model']
+        model.to(device)
+    
+    elif args.mtype=='onnx':
+        # =============== onnx 模型 ======================
+        device = torch.device(f"cuda")
+        model = ONNX(args.weight, 0).inference
+    
+    # ================ trt 模型 ======================
+    elif args.mtype=='trt':
+        cuda.init()
+        model = TrtModel(args.weight)
+        
+    compute_metric(model, dataloader, args.mtype)
+    del model
+    torch.cuda.empty_cache()
+
+
+def main_draw(args):
+    if args.mtype == "ori":
+    # # ================ 原始模型 =======================        
+        device = torch.device(f"cuda")
+        model = TorchModel(args.weight, device).inference
+    
+    elif args.mtype=='qat':
+        # =============== qat 模型 ==========================
+        device = torch.device(f"cuda")
+        model = torch.load(args.weight, map_location="cpu")['model']
+        model.to(device)
+    
+    elif args.mtype=='onnx':
+        # =============== onnx 模型 ======================
+        device = torch.device(f"cuda")
+        model = ONNX(args.weight, 0).inference
+    
+    # ================ trt 模型 ======================
+    elif args.mtype=='trt':
+        cuda.init()
+        model = TrtModel(args.weight)
+        
+    draw_boxes(model, args.image_path, args.mtype)
+    del model
+    torch.cuda.empty_cache()
+    
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    
+    subps = parser.add_subparsers(dest="cmd")
+    
+    eval_m = subps.add_parser("eval")
+    eval_m.add_argument("--data_list", default="/home/temp/coco2017/val2017.txt", type=str, help="测试数据列表文件路径")
+    eval_m.add_argument("--weight", required=True, type=str, help="权重路径")
+    eval_m.add_argument("--mtype", required=True, type=str, choices=['onnx', 'qat', 'trt', 'ori'],
+                        help="模型类型，ori原始模型，qat带有量化节点的模型，onnx模型，tensorrt模型.")
+    
+    draw = subps.add_parser("draw")
+    draw.add_argument("--weight", required=True, type=str, help="输入trt权重路径")
+    draw.add_argument("--image_path", required=True, type=str, help="待检测图片")
+    draw.add_argument("--mtype", required=True, type=str)
+    
+    args = parser.parse_args()
+    
+    if args.cmd == "eval":
+        print(args)
+        main_eval(args)
+    elif args.cmd == "draw":
+        print(args)
+        main_draw(args)

export.py

+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+"""
+Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit
+
+Format                      | `export.py --include`         | Model
+---                         | ---                           | ---
+PyTorch                     | -                             | yolov5s.pt
+TorchScript                 | `torchscript`                 | yolov5s.torchscript
+ONNX                        | `onnx`                        | yolov5s.onnx
+OpenVINO                    | `openvino`                    | yolov5s_openvino_model/
+TensorRT                    | `engine`                      | yolov5s.engine
+CoreML                      | `coreml`                      | yolov5s.mlmodel
+TensorFlow SavedModel       | `saved_model`                 | yolov5s_saved_model/
+TensorFlow GraphDef         | `pb`                          | yolov5s.pb
+TensorFlow Lite             | `tflite`                      | yolov5s.tflite
+TensorFlow Edge TPU         | `edgetpu`                     | yolov5s_edgetpu.tflite
+TensorFlow.js               | `tfjs`                        | yolov5s_web_model/
+PaddlePaddle                | `paddle`                      | yolov5s_paddle_model/
+
+Requirements:
+    $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu  # CPU
+    $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow  # GPU
+
+Usage:
+    $ python export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...
+
+Inference:
+    $ python detect.py --weights yolov5s.pt                 # PyTorch
+                                 yolov5s.torchscript        # TorchScript
+                                 yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
+                                 yolov5s_openvino_model     # OpenVINO
+                                 yolov5s.engine             # TensorRT
+                                 yolov5s.mlmodel            # CoreML (macOS-only)
+                                 yolov5s_saved_model        # TensorFlow SavedModel
+                                 yolov5s.pb                 # TensorFlow GraphDef
+                                 yolov5s.tflite             # TensorFlow Lite
+                                 yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
+                                 yolov5s_paddle_model       # PaddlePaddle
+
+TensorFlow.js:
+    $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
+    $ npm install
+    $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
+    $ npm start
+"""
+
+import argparse
+import contextlib
+import json
+import os
+import platform
+import re
+import subprocess
+import sys
+import time
+import warnings
+from pathlib import Path
+
+import pandas as pd
+import torch
+from torch.utils.mobile_optimizer import optimize_for_mobile
+
+FILE = Path(__file__).resolve()
+ROOT = FILE.parents[0]  # YOLOv5 root directory
+if str(ROOT) not in sys.path:
+    sys.path.append(str(ROOT))  # add ROOT to PATH
+if platform.system() != "Windows":
+    ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
+
+from models.experimental import attempt_load
+from models.yolo import ClassificationModel, Detect, DetectionModel, SegmentationModel
+from utils.dataloaders import LoadImages
+from utils.general import (
+    LOGGER,
+    Profile,
+    check_dataset,
+    check_img_size,
+    check_requirements,
+    check_version,
+    check_yaml,
+    colorstr,
+    file_size,
+    get_default_args,
+    print_args,
+    url2file,
+    yaml_save,
+)
+from utils.torch_utils import select_device, smart_inference_mode
+
+MACOS = platform.system() == "Darwin"  # macOS environment
+
+
+class iOSModel(torch.nn.Module):
+    def __init__(self, model, im):
+        """Initializes an iOS compatible model with normalization based on image dimensions."""
+        super().__init__()
+        b, c, h, w = im.shape  # batch, channel, height, width
+        self.model = model
+        self.nc = model.nc  # number of classes
+        if w == h:
+            self.normalize = 1.0 / w
+        else:
+            self.normalize = torch.tensor([1.0 / w, 1.0 / h, 1.0 / w, 1.0 / h])  # broadcast (slower, smaller)
+            # np = model(im)[0].shape[1]  # number of points
+            # self.normalize = torch.tensor([1. / w, 1. / h, 1. / w, 1. / h]).expand(np, 4)  # explicit (faster, larger)
+
+    def forward(self, x):
+        """Runs forward pass on the input tensor, returning class confidences and normalized coordinates."""
+        xywh, conf, cls = self.model(x)[0].squeeze().split((4, 1, self.nc), 1)
+        return cls * conf, xywh * self.normalize  # confidence (3780, 80), coordinates (3780, 4)
+
+
+def export_formats():
+    """Returns a DataFrame of supported YOLOv5 model export formats and their properties."""
+    x = [
+        ["PyTorch", "-", ".pt", True, True],
+        ["TorchScript", "torchscript", ".torchscript", True, True],
+        ["ONNX", "onnx", ".onnx", True, True],
+        ["OpenVINO", "openvino", "_openvino_model", True, False],
+        ["TensorRT", "engine", ".engine", False, True],
+        ["CoreML", "coreml", ".mlmodel", True, False],
+        ["TensorFlow SavedModel", "saved_model", "_saved_model", True, True],
+        ["TensorFlow GraphDef", "pb", ".pb", True, True],
+        ["TensorFlow Lite", "tflite", ".tflite", True, False],
+        ["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", False, False],
+        ["TensorFlow.js", "tfjs", "_web_model", False, False],
+        ["PaddlePaddle", "paddle", "_paddle_model", True, True],
+    ]
+    return pd.DataFrame(x, columns=["Format", "Argument", "Suffix", "CPU", "GPU"])
+
+
+def try_export(inner_func):
+    """Decorator @try_export for YOLOv5 model export functions that logs success/failure, time taken, and file size."""
+    inner_args = get_default_args(inner_func)
+
+    def outer_func(*args, **kwargs):
+        prefix = inner_args["prefix"]
+        try:
+            with Profile() as dt:
+                f, model = inner_func(*args, **kwargs)
+            LOGGER.info(f"{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)")
+            return f, model
+        except Exception as e:
+            LOGGER.info(f"{prefix} export failure ❌ {dt.t:.1f}s: {e}")
+            return None, None
+
+    return outer_func
+
+
+@try_export
+def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")):
+    """Exports YOLOv5 model to TorchScript format, optionally optimized for mobile, with image shape and stride
+    metadata.
+    """
+    LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...")
+    f = file.with_suffix(".torchscript")
+
+    ts = torch.jit.trace(model, im, strict=False)
+    d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
+    extra_files = {"config.txt": json.dumps(d)}  # torch._C.ExtraFilesMap()
+    if optimize:  # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
+        optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
+    else:
+        ts.save(str(f), _extra_files=extra_files)
+    return f, None
+
+
+@try_export
+def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX:")):
+    """Exports a YOLOv5 model to ONNX format with dynamic axes and optional simplification."""
+    check_requirements("onnx>=1.12.0")
+    import onnx
+
+    LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__}...")
+    f = str(file.with_suffix(".onnx"))
+
+    output_names = ["output0", "output1"] if isinstance(model, SegmentationModel) else ["output0"]
+    if dynamic:
+        dynamic = {"images": {0: "batch", 2: "height", 3: "width"}}  # shape(1,3,640,640)
+        if isinstance(model, SegmentationModel):
+            dynamic["output0"] = {0: "batch", 1: "anchors"}  # shape(1,25200,85)
+            dynamic["output1"] = {0: "batch", 2: "mask_height", 3: "mask_width"}  # shape(1,32,160,160)
+        elif isinstance(model, DetectionModel):
+            dynamic["output0"] = {0: "batch", 1: "anchors"}  # shape(1,25200,85)
+
+    torch.onnx.export(
+        model.cpu() if dynamic else model,  # --dynamic only compatible with cpu
+        im.cpu() if dynamic else im,
+        f,
+        verbose=False,
+        opset_version=opset,
+        do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
+        input_names=["images"],
+        output_names=output_names,
+        dynamic_axes=dynamic or None,
+    )
+
+    # Checks
+    model_onnx = onnx.load(f)  # load onnx model
+    onnx.checker.check_model(model_onnx)  # check onnx model
+
+    # Metadata
+    d = {"stride": int(max(model.stride)), "names": model.names}
+    for k, v in d.items():
+        meta = model_onnx.metadata_props.add()
+        meta.key, meta.value = k, str(v)
+    onnx.save(model_onnx, f)
+
+    # Simplify
+    if simplify:
+        try:
+            cuda = torch.cuda.is_available()
+            check_requirements(("onnxruntime-gpu" if cuda else "onnxruntime", "onnx-simplifier>=0.4.1"))
+            import onnxsim
+
+            LOGGER.info(f"{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...")
+            model_onnx, check = onnxsim.simplify(model_onnx)
+            assert check, "assert check failed"
+            onnx.save(model_onnx, f)
+        except Exception as e:
+            LOGGER.info(f"{prefix} simplifier failure: {e}")
+    return f, model_onnx
+
+
+@try_export
+def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")):
+    # YOLOv5 OpenVINO export
+    check_requirements("openvino-dev>=2023.0")  # requires openvino-dev: https://pypi.org/project/openvino-dev/
+    import openvino.runtime as ov  # noqa
+    from openvino.tools import mo  # noqa
+
+    LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...")
+    f = str(file).replace(file.suffix, f"_{'int8_' if int8 else ''}openvino_model{os.sep}")
+    f_onnx = file.with_suffix(".onnx")
+    f_ov = str(Path(f) / file.with_suffix(".xml").name)
+
+    ov_model = mo.convert_model(f_onnx, model_name=file.stem, framework="onnx", compress_to_fp16=half)  # export
+
+    if int8:
+        check_requirements("nncf>=2.5.0")  # requires at least version 2.5.0 to use the post-training quantization
+        import nncf
+        import numpy as np
+
+        from utils.dataloaders import create_dataloader
+
+        def gen_dataloader(yaml_path, task="train", imgsz=640, workers=4):
+            data_yaml = check_yaml(yaml_path)
+            data = check_dataset(data_yaml)
+            dataloader = create_dataloader(
+                data[task], imgsz=imgsz, batch_size=1, stride=32, pad=0.5, single_cls=False, rect=False, workers=workers
+            )[0]
+            return dataloader
+
+        # noqa: F811
+
+        def transform_fn(data_item):
+            """
+            Quantization transform function.
+
+            Extracts and preprocess input data from dataloader item for quantization.
+            Parameters:
+               data_item: Tuple with data item produced by DataLoader during iteration
+            Returns:
+                input_tensor: Input data for quantization
+            """
+            assert data_item[0].dtype == torch.uint8, "input image must be uint8 for the quantization preprocessing"
+
+            img = data_item[0].numpy().astype(np.float32)  # uint8 to fp16/32
+            img /= 255.0  # 0 - 255 to 0.0 - 1.0
+            return np.expand_dims(img, 0) if img.ndim == 3 else img
+
+        ds = gen_dataloader(data)
+        quantization_dataset = nncf.Dataset(ds, transform_fn)
+        ov_model = nncf.quantize(ov_model, quantization_dataset, preset=nncf.QuantizationPreset.MIXED)
+
+    ov.serialize(ov_model, f_ov)  # save
+    yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata)  # add metadata.yaml
+    return f, None
+
+
+@try_export
+def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
+    """Exports a YOLOv5 model to PaddlePaddle format using X2Paddle, saving to `save_dir` and adding a metadata.yaml
+    file.
+    """
+    check_requirements(("paddlepaddle", "x2paddle"))
+    import x2paddle
+    from x2paddle.convert import pytorch2paddle
+
+    LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...")
+    f = str(file).replace(".pt", f"_paddle_model{os.sep}")
+
+    pytorch2paddle(module=model, save_dir=f, jit_type="trace", input_examples=[im])  # export
+    yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata)  # add metadata.yaml
+    return f, None
+
+
+@try_export
+def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")):
+    """Exports YOLOv5 model to CoreML format with optional NMS, INT8, and FP16 support; requires coremltools."""
+    check_requirements("coremltools")
+    import coremltools as ct
+
+    LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
+    f = file.with_suffix(".mlmodel")
+
+    if nms:
+        model = iOSModel(model, im)
+    ts = torch.jit.trace(model, im, strict=False)  # TorchScript model
+    ct_model = ct.convert(ts, inputs=[ct.ImageType("image", shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
+    bits, mode = (8, "kmeans_lut") if int8 else (16, "linear") if half else (32, None)
+    if bits < 32:
+        if MACOS:  # quantization only supported on macOS
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=DeprecationWarning)  # suppress numpy==1.20 float warning
+                ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
+        else:
+            print(f"{prefix} quantization only supported on macOS, skipping...")
+    ct_model.save(f)
+    return f, ct_model
+
+
+@try_export
+def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr("TensorRT:")):
+    """
+    Exports a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0.
+
+    https://developer.nvidia.com/tensorrt
+    """
+    assert im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. `python export.py --device 0`"
+    try:
+        import tensorrt as trt
+    except Exception:
+        if platform.system() == "Linux":
+            check_requirements("nvidia-tensorrt", cmds="-U --index-url https://pypi.ngc.nvidia.com")
+        import tensorrt as trt
+
+    if trt.__version__[0] == "7":  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
+        grid = model.model[-1].anchor_grid
+        model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
+        export_onnx(model, im, file, 12, dynamic, simplify)  # opset 12
+        model.model[-1].anchor_grid = grid
+    else:  # TensorRT >= 8
+        check_version(trt.__version__, "8.0.0", hard=True)  # require tensorrt>=8.0.0
+        export_onnx(model, im, file, 12, dynamic, simplify)  # opset 12
+    onnx = file.with_suffix(".onnx")
+
+    LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...")
+    assert onnx.exists(), f"failed to export ONNX file: {onnx}"
+    f = file.with_suffix(".engine")  # TensorRT engine file
+    logger = trt.Logger(trt.Logger.INFO)
+    if verbose:
+        logger.min_severity = trt.Logger.Severity.VERBOSE
+
+    builder = trt.Builder(logger)
+    config = builder.create_builder_config()
+    config.max_workspace_size = workspace * 1 << 30
+    # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30)  # fix TRT 8.4 deprecation notice
+
+    flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+    network = builder.create_network(flag)
+    parser = trt.OnnxParser(network, logger)
+    if not parser.parse_from_file(str(onnx)):
+        raise RuntimeError(f"failed to load ONNX file: {onnx}")
+
+    inputs = [network.get_input(i) for i in range(network.num_inputs)]
+    outputs = [network.get_output(i) for i in range(network.num_outputs)]
+    for inp in inputs:
+        LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}')
+    for out in outputs:
+        LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}')
+
+    if dynamic:
+        if im.shape[0] <= 1:
+            LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument")
+        profile = builder.create_optimization_profile()
+        for inp in inputs:
+            profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)
+        config.add_optimization_profile(profile)
+
+    LOGGER.info(f"{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}")
+    if builder.platform_has_fast_fp16 and half:
+        config.set_flag(trt.BuilderFlag.FP16)
+    with builder.build_engine(network, config) as engine, open(f, "wb") as t:
+        t.write(engine.serialize())
+    return f, None
+
+
+@try_export
+def export_saved_model(
+    model,
+    im,
+    file,
+    dynamic,
+    tf_nms=False,
+    agnostic_nms=False,
+    topk_per_class=100,
+    topk_all=100,
+    iou_thres=0.45,
+    conf_thres=0.25,
+    keras=False,
+    prefix=colorstr("TensorFlow SavedModel:"),
+):
+    # YOLOv5 TensorFlow SavedModel export
+    try:
+        import tensorflow as tf
+    except Exception:
+        check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}")
+        import tensorflow as tf
+    from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
+
+    from models.tf import TFModel
+
+    LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
+    if tf.__version__ > "2.13.1":
+        helper_url = "https://github.com/ultralytics/yolov5/issues/12489"
+        LOGGER.info(
+            f"WARNING ⚠️ using Tensorflow {tf.__version__} > 2.13.1 might cause issue when exporting the model to tflite {helper_url}"
+        )  # handling issue https://github.com/ultralytics/yolov5/issues/12489
+    f = str(file).replace(".pt", "_saved_model")
+    batch_size, ch, *imgsz = list(im.shape)  # BCHW
+
+    tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
+    im = tf.zeros((batch_size, *imgsz, ch))  # BHWC order for TensorFlow
+    _ = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
+    inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if dynamic else batch_size)
+    outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
+    keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)
+    keras_model.trainable = False
+    keras_model.summary()
+    if keras:
+        keras_model.save(f, save_format="tf")
+    else:
+        spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)
+        m = tf.function(lambda x: keras_model(x))  # full model
+        m = m.get_concrete_function(spec)
+        frozen_func = convert_variables_to_constants_v2(m)
+        tfm = tf.Module()
+        tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x), [spec])
+        tfm.__call__(im)
+        tf.saved_model.save(
+            tfm,
+            f,
+            options=tf.saved_model.SaveOptions(experimental_custom_gradients=False)
+            if check_version(tf.__version__, "2.6")
+            else tf.saved_model.SaveOptions(),
+        )
+    return f, keras_model
+
+
+@try_export
+def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")):
+    """Exports YOLOv5 model to TensorFlow GraphDef *.pb format; see https://github.com/leimao/Frozen_Graph_TensorFlow for details."""
+    import tensorflow as tf
+    from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
+
+    LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
+    f = file.with_suffix(".pb")
+
+    m = tf.function(lambda x: keras_model(x))  # full model
+    m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
+    frozen_func = convert_variables_to_constants_v2(m)
+    frozen_func.graph.as_graph_def()
+    tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
+    return f, None
+
+
+@try_export
+def export_tflite(
+    keras_model, im, file, int8, per_tensor, data, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:")
+):
+    # YOLOv5 TensorFlow Lite export
+    import tensorflow as tf
+
+    LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
+    batch_size, ch, *imgsz = list(im.shape)  # BCHW
+    f = str(file).replace(".pt", "-fp16.tflite")
+
+    converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
+    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
+    converter.target_spec.supported_types = [tf.float16]
+    converter.optimizations = [tf.lite.Optimize.DEFAULT]
+    if int8:
+        from models.tf import representative_dataset_gen
+
+        dataset = LoadImages(check_dataset(check_yaml(data))["train"], img_size=imgsz, auto=False)
+        converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)
+        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
+        converter.target_spec.supported_types = []
+        converter.inference_input_type = tf.uint8  # or tf.int8
+        converter.inference_output_type = tf.uint8  # or tf.int8
+        converter.experimental_new_quantizer = True
+        if per_tensor:
+            converter._experimental_disable_per_channel = True
+        f = str(file).replace(".pt", "-int8.tflite")
+    if nms or agnostic_nms:
+        converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS)
+
+    tflite_model = converter.convert()
+    open(f, "wb").write(tflite_model)
+    return f, None
+
+
+@try_export
+def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
+    """
+    Exports a YOLOv5 model to Edge TPU compatible TFLite format; requires Linux and Edge TPU compiler.
+
+    https://coral.ai/docs/edgetpu/models-intro/
+    """
+    cmd = "edgetpu_compiler --version"
+    help_url = "https://coral.ai/docs/edgetpu/compiler/"
+    assert platform.system() == "Linux", f"export only supported on Linux. See {help_url}"
+    if subprocess.run(f"{cmd} > /dev/null 2>&1", shell=True).returncode != 0:
+        LOGGER.info(f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}")
+        sudo = subprocess.run("sudo --version >/dev/null", shell=True).returncode == 0  # sudo installed on system
+        for c in (
+            "curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -",
+            'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',
+            "sudo apt-get update",
+            "sudo apt-get install edgetpu-compiler",
+        ):
+            subprocess.run(c if sudo else c.replace("sudo ", ""), shell=True, check=True)
+    ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]
+
+    LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
+    f = str(file).replace(".pt", "-int8_edgetpu.tflite")  # Edge TPU model
+    f_tfl = str(file).replace(".pt", "-int8.tflite")  # TFLite model
+
+    subprocess.run(
+        [
+            "edgetpu_compiler",
+            "-s",
+            "-d",
+            "-k",
+            "10",
+            "--out_dir",
+            str(file.parent),
+            f_tfl,
+        ],
+        check=True,
+    )
+    return f, None
+
+
+@try_export
+def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
+    """Exports a YOLOv5 model to TensorFlow.js format, optionally with uint8 quantization."""
+    check_requirements("tensorflowjs")
+    import tensorflowjs as tfjs
+
+    LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
+    f = str(file).replace(".pt", "_web_model")  # js dir
+    f_pb = file.with_suffix(".pb")  # *.pb path
+    f_json = f"{f}/model.json"  # *.json path
+
+    args = [
+        "tensorflowjs_converter",
+        "--input_format=tf_frozen_model",
+        "--quantize_uint8" if int8 else "",
+        "--output_node_names=Identity,Identity_1,Identity_2,Identity_3",
+        str(f_pb),
+        f,
+    ]
+    subprocess.run([arg for arg in args if arg], check=True)
+
+    json = Path(f_json).read_text()
+    with open(f_json, "w") as j:  # sort JSON Identity_* in ascending order
+        subst = re.sub(
+            r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
+            r'"Identity.?.?": {"name": "Identity.?.?"}, '
+            r'"Identity.?.?": {"name": "Identity.?.?"}, '
+            r'"Identity.?.?": {"name": "Identity.?.?"}}}',
+            r'{"outputs": {"Identity": {"name": "Identity"}, '
+            r'"Identity_1": {"name": "Identity_1"}, '
+            r'"Identity_2": {"name": "Identity_2"}, '
+            r'"Identity_3": {"name": "Identity_3"}}}',
+            json,
+        )
+        j.write(subst)
+    return f, None
+
+
+def add_tflite_metadata(file, metadata, num_outputs):
+    """
+    Adds TFLite metadata to a model file, supporting multiple outputs, as specified by TensorFlow guidelines.
+
+    https://www.tensorflow.org/lite/models/convert/metadata
+    """
+    with contextlib.suppress(ImportError):
+        # check_requirements('tflite_support')
+        from tflite_support import flatbuffers
+        from tflite_support import metadata as _metadata
+        from tflite_support import metadata_schema_py_generated as _metadata_fb
+
+        tmp_file = Path("/tmp/meta.txt")
+        with open(tmp_file, "w") as meta_f:
+            meta_f.write(str(metadata))
+
+        model_meta = _metadata_fb.ModelMetadataT()
+        label_file = _metadata_fb.AssociatedFileT()
+        label_file.name = tmp_file.name
+        model_meta.associatedFiles = [label_file]
+
+        subgraph = _metadata_fb.SubGraphMetadataT()
+        subgraph.inputTensorMetadata = [_metadata_fb.TensorMetadataT()]
+        subgraph.outputTensorMetadata = [_metadata_fb.TensorMetadataT()] * num_outputs
+        model_meta.subgraphMetadata = [subgraph]
+
+        b = flatbuffers.Builder(0)
+        b.Finish(model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
+        metadata_buf = b.Output()
+
+        populator = _metadata.MetadataPopulator.with_model_file(file)
+        populator.load_metadata_buffer(metadata_buf)
+        populator.load_associated_files([str(tmp_file)])
+        populator.populate()
+        tmp_file.unlink()
+
+
+def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")):
+    """Converts a PyTorch YOLOv5 model to CoreML format with NMS, handling different input/output shapes and saving the
+    model.
+    """
+    import coremltools as ct
+    from PIL import Image
+
+    print(f"{prefix} starting pipeline with coremltools {ct.__version__}...")
+    batch_size, ch, h, w = list(im.shape)  # BCHW
+    t = time.time()
+
+    # YOLOv5 Output shapes
+    spec = model.get_spec()
+    out0, out1 = iter(spec.description.output)
+    if platform.system() == "Darwin":
+        img = Image.new("RGB", (w, h))  # img(192 width, 320 height)
+        # img = torch.zeros((*opt.img_size, 3)).numpy()  # img size(320,192,3) iDetection
+        out = model.predict({"image": img})
+        out0_shape, out1_shape = out[out0.name].shape, out[out1.name].shape
+    else:  # linux and windows can not run model.predict(), get sizes from pytorch output y
+        s = tuple(y[0].shape)
+        out0_shape, out1_shape = (s[1], s[2] - 5), (s[1], 4)  # (3780, 80), (3780, 4)
+
+    # Checks
+    nx, ny = spec.description.input[0].type.imageType.width, spec.description.input[0].type.imageType.height
+    na, nc = out0_shape
+    # na, nc = out0.type.multiArrayType.shape  # number anchors, classes
+    assert len(names) == nc, f"{len(names)} names found for nc={nc}"  # check
+
+    # Define output shapes (missing)
+    out0.type.multiArrayType.shape[:] = out0_shape  # (3780, 80)
+    out1.type.multiArrayType.shape[:] = out1_shape  # (3780, 4)
+    # spec.neuralNetwork.preprocessing[0].featureName = '0'
+
+    # Flexible input shapes
+    # from coremltools.models.neural_network import flexible_shape_utils
+    # s = [] # shapes
+    # s.append(flexible_shape_utils.NeuralNetworkImageSize(320, 192))
+    # s.append(flexible_shape_utils.NeuralNetworkImageSize(640, 384))  # (height, width)
+    # flexible_shape_utils.add_enumerated_image_sizes(spec, feature_name='image', sizes=s)
+    # r = flexible_shape_utils.NeuralNetworkImageSizeRange()  # shape ranges
+    # r.add_height_range((192, 640))
+    # r.add_width_range((192, 640))
+    # flexible_shape_utils.update_image_size_range(spec, feature_name='image', size_range=r)
+
+    # Print
+    print(spec.description)
+
+    # Model from spec
+    model = ct.models.MLModel(spec)
+
+    # 3. Create NMS protobuf
+    nms_spec = ct.proto.Model_pb2.Model()
+    nms_spec.specificationVersion = 5
+    for i in range(2):
+        decoder_output = model._spec.description.output[i].SerializeToString()
+        nms_spec.description.input.add()
+        nms_spec.description.input[i].ParseFromString(decoder_output)
+        nms_spec.description.output.add()
+        nms_spec.description.output[i].ParseFromString(decoder_output)
+
+    nms_spec.description.output[0].name = "confidence"
+    nms_spec.description.output[1].name = "coordinates"
+
+    output_sizes = [nc, 4]
+    for i in range(2):
+        ma_type = nms_spec.description.output[i].type.multiArrayType
+        ma_type.shapeRange.sizeRanges.add()
+        ma_type.shapeRange.sizeRanges[0].lowerBound = 0
+        ma_type.shapeRange.sizeRanges[0].upperBound = -1
+        ma_type.shapeRange.sizeRanges.add()
+        ma_type.shapeRange.sizeRanges[1].lowerBound = output_sizes[i]
+        ma_type.shapeRange.sizeRanges[1].upperBound = output_sizes[i]
+        del ma_type.shape[:]
+
+    nms = nms_spec.nonMaximumSuppression
+    nms.confidenceInputFeatureName = out0.name  # 1x507x80
+    nms.coordinatesInputFeatureName = out1.name  # 1x507x4
+    nms.confidenceOutputFeatureName = "confidence"
+    nms.coordinatesOutputFeatureName = "coordinates"
+    nms.iouThresholdInputFeatureName = "iouThreshold"
+    nms.confidenceThresholdInputFeatureName = "confidenceThreshold"
+    nms.iouThreshold = 0.45
+    nms.confidenceThreshold = 0.25
+    nms.pickTop.perClass = True
+    nms.stringClassLabels.vector.extend(names.values())
+    nms_model = ct.models.MLModel(nms_spec)
+
+    # 4. Pipeline models together
+    pipeline = ct.models.pipeline.Pipeline(
+        input_features=[
+            ("image", ct.models.datatypes.Array(3, ny, nx)),
+            ("iouThreshold", ct.models.datatypes.Double()),
+            ("confidenceThreshold", ct.models.datatypes.Double()),
+        ],
+        output_features=["confidence", "coordinates"],
+    )
+    pipeline.add_model(model)
+    pipeline.add_model(nms_model)
+
+    # Correct datatypes
+    pipeline.spec.description.input[0].ParseFromString(model._spec.description.input[0].SerializeToString())
+    pipeline.spec.description.output[0].ParseFromString(nms_model._spec.description.output[0].SerializeToString())
+    pipeline.spec.description.output[1].ParseFromString(nms_model._spec.description.output[1].SerializeToString())
+
+    # Update metadata
+    pipeline.spec.specificationVersion = 5
+    pipeline.spec.description.metadata.versionString = "https://github.com/ultralytics/yolov5"
+    pipeline.spec.description.metadata.shortDescription = "https://github.com/ultralytics/yolov5"
+    pipeline.spec.description.metadata.author = "glenn.jocher@ultralytics.com"
+    pipeline.spec.description.metadata.license = "https://github.com/ultralytics/yolov5/blob/master/LICENSE"
+    pipeline.spec.description.metadata.userDefined.update(
+        {
+            "classes": ",".join(names.values()),
+            "iou_threshold": str(nms.iouThreshold),
+            "confidence_threshold": str(nms.confidenceThreshold),
+        }
+    )
+
+    # Save the model
+    f = file.with_suffix(".mlmodel")  # filename
+    model = ct.models.MLModel(pipeline.spec)
+    model.input_description["image"] = "Input image"
+    model.input_description["iouThreshold"] = f"(optional) IOU Threshold override (default: {nms.iouThreshold})"
+    model.input_description[
+        "confidenceThreshold"
+    ] = f"(optional) Confidence Threshold override (default: {nms.confidenceThreshold})"
+    model.output_description["confidence"] = 'Boxes × Class confidence (see user-defined metadata "classes")'
+    model.output_description["coordinates"] = "Boxes × [x, y, width, height] (relative to image size)"
+    model.save(f)  # pipelined
+    print(f"{prefix} pipeline success ({time.time() - t:.2f}s), saved as {f} ({file_size(f):.1f} MB)")
+
+
+@smart_inference_mode()
+def run(
+    data=ROOT / "data/coco128.yaml",  # 'dataset.yaml path'
+    weights=ROOT / "yolov5s.pt",  # weights path
+    imgsz=(640, 640),  # image (height, width)
+    batch_size=1,  # batch size
+    device="cpu",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+    include=("torchscript", "onnx"),  # include formats
+    half=False,  # FP16 half-precision export
+    inplace=False,  # set YOLOv5 Detect() inplace=True
+    keras=False,  # use Keras
+    optimize=False,  # TorchScript: optimize for mobile
+    int8=False,  # CoreML/TF INT8 quantization
+    per_tensor=False,  # TF per tensor quantization
+    dynamic=False,  # ONNX/TF/TensorRT: dynamic axes
+    simplify=False,  # ONNX: simplify model
+    opset=12,  # ONNX: opset version
+    verbose=False,  # TensorRT: verbose log
+    workspace=4,  # TensorRT: workspace size (GB)
+    nms=False,  # TF: add NMS to model
+    agnostic_nms=False,  # TF: add agnostic NMS to model
+    topk_per_class=100,  # TF.js NMS: topk per class to keep
+    topk_all=100,  # TF.js NMS: topk for all classes to keep
+    iou_thres=0.45,  # TF.js NMS: IoU threshold
+    conf_thres=0.25,  # TF.js NMS: confidence threshold
+):
+    t = time.time()
+    include = [x.lower() for x in include]  # to lowercase
+    fmts = tuple(export_formats()["Argument"][1:])  # --include arguments
+    flags = [x in include for x in fmts]
+    assert sum(flags) == len(include), f"ERROR: Invalid --include {include}, valid --include arguments are {fmts}"
+    jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags  # export booleans
+    file = Path(url2file(weights) if str(weights).startswith(("http:/", "https:/")) else weights)  # PyTorch weights
+
+    # Load PyTorch model
+    device = select_device(device)
+    if half:
+        assert device.type != "cpu" or coreml, "--half only compatible with GPU export, i.e. use --device 0"
+        assert not dynamic, "--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both"
+    model = attempt_load(weights, device=device, inplace=True, fuse=True)  # load FP32 model
+
+    # Checks
+    imgsz *= 2 if len(imgsz) == 1 else 1  # expand
+    if optimize:
+        assert device.type == "cpu", "--optimize not compatible with cuda devices, i.e. use --device cpu"
+
+    # Input
+    gs = int(max(model.stride))  # grid size (max stride)
+    imgsz = [check_img_size(x, gs) for x in imgsz]  # verify img_size are gs-multiples
+    im = torch.zeros(batch_size, 3, *imgsz).to(device)  # image size(1,3,320,192) BCHW iDetection
+
+    # Update model
+    model.eval()
+    for k, m in model.named_modules():
+        if isinstance(m, Detect):
+            m.inplace = inplace
+            m.dynamic = dynamic
+            m.export = True
+
+    for _ in range(2):
+        y = model(im)  # dry runs
+    if half and not coreml:
+        im, model = im.half(), model.half()  # to FP16
+    shape = tuple((y[0] if isinstance(y, tuple) else y).shape)  # model output shape
+    metadata = {"stride": int(max(model.stride)), "names": model.names}  # model metadata
+    LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")
+
+    # Exports
+    f = [""] * len(fmts)  # exported filenames
+    warnings.filterwarnings(action="ignore", category=torch.jit.TracerWarning)  # suppress TracerWarning
+    if jit:  # TorchScript
+        f[0], _ = export_torchscript(model, im, file, optimize)
+    if engine:  # TensorRT required before ONNX
+        f[1], _ = export_engine(model, im, file, half, dynamic, simplify, workspace, verbose)
+    if onnx or xml:  # OpenVINO requires ONNX
+        f[2], _ = export_onnx(model, im, file, opset, dynamic, simplify)
+    if xml:  # OpenVINO
+        f[3], _ = export_openvino(file, metadata, half, int8, data)
+    if coreml:  # CoreML
+        f[4], ct_model = export_coreml(model, im, file, int8, half, nms)
+        if nms:
+            pipeline_coreml(ct_model, im, file, model.names, y)
+    if any((saved_model, pb, tflite, edgetpu, tfjs)):  # TensorFlow formats
+        assert not tflite or not tfjs, "TFLite and TF.js models must be exported separately, please pass only one type."
+        assert not isinstance(model, ClassificationModel), "ClassificationModel export to TF formats not yet supported."
+        f[5], s_model = export_saved_model(
+            model.cpu(),
+            im,
+            file,
+            dynamic,
+            tf_nms=nms or agnostic_nms or tfjs,
+            agnostic_nms=agnostic_nms or tfjs,
+            topk_per_class=topk_per_class,
+            topk_all=topk_all,
+            iou_thres=iou_thres,
+            conf_thres=conf_thres,
+            keras=keras,
+        )
+        if pb or tfjs:  # pb prerequisite to tfjs
+            f[6], _ = export_pb(s_model, file)
+        if tflite or edgetpu:
+            f[7], _ = export_tflite(
+                s_model, im, file, int8 or edgetpu, per_tensor, data=data, nms=nms, agnostic_nms=agnostic_nms
+            )
+            if edgetpu:
+                f[8], _ = export_edgetpu(file)
+            add_tflite_metadata(f[8] or f[7], metadata, num_outputs=len(s_model.outputs))
+        if tfjs:
+            f[9], _ = export_tfjs(file, int8)
+    if paddle:  # PaddlePaddle
+        f[10], _ = export_paddle(model, im, file, metadata)
+
+    # Finish
+    f = [str(x) for x in f if x]  # filter out '' and None
+    if any(f):
+        cls, det, seg = (isinstance(model, x) for x in (ClassificationModel, DetectionModel, SegmentationModel))  # type
+        det &= not seg  # segmentation models inherit from SegmentationModel(DetectionModel)
+        dir = Path("segment" if seg else "classify" if cls else "")
+        h = "--half" if half else ""  # --half FP16 inference arg
+        s = (
+            "# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference"
+            if cls
+            else "# WARNING ⚠️ SegmentationModel not yet supported for PyTorch Hub AutoShape inference"
+            if seg
+            else ""
+        )
+        LOGGER.info(
+            f'\nExport complete ({time.time() - t:.1f}s)'
+            f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
+            f"\nDetect:          python {dir / ('detect.py' if det else 'predict.py')} --weights {f[-1]} {h}"
+            f"\nValidate:        python {dir / 'val.py'} --weights {f[-1]} {h}"
+            f"\nPyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}')  {s}"
+            f'\nVisualize:       https://netron.app'
+        )
+    return f  # return list of exported files/dirs
+
+
+def parse_opt(known=False):
+    """Parses command-line arguments for YOLOv5 model export configurations, returning the parsed options."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
+    parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model.pt path(s)")
+    parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640, 640], help="image (h, w)")
+    parser.add_argument("--batch-size", type=int, default=1, help="batch size")
+    parser.add_argument("--device", default="cpu", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
+    parser.add_argument("--half", action="store_true", help="FP16 half-precision export")
+    parser.add_argument("--inplace", action="store_true", help="set YOLOv5 Detect() inplace=True")
+    parser.add_argument("--keras", action="store_true", help="TF: use Keras")
+    parser.add_argument("--optimize", action="store_true", help="TorchScript: optimize for mobile")
+    parser.add_argument("--int8", action="store_true", help="CoreML/TF/OpenVINO INT8 quantization")
+    parser.add_argument("--per-tensor", action="store_true", help="TF per-tensor quantization")
+    parser.add_argument("--dynamic", action="store_true", help="ONNX/TF/TensorRT: dynamic axes")
+    parser.add_argument("--simplify", action="store_true", help="ONNX: simplify model")
+    parser.add_argument("--opset", type=int, default=17, help="ONNX: opset version")
+    parser.add_argument("--verbose", action="store_true", help="TensorRT: verbose log")
+    parser.add_argument("--workspace", type=int, default=4, help="TensorRT: workspace size (GB)")
+    parser.add_argument("--nms", action="store_true", help="TF: add NMS to model")
+    parser.add_argument("--agnostic-nms", action="store_true", help="TF: add agnostic NMS to model")
+    parser.add_argument("--topk-per-class", type=int, default=100, help="TF.js NMS: topk per class to keep")
+    parser.add_argument("--topk-all", type=int, default=100, help="TF.js NMS: topk for all classes to keep")
+    parser.add_argument("--iou-thres", type=float, default=0.45, help="TF.js NMS: IoU threshold")
+    parser.add_argument("--conf-thres", type=float, default=0.25, help="TF.js NMS: confidence threshold")
+    parser.add_argument(
+        "--include",
+        nargs="+",
+        default=["torchscript"],
+        help="torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle",
+    )
+    opt = parser.parse_known_args()[0] if known else parser.parse_args()
+    print_args(vars(opt))
+    return opt
+
+
+def main(opt):
+    """Executes the YOLOv5 model inference or export with specified weights and options."""
+    for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
+        run(**vars(opt))
+
+
+if __name__ == "__main__":
+    opt = parse_opt()
+    main(opt)

hubconf.py

+# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
+"""
+PyTorch Hub models https://pytorch.org/hub/ultralytics_yolov5
+
+Usage:
+    import torch
+    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # official model
+    model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s')  # from branch
+    model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt')  # custom/local model
+    model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local')  # local repo
+"""
+
+import torch
+
+
+def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
+    """
+    Creates or loads a YOLOv5 model.
+
+    Arguments:
+        name (str): model name 'yolov5s' or path 'path/to/best.pt'
+        pretrained (bool): load pretrained weights into the model
+        channels (int): number of input channels
+        classes (int): number of model classes
+        autoshape (bool): apply YOLOv5 .autoshape() wrapper to model
+        verbose (bool): print all information to screen
+        device (str, torch.device, None): device to use for model parameters
+
+    Returns:
+        YOLOv5 model
+    """
+    from pathlib import Path
+
+    from models.common import AutoShape, DetectMultiBackend
+    from models.experimental import attempt_load
+    from models.yolo import ClassificationModel, DetectionModel, SegmentationModel
+    from utils.downloads import attempt_download
+    from utils.general import LOGGER, ROOT, check_requirements, intersect_dicts, logging
+    from utils.torch_utils import select_device
+
+    if not verbose:
+        LOGGER.setLevel(logging.WARNING)
+    check_requirements(ROOT / "requirements.txt", exclude=("opencv-python", "tensorboard", "thop"))
+    name = Path(name)
+    path = name.with_suffix(".pt") if name.suffix == "" and not name.is_dir() else name  # checkpoint path
+    try:
+        device = select_device(device)
+        if pretrained and channels == 3 and classes == 80:
+            try:
+                model = DetectMultiBackend(path, device=device, fuse=autoshape)  # detection model
+                if autoshape:
+                    if model.pt and isinstance(model.model, ClassificationModel):
+                        LOGGER.warning(
+                            "WARNING ⚠️ YOLOv5 ClassificationModel is not yet AutoShape compatible. "
+                            "You must pass torch tensors in BCHW to this model, i.e. shape(1,3,224,224)."
+                        )
+                    elif model.pt and isinstance(model.model, SegmentationModel):
+                        LOGGER.warning(
+                            "WARNING ⚠️ YOLOv5 SegmentationModel is not yet AutoShape compatible. "
+                            "You will not be able to run inference with this model."
+                        )
+                    else:
+                        model = AutoShape(model)  # for file/URI/PIL/cv2/np inputs and NMS
+            except Exception:
+                model = attempt_load(path, device=device, fuse=False)  # arbitrary model
+        else:
+            cfg = list((Path(__file__).parent / "models").rglob(f"{path.stem}.yaml"))[0]  # model.yaml path
+            model = DetectionModel(cfg, channels, classes)  # create model
+            if pretrained:
+                ckpt = torch.load(attempt_download(path), map_location=device)  # load
+                csd = ckpt["model"].float().state_dict()  # checkpoint state_dict as FP32
+                csd = intersect_dicts(csd, model.state_dict(), exclude=["anchors"])  # intersect
+                model.load_state_dict(csd, strict=False)  # load
+                if len(ckpt["model"].names) == classes:
+                    model.names = ckpt["model"].names  # set class names attribute
+        if not verbose:
+            LOGGER.setLevel(logging.INFO)  # reset to default
+        return model.to(device)
+
+    except Exception as e:
+        help_url = "https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading"
+        s = f"{e}. Cache may be out of date, try `force_reload=True` or see {help_url} for help."
+        raise Exception(s) from e
+
+
+def custom(path="path/to/model.pt", autoshape=True, _verbose=True, device=None):
+    """Loads a custom or local YOLOv5 model from a given path with optional autoshaping and device specification."""
+    return _create(path, autoshape=autoshape, verbose=_verbose, device=device)
+
+
+def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """Instantiates the YOLOv5-nano model with options for pretraining, input channels, class count, autoshaping,
+    verbosity, and device.
+    """
+    return _create("yolov5n", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """Creates YOLOv5-small model with options for pretraining, input channels, class count, autoshaping, verbosity, and
+    device.
+    """
+    return _create("yolov5s", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """Instantiates the YOLOv5-medium model with customizable pretraining, channel count, class count, autoshaping,
+    verbosity, and device.
+    """
+    return _create("yolov5m", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """Creates YOLOv5-large model with options for pretraining, channels, classes, autoshaping, verbosity, and device
+    selection.
+    """
+    return _create("yolov5l", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """Instantiates the YOLOv5-xlarge model with customizable pretraining, channel count, class count, autoshaping,
+    verbosity, and device.
+    """
+    return _create("yolov5x", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """Creates YOLOv5-nano-P6 model with options for pretraining, channels, classes, autoshaping, verbosity, and
+    device.
+    """
+    return _create("yolov5n6", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """Instantiate YOLOv5-small-P6 model with options for pretraining, input channels, number of classes, autoshaping,
+    verbosity, and device selection.
+    """
+    return _create("yolov5s6", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """Creates YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity,
+    and device.
+    """
+    return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """Instantiates the YOLOv5-large-P6 model with customizable pretraining, channel and class counts, autoshaping,
+    verbosity, and device selection.
+    """
+    return _create("yolov5l6", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """Creates YOLOv5-xlarge-P6 model with options for pretraining, channels, classes, autoshaping, verbosity, and
+    device.
+    """
+    return _create("yolov5x6", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+if __name__ == "__main__":
+    import argparse
+    from pathlib import Path
+
+    import numpy as np
+    from PIL import Image
+
+    from utils.general import cv2, print_args
+
+    # Argparser
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, default="yolov5s", help="model name")
+    opt = parser.parse_args()
+    print_args(vars(opt))
+
+    # Model
+    model = _create(name=opt.model, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True)
+    # model = custom(path='path/to/model.pt')  # custom
+
+    # Images
+    imgs = [
+        "data/images/zidane.jpg",  # filename
+        Path("data/images/zidane.jpg"),  # Path
+        "https://ultralytics.com/images/zidane.jpg",  # URI
+        cv2.imread("data/images/bus.jpg")[:, :, ::-1],  # OpenCV
+        Image.open("data/images/bus.jpg"),  # PIL
+        np.zeros((320, 640, 3)),
+    ]  # numpy
+
+    # Inference
+    results = model(imgs, size=320)  # batched inference
+
+    # Results
+    results.print()
+    results.save()