legend-labels Adjust legend labels for classes without instances.

NauchtanRobotics · Oct 14, 2021 · 1bfd3d7 · 1bfd3d7
1 parent d0bfeb3
commit 1bfd3d7
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/utils/metrics.py b/utils/metrics.py
@@ -71,6 +71,9 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names
 
     # Compute F1 (harmonic mean of precision and recall)
     f1 = 2 * p * r / (p + r + 1e-16)
+    # Adjust legend labels for gaps
+    names = {k: v for k, v in names.items() if k in unique_classes}
+    names = {i: v for i, v in enumerate(names.values())}
     if plot:
         plot_pr_curve(px, py, ap, Path(save_dir) / 'PR_curve.png', names)
         plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1')