diff --git a/utils/metrics.py b/utils/metrics.py index a0ae78b12445..145c9d902e58 100644 --- a/utils/metrics.py +++ b/utils/metrics.py @@ -71,8 +71,8 @@ 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) - names = [v for k, v in names.items() if k in unique_classes] # List: only series that have data - names = {i: v for i, v in enumerate(names)} # back to dict; re-indexed to start at 0 + names = [v for k, v in names.items() if k in unique_classes] # List: only classes that have data + names = {i: v for i, v in enumerate(names)} # to dict 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')