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 <glenn.jocher@ultralytics.com>pull/1189/head
parent
f532bc68a9
commit
83deec133d
|
@ -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 @@ if __name__ == '__main__':
|
|||
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')
|
||||
|
|
Loading…
Reference in New Issue