glenn-jocher commited on
Commit
ad05e37
·
unverified ·
1 Parent(s): 1620669

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

Browse files

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)

Files changed (1) hide show
  1. detect.py +11 -8
detect.py CHANGED
@@ -17,6 +17,7 @@ from utils.torch_utils import select_device, load_classifier, time_synchronized
17
 
18
  def detect(save_img=False):
19
  source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
 
20
  webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
21
  ('rtsp://', 'rtmp://', 'http://'))
22
 
@@ -49,7 +50,6 @@ def detect(save_img=False):
49
  cudnn.benchmark = True # set True to speed up constant image size inference
50
  dataset = LoadStreams(source, img_size=imgsz, stride=stride)
51
  else:
52
- save_img = True
53
  dataset = LoadImages(source, img_size=imgsz, stride=stride)
54
 
55
  # Get names and colors
@@ -124,17 +124,19 @@ def detect(save_img=False):
124
  if save_img:
125
  if dataset.mode == 'image':
126
  cv2.imwrite(save_path, im0)
127
- else: # 'video'
128
  if vid_path != save_path: # new video
129
  vid_path = save_path
130
  if isinstance(vid_writer, cv2.VideoWriter):
131
  vid_writer.release() # release previous video writer
132
-
133
- fourcc = 'mp4v' # output video codec
134
- fps = vid_cap.get(cv2.CAP_PROP_FPS)
135
- w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
136
- h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
137
- vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
 
 
138
  vid_writer.write(im0)
139
 
140
  if save_txt or save_img:
@@ -155,6 +157,7 @@ if __name__ == '__main__':
155
  parser.add_argument('--view-img', action='store_true', help='display results')
156
  parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
157
  parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
 
158
  parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
159
  parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
160
  parser.add_argument('--augment', action='store_true', help='augmented inference')
 
17
 
18
  def detect(save_img=False):
19
  source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
20
+ save_img = not opt.nosave and not source.endswith('.txt') # save inference images
21
  webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
22
  ('rtsp://', 'rtmp://', 'http://'))
23
 
 
50
  cudnn.benchmark = True # set True to speed up constant image size inference
51
  dataset = LoadStreams(source, img_size=imgsz, stride=stride)
52
  else:
 
53
  dataset = LoadImages(source, img_size=imgsz, stride=stride)
54
 
55
  # Get names and colors
 
124
  if save_img:
125
  if dataset.mode == 'image':
126
  cv2.imwrite(save_path, im0)
127
+ else: # 'video' or 'stream'
128
  if vid_path != save_path: # new video
129
  vid_path = save_path
130
  if isinstance(vid_writer, cv2.VideoWriter):
131
  vid_writer.release() # release previous video writer
132
+ if vid_cap: # video
133
+ fps = vid_cap.get(cv2.CAP_PROP_FPS)
134
+ w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
135
+ h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
136
+ else: # stream
137
+ fps, w, h = 30, im0.shape[1], im0.shape[0]
138
+ save_path += '.mp4'
139
+ vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
140
  vid_writer.write(im0)
141
 
142
  if save_txt or save_img:
 
157
  parser.add_argument('--view-img', action='store_true', help='display results')
158
  parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
159
  parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
160
+ parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
161
  parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
162
  parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
163
  parser.add_argument('--augment', action='store_true', help='augmented inference')