From 14e2d38959e6a504a2b6519c09edd216b538e400 Mon Sep 17 00:00:00 2001 From: ketan-b <54092325+ketan-b@users.noreply.github.com> Date: Sat, 3 Jul 2021 00:38:57 +0530 Subject: [PATCH 1/4] 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! --- detect.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/detect.py b/detect.py index 808f3584c93d..ce496e0873bc 100644 --- a/detect.py +++ b/detect.py @@ -77,7 +77,10 @@ def run(weights='yolov5s.pt', # model.pt path(s) modelc.load_state_dict(torch.load('resnet50.pt', map_location=device)['model']).to(device).eval() # Set Dataloader - vid_path, vid_writer = None, None + if source.endswith('.txt'): + with open(source, 'r') as f: + sources_len = len([x.strip() for x in f.read().strip().splitlines() if len(x.strip())]) + vid_path, vid_writer = [None]*sources_len, [None]*sources_len if webcam: view_img = check_imshow() cudnn.benchmark = True # set True to speed up constant image size inference @@ -156,12 +159,12 @@ def run(weights='yolov5s.pt', # model.pt path(s) # Save results (image with detections) if save_img: if dataset.mode == 'image': - cv2.imwrite(save_path, im0) + 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 +172,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 '' From ac92ea596a0b793fde4c45a4dd808767ab70c26e Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 3 Jul 2021 20:30:09 +0200 Subject: [PATCH 2/4] Cleanup list lengths --- detect.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/detect.py b/detect.py index ce496e0873bc..192c5dd2d209 100644 --- a/detect.py +++ b/detect.py @@ -77,16 +77,14 @@ def run(weights='yolov5s.pt', # model.pt path(s) modelc.load_state_dict(torch.load('resnet50.pt', map_location=device)['model']).to(device).eval() # Set Dataloader - if source.endswith('.txt'): - with open(source, 'r') as f: - sources_len = len([x.strip() for x in f.read().strip().splitlines() if len(x.strip())]) - vid_path, vid_writer = [None]*sources_len, [None]*sources_len 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) + vid_path, vid_writer = [None] * len(dataset.sources), [None] * len(dataset.sources) else: dataset = LoadImages(source, img_size=imgsz, stride=stride) + vid_path, vid_writer = None, None # Run inference if device.type != 'cpu': @@ -159,7 +157,7 @@ def run(weights='yolov5s.pt', # model.pt path(s) # Save results (image with detections) if save_img: if dataset.mode == 'image': - cv2.imwrite(save_path, im0) + cv2.imwrite(save_path, im0) else: # 'video' or 'stream' if vid_path[i] != save_path: # new video vid_path[i] = save_path From 9f50ec3ff7eea672a65ae751312383dd4e9ecef8 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 4 Jul 2021 12:44:47 +0200 Subject: [PATCH 3/4] batch size variable --- detect.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/detect.py b/detect.py index 192c5dd2d209..a4542f7e8802 100644 --- a/detect.py +++ b/detect.py @@ -76,15 +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 + # 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) - vid_path, vid_writer = [None] * len(dataset.sources), [None] * len(dataset.sources) + bs = len(dataset) # batch_size else: dataset = LoadImages(source, img_size=imgsz, stride=stride) - vid_path, vid_writer = None, None + bs = 1 # batch_size + vid_path, vid_writer = [None] * bs, [None] * bs # Run inference if device.type != 'cpu': From c1c47ecc831f2725559abb7ed8a151e101e31c52 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 4 Jul 2021 12:46:55 +0200 Subject: [PATCH 4/4] Update datasets.py --- utils/datasets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/datasets.py b/utils/datasets.py index f7315522e375..8560f7cfeb88 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -352,7 +352,7 @@ def __next__(self): 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):