Save webcam results, add --nosave option (#2598)
Browse filesThis 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)
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 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
137 |
-
|
|
|
|
|
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')
|