zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Save webcam results, add --nosave option (#2598) 

This updates the default detect.py behavior to automatically save all inference images/videos/webcams unless the new argument --nosave is used (python detect.py --nosave) or unless a list of streaming sources is passed (python detect.py --source streams.txt)
This commit is contained in:
Glenn Jocher 2021-03-25 14:09:49 +01:00 committed by GitHub
parent 16206692f2
commit ad05e37d99
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 11 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
detect.py
View File

@ -17,6 +17,7 @@ from utils.torch_utils import select_device, load_classifier, time_synchronized
def detect(save_img=False): def detect(save_img=False):
source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
save_img = not opt.nosave and not source.endswith('.txt') # save inference images
webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith( webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
('rtsp://', 'rtmp://', 'http://')) ('rtsp://', 'rtmp://', 'http://'))
@ -49,7 +50,6 @@ def detect(save_img=False):
cudnn.benchmark = True # set True to speed up constant image size inference cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride) dataset = LoadStreams(source, img_size=imgsz, stride=stride)
else: else:
save_img = True
dataset = LoadImages(source, img_size=imgsz, stride=stride) dataset = LoadImages(source, img_size=imgsz, stride=stride)
# Get names and colors # Get names and colors
@ -124,17 +124,19 @@ def detect(save_img=False):
if save_img: if save_img:
if dataset.mode == 'image': if dataset.mode == 'image':
cv2.imwrite(save_path, im0) cv2.imwrite(save_path, im0)
else: # 'video' else: # 'video' or 'stream'
if vid_path != save_path: # new video if vid_path != save_path: # new video
vid_path = save_path vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter): if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer vid_writer.release() # release previous video writer
if vid_cap: # video
fourcc = 'mp4v' # output video codec fps = vid_cap.get(cv2.CAP_PROP_FPS)
fps = vid_cap.get(cv2.CAP_PROP_FPS) w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) else: # stream
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)) 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.write(im0)
if save_txt or save_img: if save_txt or save_img:
@ -155,6 +157,7 @@ if __name__ == '__main__':
parser.add_argument('--view-img', action='store_true', help='display results') parser.add_argument('--view-img', action='store_true', help='display results')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') 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-conf', action='store_true', help='save confidences in --save-txt labels')
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3') parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') 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('--augment', action='store_true', help='augmented inference')