Tian Wang commited on
Commit
7e1047a
·
1 Parent(s): f746e93

Show detection and classification confidence on labels

Browse files
Files changed (1) hide show
  1. src/inference/solve.py +4 -1
src/inference/solve.py CHANGED
@@ -373,7 +373,10 @@ class SetSolver:
373
  color = SET_COLORS[color_idx % len(SET_COLORS)]
374
  draw.rectangle([x1, y1, x2, y2], outline=color, width=4)
375
 
376
- label = card_to_chinese(attrs)
 
 
 
377
  draw.text((x1, y1 - 20), label, fill=color, font=font)
378
 
379
  # Draw Set info
 
373
  color = SET_COLORS[color_idx % len(SET_COLORS)]
374
  draw.rectangle([x1, y1, x2, y2], outline=color, width=4)
375
 
376
+ det_conf = c["detection"]["confidence"]
377
+ cls_conf = min(attrs.get("number_conf", 1), attrs.get("color_conf", 1),
378
+ attrs.get("shape_conf", 1), attrs.get("fill_conf", 1))
379
+ label = f"{card_to_chinese(attrs)} d{det_conf:.0%} c{cls_conf:.0%}"
380
  draw.text((x1, y1 - 20), label, fill=color, font=font)
381
 
382
  # Draw Set info