diff --git a/detect.py b/detect.py index 808f3584..a4542f7e 100644 --- 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 ''