diff --git a/detect.py b/detect.py index f791faa0908..474a4570f05 100644 --- a/detect.py +++ b/detect.py @@ -82,6 +82,7 @@ def run( save_conf=False, # save confidences in --save-txt labels save_crop=False, # save cropped prediction boxes nosave=False, # do not save images/videos + savepositives=False, # save only positive detections classes=None, # filter by class: --class 0, or --class 0 2 3 agnostic_nms=False, # class-agnostic NMS augment=False, # augmented inference @@ -234,7 +235,7 @@ def write_to_csv(image_name, prediction, confidence): cv2.waitKey(1) # 1 millisecond # Save results (image with detections) - if save_img: + if save_img and ((savepositives and len(det)) or not savepositives): if dataset.mode == "image": cv2.imwrite(save_path, im0) else: # 'video' or 'stream' @@ -282,6 +283,7 @@ def parse_opt(): parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels") parser.add_argument("--save-crop", action="store_true", help="save cropped prediction boxes") parser.add_argument("--nosave", action="store_true", help="do not save images/videos") + parser.add_argument("--savepositives", action="store_true", help="save files only with a positive detection") parser.add_argument("--classes", nargs="+", type=int, help="filter by class: --classes 0, or --classes 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")