Adding --save-dir and --save-conf options to test.py (#1182)
* Adding --output and --save-conf options to test.py * Update help fields * Update test.py * Make arguments and comments uniform with test.py * Remove previous and print save_dir on finish Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
1 parent
0a52ae16dc
commit
19c8b2c9b9
2 files changed
+27
-15
No files matched your search
@@ -19,15 +19,15 @@ from utils.torch_utils import select_device, load_classifier, time_synchronized
|
||||
|
||||
def detect(save_img=False):
|
||||
out, source, weights, view_img, save_txt, imgsz = \
|
||||
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
|
||||
opt.save_dir, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
|
||||
webcam = source.isnumeric() or source.startswith(('rtsp://', 'rtmp://', 'http://')) or source.endswith('.txt')
|
||||
|
||||
# Initialize
|
||||
set_logging()
|
||||
device = select_device(opt.device)
|
||||
if os.path.exists(out):
|
||||
shutil.rmtree(out) # delete output folder
|
||||
os.makedirs(out) # make new output folder
|
||||
if os.path.exists(out): # output dir
|
||||
shutil.rmtree(out) # delete dir
|
||||
os.makedirs(out) # make new dir
|
||||
half = device.type != 'cpu' # half precision only supported on CUDA
|
||||
|
||||
# Load model
|
||||
@@ -148,14 +148,14 @@ if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
|
||||
parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam
|
||||
parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder
|
||||
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
||||
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
|
||||
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
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='output confidences in --save-txt labels')
|
||||
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
|
||||
parser.add_argument('--save-dir', type=str, default='inference/output', help='directory to save results')
|
||||
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')
|
||||
|
||||
Reference in new issue
Block a user