Add CLI Fire for detection script

detect.py

@@ -24,13 +24,14 @@
                                          yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
 """
 
-import argparse
 import os
 import sys
 from pathlib import Path
+from typing import Optional, Sequence, Union
 
 import torch
 import torch.backends.cudnn as cudnn
+from fire import Fire
 
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[0]  # YOLOv5 root directory
@@ -48,34 +49,67 @@
 
 @torch.no_grad()
 def run(
-        weights=ROOT / 'yolov5s.pt',  # model.pt path(s)
-        source=ROOT / 'data/images',  # file/dir/URL/glob, 0 for 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_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
+    weights: Union[str, Path] = ROOT / 'yolov5s.pt',
+    source: Union[str, Path] = ROOT / 'data/images',
+    data=ROOT / 'data/coco128.yaml',
+    imgsz: Sequence[int] = (640, 640),
+    conf_thres: float = 0.25,
+    iou_thres: float = 0.45,
+    max_det: int = 1000,
+    device: str = '',
+    view_img: bool = False,
+    save_txt: bool = False,
+    save_conf: bool = False,
+    save_crop: bool = False,
+    nosave: bool = False,
+    classes: Optional[Sequence] = None,
+    agnostic_nms: bool = False,
+    augment: bool = False,
+    visualize: bool = False,
+    update: bool = False,
+    project: Union[str, Path] = ROOT / 'runs/detect',
+    name: str = 'exp',
+    exist_ok: bool = False,
+    line_thickness: int = 3,
+    hide_labels: bool = False,
+    hide_conf: bool = False,
+    half: bool = False,
+    dnn: bool = False,
 ):
+    """
+    Args:
+        weights: model path(s)
+        source: file/dir/URL/glob, 0 for webcam
+        data: dataset.yaml path
+        imgsz: inference size (height, width)
+        conf_thres: confidence threshold
+        iou_thres: NMS IoU threshold
+        max_det: maximum detections per image
+        device: cuda device, i.e. 0 or 0,1,2,3 or cpu
+        view_img: show results
+        save_txt: save results to *.txt
+        save_conf: save confidences in `save_txt` labels
+        save_crop: save cropped prediction boxes
+        nosave: do not save images/videos
+        classes: filter by class: `classes=0`, or `classes 0 2 3`
+        agnostic_nms: class-agnostic NMS
+        augment: augmented inference
+        visualize: visualize features
+        update: update all models
+        project: save results to project/name
+        name: save results to project/name
+        exist_ok: existing project/name ok, do not increment
+        line_thickness: bounding box thickness (pixels)
+        hide_labels: hide labels
+        hide_conf: hide confidences
+        half: use FP16 half-precision inference
+        dnn: use OpenCV DNN for ONNX inference
+    """
+    print_args()
     source = str(source)
+    if not isinstance(imgsz, Sequence):
+        imgsz = [imgsz]
+    imgsz *= 2 if len(imgsz) == 1 else 1  # expand
     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://'))
@@ -208,45 +242,6 @@ def run(
         strip_optimizer(weights)  # update model (to fix SourceChangeWarning)
 
 
-def parse_opt():
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)')
-    parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob, 0 for 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-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')
-    opt = parser.parse_args()
-    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
-    print_args(vars(opt))
-    return opt
-
-
-def main(opt):
-    check_requirements(exclude=('tensorboard', 'thop'))
-    run(**vars(opt))
-
-
 if __name__ == "__main__":
-    opt = parse_opt()
-    main(opt)
+    check_requirements(exclude=('tensorboard', 'thop'))
+    Fire(run)

requirements.txt

@@ -11,6 +11,7 @@ scipy>=1.4.1
 torch>=1.7.0
 torchvision>=0.8.1
 tqdm>=4.41.0
+fire
 
 # Logging -------------------------------------
 tensorboard>=2.4.1