added cmd line arg for saving only positives

Added a simple command line argument (default False) where if set to true, only the images and videos which contain a positive identification are saved to the output. 

Location in code is not extremely efficient, but is present at that location for readability and minimal intrusiveness.

Signed-off-by: GarbageHaus <90212322+GarbageHaus@users.noreply.github.com>
pull/13005/head
GarbageHaus 2024-05-12 18:48:41 -04:00 committed by GitHub
parent 28e06aa391
commit 1bff0142e5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions

View File

@ -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
@ -233,7 +234,7 @@ def run(
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'
@ -281,6 +282,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")