From a9d398e812d50380f976e4c19f3edcd50ffd4101 Mon Sep 17 00:00:00 2001 From: Oleg Boiko Date: Mon, 19 Oct 2020 16:38:47 -0700 Subject: [PATCH 1/5] Adding save_conf options for detect.py and test.py --- detect.py | 10 +++++++--- test.py | 16 +++++++++++++--- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/detect.py b/detect.py index a98fe7394854..9d94fa97738e 100644 --- a/detect.py +++ b/detect.py @@ -18,8 +18,8 @@ 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 + out, source, weights, view_img, save_txt, imgsz, save_conf = \ + opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, opt.save_conf webcam = source.isnumeric() or source.startswith(('rtsp://', 'rtmp://', 'http://')) or source.endswith('.txt') # Initialize @@ -105,7 +105,10 @@ def detect(save_img=False): if save_txt: # Write to file xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh with open(txt_path + '.txt', 'a') as f: - f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format + if save_conf: + f.write(('%g ' * 6 + '\n') % (cls, conf, *xywh)) # label format includes conf + else: + f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format does not include conf if save_img or view_img: # Add bbox to image label = '%s %.2f' % (names[int(cls)], conf) @@ -158,6 +161,7 @@ def detect(save_img=False): parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--update', action='store_true', help='update all models') + parser.add_argument('--save-conf', action='store_false', help='put confidence score next to class in label*.txt') opt = parser.parse_args() print(opt) diff --git a/test.py b/test.py index 9e79a769f884..942986b71178 100644 --- a/test.py +++ b/test.py @@ -32,6 +32,7 @@ def test(data, dataloader=None, save_dir=Path(''), # for saving images save_txt=False, # for auto-labelling + save_conf=False, plots=True): # Initialize/load model and set device training = model is not None @@ -133,7 +134,10 @@ def test(data, for *xyxy, conf, cls in x: xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh with open(str(out / Path(paths[si]).stem) + '.txt', 'a') as f: - f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format + if save_conf: + f.write(('%g ' * 6 + '\n') % (cls, conf, *xywh)) # label format includes conf + else: + f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format does not include conf # Clip boxes to image bounds clip_coords(pred, (height, width)) @@ -262,7 +266,9 @@ def test(data, parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset') parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--verbose', action='store_true', help='report mAP by class') - parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') + parser.add_argument('--save-txt', action='store_false', help='save results to *.txt') + parser.add_argument('--save-conf', action='store_false', help='put confidence score next to class in label*.txt') + parser.add_argument('--output', type=str, default='', help='output folder') # output folder opt = parser.parse_args() opt.save_json |= opt.data.endswith('coco.yaml') opt.data = check_file(opt.data) # check file @@ -278,7 +284,11 @@ def test(data, opt.save_json, opt.single_cls, opt.augment, - opt.verbose) + opt.verbose, + save_dir=Path(opt.output), + save_txt=opt.save_txt, + save_conf=opt.save_conf, + ) elif opt.task == 'study': # run over a range of settings and save/plot for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: From 8bed5297b371eb57ea444f7de23deca846a54ae8 Mon Sep 17 00:00:00 2001 From: Oleg Boiko Date: Mon, 19 Oct 2020 17:17:03 -0700 Subject: [PATCH 2/5] store_false -> store_true --- detect.py | 2 +- test.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/detect.py b/detect.py index 9d94fa97738e..0de49714449e 100644 --- a/detect.py +++ b/detect.py @@ -161,7 +161,7 @@ def detect(save_img=False): parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--update', action='store_true', help='update all models') - parser.add_argument('--save-conf', action='store_false', help='put confidence score next to class in label*.txt') + parser.add_argument('--save-conf', action='store_true', help='put confidence score next to class in label*.txt') opt = parser.parse_args() print(opt) diff --git a/test.py b/test.py index 942986b71178..0944f38969e5 100644 --- a/test.py +++ b/test.py @@ -266,8 +266,8 @@ def test(data, parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset') parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--verbose', action='store_true', help='report mAP by class') - parser.add_argument('--save-txt', action='store_false', help='save results to *.txt') - parser.add_argument('--save-conf', action='store_false', help='put confidence score next to class in label*.txt') + parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') + parser.add_argument('--save-conf', action='store_true', help='put confidence score next to class in label*.txt') parser.add_argument('--output', type=str, default='', help='output folder') # output folder opt = parser.parse_args() opt.save_json |= opt.data.endswith('coco.yaml') From 2076ba52cdbac3da68bc4e484e93e2a61292b4a0 Mon Sep 17 00:00:00 2001 From: Oleg Boiko Date: Mon, 19 Oct 2020 17:24:28 -0700 Subject: [PATCH 3/5] Using save_dir to store *.txt files --- test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test.py b/test.py index 0944f38969e5..bb5190cced1a 100644 --- a/test.py +++ b/test.py @@ -44,7 +44,7 @@ def test(data, device = select_device(opt.device, batch_size=batch_size) save_txt = opt.save_txt # save *.txt labels if save_txt: - out = Path('inference/output') + out = save_dir / 'output' if os.path.exists(out): shutil.rmtree(out) # delete output folder os.makedirs(out) # make new output folder From bdac9d9fa917ef019b3df82c88fe0e81a78bc1a9 Mon Sep 17 00:00:00 2001 From: Oleg Boiko Date: Tue, 20 Oct 2020 08:44:56 -0700 Subject: [PATCH 4/5] Removing duplicate --save-conf flag --- detect.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/detect.py b/detect.py index 28ae355e2022..807dee3cce35 100644 --- a/detect.py +++ b/detect.py @@ -18,8 +18,8 @@ def detect(save_img=False): - out, source, weights, view_img, save_txt, imgsz, save_conf = \ - opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, opt.save_conf + out, source, weights, view_img, save_txt, imgsz = \ + opt.output, 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 @@ -160,7 +160,6 @@ def detect(save_img=False): parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--update', action='store_true', help='update all models') - parser.add_argument('--save-conf', action='store_true', help='put confidence score next to class in label*.txt') opt = parser.parse_args() print(opt) From 7ca8c5c717850e5ccaf468999c098517d954e3be Mon Sep 17 00:00:00 2001 From: Oleg Boiko Date: Tue, 20 Oct 2020 08:47:01 -0700 Subject: [PATCH 5/5] Aligning --save-conf code between test.py and detect.py --- test.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/test.py b/test.py index bb5190cced1a..597b56804c0f 100644 --- a/test.py +++ b/test.py @@ -133,11 +133,9 @@ def test(data, x[:, :4] = scale_coords(img[si].shape[1:], x[:, :4], shapes[si][0], shapes[si][1]) # to original for *xyxy, conf, cls in x: xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh + line = (cls, conf, *xywh) if save_conf else (cls, *xywh) # label format with open(str(out / Path(paths[si]).stem) + '.txt', 'a') as f: - if save_conf: - f.write(('%g ' * 6 + '\n') % (cls, conf, *xywh)) # label format includes conf - else: - f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format does not include conf + f.write(('%g ' * len(line) + '\n') % line) # Clip boxes to image bounds clip_coords(pred, (height, width))