|
|
@@ -17,6 +17,7 @@ from utils.torch_utils import select_device, load_classifier, time_synchronized |
|
|
|
|
|
|
|
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 |
|
|
|
save_img = not opt.nosave and not source.endswith('.txt') # save inference images |
|
|
|
webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith( |
|
|
|
('rtsp://', 'rtmp://', 'http://')) |
|
|
|
|
|
|
@@ -49,7 +50,6 @@ def detect(save_img=False): |
|
|
|
cudnn.benchmark = True # set True to speed up constant image size inference |
|
|
|
dataset = LoadStreams(source, img_size=imgsz, stride=stride) |
|
|
|
else: |
|
|
|
save_img = True |
|
|
|
dataset = LoadImages(source, img_size=imgsz, stride=stride) |
|
|
|
|
|
|
|
# Get names and colors |
|
|
@@ -124,17 +124,19 @@ def detect(save_img=False): |
|
|
|
if save_img: |
|
|
|
if dataset.mode == 'image': |
|
|
|
cv2.imwrite(save_path, im0) |
|
|
|
else: # 'video' |
|
|
|
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 |
|
|
|
|
|
|
|
fourcc = 'mp4v' # output video codec |
|
|
|
fps = vid_cap.get(cv2.CAP_PROP_FPS) |
|
|
|
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
|
|
|
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
|
|
|
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)) |
|
|
|
if vid_cap: # video |
|
|
|
fps = vid_cap.get(cv2.CAP_PROP_FPS) |
|
|
|
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
|
|
|
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
|
|
|
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) |
|
|
|
|
|
|
|
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('--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('--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('--agnostic-nms', action='store_true', help='class-agnostic NMS') |
|
|
|
parser.add_argument('--augment', action='store_true', help='augmented inference') |