Streaming Classification support (#9106)

* Streaming Classification support

* Streaming Classification support

* Streaming Classification support

classify/predict.py

@@ -1,12 +1,15 @@
 # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
 """
-Run YOLOv5 classification inference on images, videos, directories, and globs.
+Run YOLOv5 classification inference on images, videos, directories, globs, YouTube, webcam, streams, etc.
 
 Usage - sources:
-    $ python classify/predict.py --weights yolov5s.pt --source img.jpg        # image
-                                                               vid.mp4        # video
-                                                               path/          # directory
-                                                               'path/*.jpg'   # glob
+    $ python classify/predict.py --weights yolov5s-cls.pt --source 0                               # webcam
+                                                                   img.jpg                         # image
+                                                                   vid.mp4                         # video
+                                                                   path/                           # directory
+                                                                   'path/*.jpg'                    # glob
+                                                                   'https://youtu.be/Zgi9g1ksQHc'  # YouTube
+                                                                   'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream
 
 Usage - formats:
     $ python classify/predict.py --weights yolov5s-cls.pt                 # PyTorch
@@ -23,9 +26,11 @@
 
 import argparse
 import os
+import platform
 import sys
 from pathlib import Path
 
+import torch.backends.cudnn as cudnn
 import torch.nn.functional as F
 
 FILE = Path(__file__).resolve()
@@ -36,79 +41,164 @@
 
 from models.common import DetectMultiBackend
 from utils.augmentations import classify_transforms
-from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages
-from utils.general import LOGGER, Profile, check_file, check_requirements, colorstr, increment_path, print_args
+from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
+from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
+                           increment_path, print_args, strip_optimizer)
+from utils.plots import Annotator
 from utils.torch_utils import select_device, smart_inference_mode
 
 
 @smart_inference_mode()
 def run(
         weights=ROOT / 'yolov5s-cls.pt',  # model.pt path(s)
-        source=ROOT / 'data/images',  # file/dir/URL/glob
-        imgsz=224,  # inference size
+        source=ROOT / 'data/images',  # file/dir/URL/glob, 0 for webcam
+        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
+        imgsz=(224, 224),  # inference size (height, width)
         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
+        nosave=False,  # do not save images/videos
+        augment=False,  # augmented inference
+        visualize=False,  # visualize features
+        update=False,  # update all models
+        project=ROOT / 'runs/predict-cls',  # save results to project/name
+        name='exp',  # save results to project/name
+        exist_ok=False,  # existing project/name ok, do not increment
         half=False,  # use FP16 half-precision inference
         dnn=False,  # use OpenCV DNN for ONNX inference
-        project=ROOT / 'runs/predict-cls',  # save to project/name
-        name='exp',  # save to project/name
-        exist_ok=False,  # existing project/name ok, do not increment
 ):
     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('.txt') or (is_url and not is_file)
     if is_url and is_file:
         source = check_file(source)  # download
 
-    dt = Profile(), Profile(), Profile()
-    device = select_device(device)
-
     # Directories
     save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
-    save_dir.mkdir(parents=True, exist_ok=True)  # make dir
+    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
 
     # Load model
-    model = DetectMultiBackend(weights, device=device, dnn=dnn, fp16=half)
-    model.warmup(imgsz=(1, 3, imgsz, imgsz))  # warmup
-    dataset = LoadImages(source, img_size=imgsz, transforms=classify_transforms(imgsz))
-    for seen, (path, im, im0s, vid_cap, s) in enumerate(dataset):
-        # Image
+    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
+    if webcam:
+        view_img = check_imshow()
+        cudnn.benchmark = True  # set True to speed up constant image size inference
+        dataset = LoadStreams(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]))
+        bs = len(dataset)  # batch_size
+    else:
+        dataset = LoadImages(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]))
+        bs = 1  # batch_size
+    vid_path, vid_writer = [None] * bs, [None] * bs
+
+    # Run inference
+    model.warmup(imgsz=(1 if pt else bs, 3, *imgsz))  # warmup
+    seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
+    for path, im, im0s, vid_cap, s in dataset:
         with dt[0]:
-            im = im.unsqueeze(0).to(device)
-            im = im.half() if model.fp16 else im.float()
+            im = im.to(device)
+            im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
+            if len(im.shape) == 3:
+                im = im[None]  # expand for batch dim
 
         # Inference
         with dt[1]:
             results = model(im)
 
         # Post-process
         with dt[2]:
-            p = F.softmax(results, dim=1)  # probabilities
-            i = p.argsort(1, descending=True)[:, :5].squeeze().tolist()  # top 5 indices
-            # if save:
-            #    imshow_cls(im, f=save_dir / Path(path).name, verbose=True)
-            LOGGER.info(
-                f"{s}{imgsz}x{imgsz} {', '.join(f'{model.names[j]} {p[0, j]:.2f}' for j in i)}, {dt[1].dt * 1E3:.1f}ms")
+            pred = F.softmax(results, dim=1)  # probabilities
+
+        # Process predictions
+        for i, prob in enumerate(pred):  # per image
+            seen += 1
+            if webcam:  # batch_size >= 1
+                p, im0 = path[i], im0s[i].copy()
+                s += f'{i}: '
+            else:
+                p, im0 = path, im0s.copy()
+
+            p = Path(p)  # to Path
+            save_path = str(save_dir / p.name)  # im.jpg
+            s += '%gx%g ' % im.shape[2:]  # print string
+            annotator = Annotator(im0, example=str(names), pil=True)
+
+            # Print results
+            top5i = prob.argsort(0, descending=True)[:5].tolist()  # top 5 indices
+            s += f"{', '.join(f'{names[j]} {prob[j]:.2f}' for j in top5i)}, "
+
+            # Write results
+            if save_img or view_img:  # Add bbox to image
+                text = '\n'.join(f'{prob[j]:.2f} {names[j]}' for j in top5i)
+                annotator.text((64, 64), text, txt_color=(255, 255, 255))
+
+            # 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}{dt[1].dt * 1E3:.1f}ms")
 
     # Print results
-    t = tuple(x.t / (seen + 1) * 1E3 for x in dt)  # speeds per image
-    shape = (1, 3, imgsz, imgsz)
-    LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms post-process per image at shape {shape}' % t)
-    LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
-    return p
+    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():
     parser = argparse.ArgumentParser()
     parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-cls.pt', help='model path(s)')
-    parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob')
-    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=224, help='train, val image size (pixels)')
+    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=[224], help='inference size h,w')
     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('--nosave', action='store_true', help='do not save images/videos')
+    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/predict-cls', 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('--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('--project', default=ROOT / 'runs/predict-cls', help='save to project/name')
-    parser.add_argument('--name', default='exp', help='save to project/name')
-    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
     opt = parser.parse_args()
+    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
     print_args(vars(opt))
     return opt
 

detect.py

@@ -1,6 +1,6 @@
 # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
 """
-Run YOLOv5 detection inference on images, videos, directories, streams, etc.
+Run YOLOv5 detection inference on images, videos, directories, globs, YouTube, webcam, streams, etc.
 
 Usage - sources:
     $ python detect.py --weights yolov5s.pt --source 0                               # webcam

utils/augmentations.py

@@ -344,4 +344,5 @@ def classify_albumentations(augment=True,
 
 def classify_transforms(size=224):
     # Transforms to apply if albumentations not installed
+    assert isinstance(size, int), f'ERROR: classify_transforms size {size} must be integer, not (list, tuple)'
     return T.Compose([T.ToTensor(), T.Resize(size), T.CenterCrop(size), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])