Add multi-stream saving feature (#3864)

* Added the recording feature for multiple streams Thanks for the very cool repo!! I was trying to record multiple feeds at the same time, but the current version of the detector only had one video writer and one vid_path! So the streams were not being saved and only were initialized with one frame and this process didn't record the whole thing. Fix: I made a list of `vid_writer` and `vid_path` and the `i` from the loop over the `pred` took care of the writer which need to work! I hope this helps, Thanks! * Cleanup list lengths * batch size variable * Update datasets.py Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

Add multi-stream saving feature (#3864)
9d86b54e · ketan-b · GitHub · d3e9d698 · 9d86b54e · 9d86b54e
--- a/detect.py
+++ b/detect.py
@@ -76,14 +76,16 @@ def run(weights='yolov5s.pt',  # model.pt path(s)
        modelc = load_classifier(name='resnet50', n=2)  # initialize
        modelc.load_state_dict(torch.load('resnet50.pt', map_location=device)['model']).to(device).eval()

-    # Set Dataloader
-    vid_path, vid_writer = None, None
+    # 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, stride=stride)
+        bs = len(dataset)  # batch_size
    else:
        dataset = LoadImages(source, img_size=imgsz, stride=stride)
+        bs = 1  # batch_size
+    vid_path, vid_writer = [None] * bs, [None] * bs

    # Run inference
    if device.type != 'cpu':
@@ -158,10 +160,10 @@ def run(weights='yolov5s.pt',  # model.pt path(s)
                if dataset.mode == 'image':
                    cv2.imwrite(save_path, im0)
                else:  # 'video' or 'stream'
-                    if vid_path != save_path:  # new video
-                        vid_path = save_path
-                        if isinstance(vid_writer, cv2.VideoWriter):
-                            vid_writer.release()  # release previous video writer
+                    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))
@@ -169,8 +171,8 @@ def run(weights='yolov5s.pt',  # model.pt path(s)
                        else:  # stream
                            fps, w, h = 30, im0.shape[1], im0.shape[0]
                            save_path += '.mp4'
-                        vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
-                    vid_writer.write(im0)
+                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+                    vid_writer[i].write(im0)

    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 ''

--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -352,7 +352,7 @@ class LoadStreams:  # multiple IP or RTSP cameras
        return self.sources, img, img0, None

    def __len__(self):
-        return 0  # 1E12 frames = 32 streams at 30 FPS for 30 years
+        return len(self.sources)  # 1E12 frames = 32 streams at 30 FPS for 30 years


 def img2label_paths(img_paths):