update to ultralytics/yolov5 24.04.21 (#19)

* update to v5.0.1

* update to ultralytics/yolo 24.04.21

* update readme

* remove redundant file

* update onnx coremltools export workflow

.github/workflows/ci.yml

@@ -74,5 +74,6 @@ jobs:
           python yolov5/test.py --img 128 --batch 16 --weights runs/train/exp/weights/last.pt --device $di
           yolo_test --img 128 --batch 16 --weights runs/train/exp/weights/last.pt --device $di
           # export
+          pip install onnx coremltools
           python yolov5/models/export.py --weights yolov5/weights/yolov5s.pt --device $di
           yolo_export --weights yolov5/weights/yolov5s.pt --device $di

.github/workflows/package_testing.yml

@@ -70,3 +70,7 @@ jobs:
           yolo_test --img 128 --batch 16 --weights yolov5/weights/yolov5s.pt --device $di
           python yolov5/test.py --img 128 --batch 16 --weights runs/train/exp/weights/last.pt --device $di
           yolo_test --img 128 --batch 16 --weights runs/train/exp/weights/last.pt --device $di
+          # export
+          pip install onnx coremltools
+          python yolov5/models/export.py --weights yolov5/weights/yolov5s.pt --device $di
+          yolo_export --weights yolov5/weights/yolov5s.pt --device $di

README.md

@@ -11,7 +11,7 @@ You can finally install [YOLOv5 object detector](https://github.com/ultralytics/
 
 This package is up-to-date with the latest release of [ultralytics/yolov5](https://github.com/ultralytics/yolov5).
 
-<img src="https://user-images.githubusercontent.com/26833433/103594689-455e0e00-4eae-11eb-9cdf-7d753e2ceeeb.png" width="1000">
+<img src="https://user-images.githubusercontent.com/26833433/114313216-f0a5e100-9af5-11eb-8445-c682b60da2e3.png" width="1000">
 
 ## Installation
 

yolov5/__init__.py

@@ -1,5 +1,4 @@
 from yolov5.helpers import YOLOv5
 from yolov5.helpers import load_model as load
 
-
-__version__ = "5.0.0"
+__version__ = "5.0.1"

yolov5/data/argoverse_hd.yaml

@@ -1,6 +1,6 @@
 # Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
 # Train command: python train.py --data argoverse_hd.yaml
-# Default dataset location is next to /yolov5:
+# Default dataset location is next to YOLOv5:
 #   /parent_folder
 #     /argoverse
 #     /yolov5

yolov5/data/coco.yaml

@@ -1,6 +1,6 @@
 # COCO 2017 dataset http://cocodataset.org
 # Train command: python train.py --data coco.yaml
-# Default dataset location is next to /yolov5:
+# Default dataset location is next to YOLOv5:
 #   /parent_folder
 #     /coco
 #     /yolov5
@@ -30,6 +30,6 @@ names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', '
 
 # Print classes
 # with open('data/coco.yaml') as f:
-#   d = yaml.load(f, Loader=yaml.FullLoader)  # dict
+#   d = yaml.safe_load(f)  # dict
 #   for i, x in enumerate(d['names']):
 #     print(i, x)

yolov5/data/coco128.yaml

@@ -1,6 +1,6 @@
 # COCO 2017 dataset http://cocodataset.org - first 128 training images
 # Train command: python train.py --data coco128.yaml
-# Default dataset location is next to /yolov5:
+# Default dataset location is next to YOLOv5:
 #   /parent_folder
 #     /coco128
 #     /yolov5

yolov5/data/scripts/get_argoverse_hd.sh

@@ -2,7 +2,7 @@
 # Argoverse-HD dataset (ring-front-center camera) http://www.cs.cmu.edu/~mengtial/proj/streaming/
 # Download command: bash data/scripts/get_argoverse_hd.sh
 # Train command: python train.py --data argoverse_hd.yaml
-# Default dataset location is next to /yolov5:
+# Default dataset location is next to YOLOv5:
 #   /parent_folder
 #     /argoverse
 #     /yolov5

yolov5/data/scripts/get_coco.sh

@@ -2,7 +2,7 @@
 # COCO 2017 dataset http://cocodataset.org
 # Download command: bash data/scripts/get_coco.sh
 # Train command: python train.py --data coco.yaml
-# Default dataset location is next to /yolov5:
+# Default dataset location is next to YOLOv5:
 #   /parent_folder
 #     /coco
 #     /yolov5

yolov5/data/scripts/get_voc.sh

@@ -2,7 +2,7 @@
 # PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/
 # Download command: bash data/scripts/get_voc.sh
 # Train command: python train.py --data voc.yaml
-# Default dataset location is next to /yolov5:
+# Default dataset location is next to YOLOv5:
 #   /parent_folder
 #     /VOC
 #     /yolov5
@@ -29,34 +29,27 @@ echo "Completed in" $runtime "seconds"
 
 echo "Splitting dataset..."
 python3 - "$@" <<END
-import xml.etree.ElementTree as ET
-import pickle
 import os
-from os import listdir, getcwd
-from os.path import join
+import xml.etree.ElementTree as ET
+from os import getcwd
+
+sets = [('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
 
-sets=[('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
+classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog",
+           "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
 
-classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
 
+def convert_box(size, box):
+    dw = 1. / (size[0])
+    dh = 1. / (size[1])
+    x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]
+    return x * dw, y * dh, w * dw, h * dh
 
-def convert(size, box):
-    dw = 1./(size[0])
-    dh = 1./(size[1])
-    x = (box[0] + box[1])/2.0 - 1
-    y = (box[2] + box[3])/2.0 - 1
-    w = box[1] - box[0]
-    h = box[3] - box[2]
-    x = x*dw
-    w = w*dw
-    y = y*dh
-    h = h*dh
-    return (x,y,w,h)
 
 def convert_annotation(year, image_id):
-    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
-    out_file = open('VOCdevkit/VOC%s/labels/%s.txt'%(year, image_id), 'w')
-    tree=ET.parse(in_file)
+    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml' % (year, image_id))
+    out_file = open('VOCdevkit/VOC%s/labels/%s.txt' % (year, image_id), 'w')
+    tree = ET.parse(in_file)
     root = tree.getroot()
     size = root.find('size')
     w = int(size.find('width').text)
@@ -65,74 +58,58 @@ def convert_annotation(year, image_id):
     for obj in root.iter('object'):
         difficult = obj.find('difficult').text
         cls = obj.find('name').text
-        if cls not in classes or int(difficult)==1:
+        if cls not in classes or int(difficult) == 1:
             continue
         cls_id = classes.index(cls)
         xmlbox = obj.find('bndbox')
-        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
-        bb = convert((w,h), b)
+        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
+             float(xmlbox.find('ymax').text))
+        bb = convert_box((w, h), b)
         out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
 
-wd = getcwd()
 
+cwd = getcwd()
 for year, image_set in sets:
-    if not os.path.exists('VOCdevkit/VOC%s/labels/'%(year)):
-        os.makedirs('VOCdevkit/VOC%s/labels/'%(year))
-    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
-    list_file = open('%s_%s.txt'%(year, image_set), 'w')
+    if not os.path.exists('VOCdevkit/VOC%s/labels/' % year):
+        os.makedirs('VOCdevkit/VOC%s/labels/' % year)
+    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt' % (year, image_set)).read().strip().split()
+    list_file = open('%s_%s.txt' % (year, image_set), 'w')
     for image_id in image_ids:
-        list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n'%(wd, year, image_id))
+        list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n' % (cwd, year, image_id))
         convert_annotation(year, image_id)
     list_file.close()
-
 END
 
 cat 2007_train.txt 2007_val.txt 2012_train.txt 2012_val.txt >train.txt
 cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt >train.all.txt
 
-python3 - "$@" <<END
+mkdir ../VOC ../VOC/images ../VOC/images/train ../VOC/images/val
+mkdir ../VOC/labels ../VOC/labels/train ../VOC/labels/val
 
-import shutil
+python3 - "$@" <<END
 import os
-os.system('mkdir ../VOC/')
-os.system('mkdir ../VOC/images')
-os.system('mkdir ../VOC/images/train')
-os.system('mkdir ../VOC/images/val')
-
-os.system('mkdir ../VOC/labels')
-os.system('mkdir ../VOC/labels/train')
-os.system('mkdir ../VOC/labels/val')
 
-import os
 print(os.path.exists('../tmp/train.txt'))
-f = open('../tmp/train.txt', 'r')
-lines = f.readlines()
-
-for line in lines:
-    line = "/".join(line.split('/')[-5:]).strip()
-    if (os.path.exists("../" + line)):
-        os.system("cp ../"+ line + " ../VOC/images/train")
-        
-    line = line.replace('JPEGImages', 'labels')
-    line = line.replace('jpg', 'txt')
-    if (os.path.exists("../" + line)):
-        os.system("cp ../"+ line + " ../VOC/labels/train")
+with open('../tmp/train.txt', 'r') as f:
+    for line in f.readlines():
+        line = "/".join(line.split('/')[-5:]).strip()
+        if os.path.exists("../" + line):
+            os.system("cp ../" + line + " ../VOC/images/train")
 
+        line = line.replace('JPEGImages', 'labels').replace('jpg', 'txt')
+        if os.path.exists("../" + line):
+            os.system("cp ../" + line + " ../VOC/labels/train")
 
 print(os.path.exists('../tmp/2007_test.txt'))
-f = open('../tmp/2007_test.txt', 'r')
-lines = f.readlines()
-
-for line in lines:
-    line = "/".join(line.split('/')[-5:]).strip()
-    if (os.path.exists("../" + line)):
-        os.system("cp ../"+ line + " ../VOC/images/val")
-        
-    line = line.replace('JPEGImages', 'labels')
-    line = line.replace('jpg', 'txt')
-    if (os.path.exists("../" + line)):
-        os.system("cp ../"+ line + " ../VOC/labels/val")
-
+with open('../tmp/2007_test.txt', 'r') as f:
+    for line in f.readlines():
+        line = "/".join(line.split('/')[-5:]).strip()
+        if os.path.exists("../" + line):
+            os.system("cp ../" + line + " ../VOC/images/val")
+
+        line = line.replace('JPEGImages', 'labels').replace('jpg', 'txt')
+        if os.path.exists("../" + line):
+            os.system("cp ../" + line + " ../VOC/labels/val")
 END
 
 rm -rf ../tmp # remove temporary directory

yolov5/data/visdrone.yaml

@@ -0,0 +1,65 @@
+# VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset
+# Train command: python train.py --data visdrone.yaml
+# Default dataset location is next to YOLOv5:
+#   /parent_folder
+#     /VisDrone
+#     /yolov5
+
+
+# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
+train: ../VisDrone/VisDrone2019-DET-train/images  # 6471 images
+val: ../VisDrone/VisDrone2019-DET-val/images  # 548 images
+test: ../VisDrone/VisDrone2019-DET-test-dev/images  # 1610 images
+
+# number of classes
+nc: 10
+
+# class names
+names: [ 'pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor' ]
+
+
+# download command/URL (optional) --------------------------------------------------------------------------------------
+download: |
+  import os
+  from pathlib import Path
+
+  from utils.general import download
+
+
+  def visdrone2yolo(dir):
+      from PIL import Image
+      from tqdm import tqdm
+
+      def convert_box(size, box):
+          # Convert VisDrone box to YOLO xywh box
+          dw = 1. / size[0]
+          dh = 1. / size[1]
+          return (box[0] + box[2] / 2) * dw, (box[1] + box[3] / 2) * dh, box[2] * dw, box[3] * dh
+
+      (dir / 'labels').mkdir(parents=True, exist_ok=True)  # make labels directory
+      pbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}')
+      for f in pbar:
+          img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).size
+          lines = []
+          with open(f, 'r') as file:  # read annotation.txt
+              for row in [x.split(',') for x in file.read().strip().splitlines()]:
+                  if row[4] == '0':  # VisDrone 'ignored regions' class 0
+                      continue
+                  cls = int(row[5]) - 1
+                  box = convert_box(img_size, tuple(map(int, row[:4])))
+                  lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n")
+                  with open(str(f).replace(os.sep + 'annotations' + os.sep, os.sep + 'labels' + os.sep), 'w') as fl:
+                      fl.writelines(lines)  # write label.txt
+
+
+  # Download
+  dir = Path('../VisDrone')  # dataset directory
+  urls = ['https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-train.zip',
+          'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-val.zip',
+          'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-dev.zip'
+          'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-challenge.zip']
+  download(urls, dir=dir)
+
+  # Convert
+  for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':
+      visdrone2yolo(dir / d)  # convert VisDrone annotations to YOLO labels

yolov5/data/voc.yaml

@@ -1,6 +1,6 @@
 # PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/
 # Train command: python train.py --data voc.yaml
-# Default dataset location is next to /yolov5:
+# Default dataset location is next to YOLOv5:
 #   /parent_folder
 #     /VOC
 #     /yolov5

yolov5/detect.py

@@ -12,21 +12,21 @@
 from yolov5.utils.general import (apply_classifier, check_img_size,
                                   check_imshow, check_requirements,
                                   increment_path, non_max_suppression,
-                                  scale_coords, set_logging, strip_optimizer,
-                                  xyxy2xywh)
+                                  save_one_box, scale_coords, set_logging,
+                                  strip_optimizer, xyxy2xywh)
 from yolov5.utils.plots import plot_one_box
 from yolov5.utils.torch_utils import (load_classifier, select_device,
                                       time_synchronized)
 
 
-def detect(save_img=False, opt=None):
+def detect(opt):
     source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
     save_img = not opt.nosave and not source.endswith('.txt')  # save inference images
     webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
         ('rtsp://', 'rtmp://', 'http://', 'https://'))
 
     # Directories
-    save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))  # increment run
+    save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)  # increment run
     (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
 
     # Initialize
@@ -88,7 +88,7 @@ def detect(save_img=False, opt=None):
             if webcam:  # batch_size >= 1
                 p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
             else:
-                p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
+                p, s, im0, frame = path, '', im0s.copy(), getattr(dataset, 'frame', 0)
 
             p = Path(p)  # to Path
             save_path = str(save_dir / p.name)  # img.jpg
@@ -112,9 +112,13 @@ def detect(save_img=False, opt=None):
                         with open(txt_path + '.txt', 'a') as f:
                             f.write(('%g ' * len(line)).rstrip() % line + '\n')
 
-                    if save_img or view_img:  # Add bbox to image
-                        label = f'{names[int(cls)]} {conf:.2f}'
-                        plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
+                    if save_img or opt.save_crop or view_img:  # Add bbox to image
+                        c = int(cls)  # integer class
+                        label = None if opt.hide_labels else (names[c] if opt.hide_conf else f'{names[c]} {conf:.2f}')
+
+                        plot_one_box(xyxy, im0, label=label, color=colors[c], line_thickness=opt.line_thickness)
+                        if opt.save_crop:
+                            save_one_box(xyxy, im0s, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
 
             # Print time (inference + NMS)
             print(f'{s}Done. ({t2 - t1:.3f}s)')
@@ -161,6 +165,7 @@ def main():
     parser.add_argument('--view-img', action='store_true', help='display results')
     parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
     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: --class 0, or --class 0 2 3')
     parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
@@ -169,6 +174,9 @@ def main():
     parser.add_argument('--project', default='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=True, action='store_true', help='hide labels')
+    parser.add_argument('--hide-conf', default=True, action='store_true', help='hide confidences')
     opt = parser.parse_args()
     print(opt)
     #check_requirements(exclude=('pycocotools', 'thop'))