Skip to content

Commit

Permalink
Add confidence score to label txt files (ultralytics#994)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
oakhtar147 and glenn-jocher committed Oct 20, 2020
1 parent 8457594 commit a72d3fb
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion detect.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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')
Expand Down

0 comments on commit a72d3fb

Please sign in to comment.