From bc16152fccf7860f68c1b822cbfdb3e1e74c0e6e Mon Sep 17 00:00:00 2001 From: oakhtar147 Date: Fri, 18 Sep 2020 20:48:50 +0500 Subject: [PATCH] added feature to save conf in txt labels --- detect.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/detect.py b/detect.py index fd2c14d1ba2e..7072aac17ab7 100644 --- a/detect.py +++ b/detect.py @@ -19,8 +19,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') or source.startswith('http') or source.endswith('.txt') # Initialize @@ -106,7 +106,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) @@ -161,6 +164,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_true', help='put confidence score next to class in label*.txt') opt = parser.parse_args() print(opt)