Merge pull request #90 from janelia-cellmap/json_none-ify

Json none ify

src/cellmap_segmentation_challenge/evaluate.py

@@ -42,9 +42,10 @@
 
 
 HAUSDORFF_DISTANCE_MAX = np.inf
+CAST_TO_NONE = [np.nan, np.inf, -np.inf]
 MAX_MAIN_THREADS = int(os.getenv("MAX_MAIN_THREADS", 2))
-MAX_LABEL_THREADS = int(os.getenv("MAX_LABEL_THREADS", 2))
-MAX_INSTANCE_THREADS = int(os.getenv("MAX_INSTANCE_THREADS", 2))
+MAX_LABEL_THREADS = int(os.getenv("MAX_LABEL_THREADS", 4))
+MAX_INSTANCE_THREADS = int(os.getenv("MAX_INSTANCE_THREADS", 4))
 PRECOMPUTE_LIMIT = int(os.getenv("PRECOMPUTE_LIMIT", 1e9))
 DEBUG = os.getenv("DEBUG", "False") != "False"
 
@@ -64,12 +65,18 @@ def __len__(self):
         return len(self.ids)
 
 
-def score_label_single(label, pred_volume_path, truth_path, instance_classes):
+def score_label_single(
+    label,
+    pred_volume_path,
+    truth_path,
+    instance_classes,
+):
     score = score_label(
         pred_volume_path / label,
         truth_path=truth_path,
         instance_classes=instance_classes,
     )
+
     return (label, score)
 
 
@@ -698,14 +705,20 @@ def missing_volume_score(
     return scores
 
 
-def combine_scores(scores, include_missing=True, instance_classes=INSTANCE_CLASSES):
+def combine_scores(
+    scores,
+    include_missing=True,
+    instance_classes=INSTANCE_CLASSES,
+    cast_to_none=CAST_TO_NONE,
+):
     """
     Combine scores across volumes, normalizing by the number of voxels.
 
     Args:
         scores (dict): A dictionary of scores for each volume, as returned by `score_volume`.
         include_missing (bool): Whether to include missing volumes in the combined scores.
         instance_classes (list): A list of instance classes.
+        cast_to_none (list): A list of values to cast to None in the combined scores.
 
     Returns:
         dict: A dictionary of combined scores across all volumes.
@@ -719,7 +732,7 @@ def combine_scores(scores, include_missing=True, instance_classes=INSTANCE_CLASS
     scores = scores.copy()
     label_scores = {}
     total_volumes = {}
-    for these_scores in scores.values():
+    for ds, these_scores in scores.items():
         for label, this_score in these_scores.items():
             print(this_score)
             if this_score["is_missing"] and not include_missing:
@@ -734,26 +747,17 @@ def combine_scores(scores, include_missing=True, instance_classes=INSTANCE_CLASS
                         "combined_score": 0,
                     }
                     total_volumes[label] = 0
-                label_scores[label]["accuracy"] += this_score["accuracy"] * total_volume
-                label_scores[label]["hausdorff_distance"] += (
-                    this_score["hausdorff_distance"] * total_volume
-                )
-                label_scores[label]["normalized_hausdorff_distance"] += (
-                    this_score["normalized_hausdorff_distance"] * total_volume
-                )
-                label_scores[label]["combined_score"] += (
-                    this_score["combined_score"] * total_volume
-                )
-                total_volumes[label] += total_volume
             else:
                 if label not in label_scores:
                     label_scores[label] = {"iou": 0, "dice_score": 0}
                     total_volumes[label] = 0
-                label_scores[label]["iou"] += this_score["iou"] * total_volume
-                label_scores[label]["dice_score"] += (
-                    this_score["dice_score"] * total_volume
-                )
-                total_volumes[label] += total_volume
+            for key in label_scores[label].keys():
+                if this_score[key] is None:
+                    continue
+                label_scores[label][key] += this_score[key] * total_volume
+                if this_score[key] in cast_to_none:
+                    scores[ds][label][key] = None
+            total_volumes[label] += total_volume
 
     # Normalize back to the total number of voxels
     for label in label_scores:
@@ -765,6 +769,10 @@ def combine_scores(scores, include_missing=True, instance_classes=INSTANCE_CLASS
         else:
             label_scores[label]["iou"] /= total_volumes[label]
             label_scores[label]["dice_score"] /= total_volumes[label]
+        # Cast to None if the value is in `cast_to_none`
+        for key in label_scores[label]:
+            if label_scores[label][key] in cast_to_none:
+                label_scores[label][key] = None
     scores["label_scores"] = label_scores
 
     # Compute the overall score
@@ -921,14 +929,14 @@ def score_submission(
     if result_file:
         print(f"Saving collected scores to {result_file}...")
         with open(result_file, "w") as f:
-            json.dump(all_scores, f)
+            json.dump(all_scores, f, indent=4)
 
         found_result_file = str(result_file).replace(
             UPath(result_file).suffix, "_submitted_only" + UPath(result_file).suffix
         )
         print(f"Saving scores for only submitted data to {found_result_file}...")
         with open(found_result_file, "w") as f:
-            json.dump(found_scores, f)
+            json.dump(found_scores, f, indent=4)
         print(
             f"Scores saved to {result_file} and {found_result_file} in {time() - start_time:.2f} seconds"
         )