From b2f47854f46b44a7967c4786c1e95daee44cfc0f Mon Sep 17 00:00:00 2001 From: Osama Akhtar <57761675+oakhtar147@users.noreply.github.com> Date: Tue, 20 Oct 2020 17:20:27 +0500 Subject: [PATCH] Add confidence score to label txt files (#994) * add contrast to conf/class against bbox rectangle color * added feature to save conf in txt labels * Update general.py * Update detect.py * Update detect.py * Update detect.py save_conf variable removed as only used once, so we can use opt.save_conf instead. * Update detect.py * Update detect.py Co-authored-by: Glenn Jocher --- detect.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/detect.py b/detect.py index a98fe7394854..807dee3cce35 100644 --- a/detect.py +++ b/detect.py @@ -104,8 +104,9 @@ def detect(save_img=False): for *xyxy, conf, cls in reversed(det): if save_txt: # Write to file xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh + line = (cls, conf, *xywh) if opt.save_conf else (cls, *xywh) # label format with open(txt_path + '.txt', 'a') as f: - f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format + f.write(('%g ' * len(line) + '\n') % line) if save_img or view_img: # Add bbox to image label = '%s %.2f' % (names[int(cls)], conf) @@ -154,6 +155,7 @@ def detect(save_img=False): 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('--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')