Allow to choose arbitrary i for BC(i)

Author: TimoK93

ctc_metrics/scripts/evaluate.py

@@ -2,6 +2,7 @@
 from os.path import join, basename
 from multiprocessing import Pool, cpu_count
 import numpy as np
+import warnings
 
 from ctc_metrics.metrics import (
     valid, det, seg, tra, ct, tf, bc, raw_division_metrics, cca, mota, hota, idf1, chota, mtml, faf,
@@ -171,6 +172,19 @@ def calculate_metrics(
         traj["labels_comp_merged"] = new_labels
         traj["mapped_comp_merged"] = new_mapped
 
+    # Check if a manual i was defined for BC(i)
+    max_i_for_bci = 3
+    for m in metrics:
+        if m.startswith("BC("):
+            try:
+                i = int(m[3:-1])
+                if i > max_i_for_bci:
+                    max_i_for_bci = i
+                if "BC" not in metrics:
+                    metrics.append("BC")
+            except ValueError:
+                warnings.warn(f"{m} is not a valid metric identifier!.")
+
     # Prepare intermediate results
     graph_operations = {}
     if "DET" in metrics or "TRA" in metrics:
@@ -225,7 +239,7 @@ def calculate_metrics(
             traj["labels_ref"], traj["mapped_ref"], traj["mapped_comp"])
 
     if "BC" in metrics:
-        for i in range(4):
+        for i in range(max_i_for_bci+1):
             tp, fp, fn = raw_division_metrics(comp_tracks, ref_tracks,
                 traj["mapped_ref"], traj["mapped_comp"],
                 i=i)
@@ -243,13 +257,13 @@ def calculate_metrics(
 
     if "CT" in metrics and "BC" in metrics and \
             "CCA" in metrics and "TF" in metrics:
-        for i in range(4):
+        for i in range(max_i_for_bci+1):
             results[f"BIO({i})"] = bio(
                 results["CT"], results["TF"],
                 results[f"BC({i})"], results["CCA"])
 
     if "BIO" in results and "LNK" in results:
-        for i in range(4):
+        for i in range(max_i_for_bci+1):
             results[f"OP_CLB({i})"] = op_clb(
                 results["LNK"], results[f"BIO({i})"])
 
@@ -365,7 +379,7 @@ def parse_args():
     parser.add_argument('--tra', action="store_true")
     parser.add_argument('--ct', action="store_true")
     parser.add_argument('--tf', action="store_true")
-    parser.add_argument('--bc', action="store_true")
+    parser.add_argument('--bc', type=int, default=0)
     parser.add_argument('--cca', action="store_true")
     parser.add_argument('--mota', action="store_true")
     parser.add_argument('--hota', action="store_true")
@@ -391,7 +405,7 @@ def main():
         ("TRA", args.tra),
         ("CT", args.ct),
         ("TF", args.tf),
-        ("BC", args.bc),
+        (f"BC({args.bc})", args.bc),
         ("CCA", args.cca),
         ("MOTA", args.mota),
         ("HOTA", args.hota),