From f2b42620dbc329a3b73e6878705142a0a8d3ea0e Mon Sep 17 00:00:00 2001
From: Jordao Bragantini <jordao.bragantini@czbiohub.org>
Date: Thu, 19 Sep 2024 13:02:37 -0700
Subject: [PATCH] WIP evaluating superpixel results

---
 notebooks/superpixel_stats.py      | 41 ++++++++++++++++++++++
 src/cellcanvas_spp/ground_truth.py | 55 ++++++++++++++++++++++++++----
 2 files changed, 89 insertions(+), 7 deletions(-)
 create mode 100644 notebooks/superpixel_stats.py

diff --git a/notebooks/superpixel_stats.py b/notebooks/superpixel_stats.py
new file mode 100644
index 0000000..12a28e7
--- /dev/null
+++ b/notebooks/superpixel_stats.py
@@ -0,0 +1,41 @@
+# %%
+# The following magic causes the notebook to reload external python modules upon execution of a cell
+# This is useful when developing modules in parallel to the notebook
+
+import copick
+from pathlib import Path
+from cellcanvas_spp.ground_truth import copick_to_ground_truth_image, ground_truth_stats
+
+import numpy as np
+import pandas as pd
+from tifffile import imread
+
+# %%
+DATA_DIR = Path('/Users/jordao.bragantini/Softwares/superpixels/notebooks/data/copick_10439')
+
+root = copick.from_file(DATA_DIR / "synthetic_data_10439_dataportal.json")
+
+runs = ["16193", "16191"]
+
+stats = []
+
+for run in runs:
+    gt = copick_to_ground_truth_image(root, run)
+    segm = imread(DATA_DIR / f"segm_{run}.tif")
+    import napari
+    viewer = napari.view_labels(segm)
+    viewer.add_labels(gt)
+
+    df = ground_truth_stats(segm, gt)
+    df["run"] = run
+    stats.append(df)
+
+df = pd.concat(stats)
+
+
+# %%
+
+print(df.describe())
+
+
+# %%
diff --git a/src/cellcanvas_spp/ground_truth.py b/src/cellcanvas_spp/ground_truth.py
index 0a11304..639965e 100644
--- a/src/cellcanvas_spp/ground_truth.py
+++ b/src/cellcanvas_spp/ground_truth.py
@@ -1,7 +1,11 @@
 import pandas as pd
+import copick
+from copick.impl.filesystem import CopickRoot
 from numpy.typing import ArrayLike
 from skimage.measure import regionprops_table
 import numpy as np
+import zarr
+
 
 def intensity_voxel_counts(
         regionmask: ArrayLike, 
@@ -85,12 +89,19 @@ def get_gt_label_per_super_pixel(row) :
     # Function to from the 8 label counts to single label (the class with most pixels, or background)
 
     counts = row.values
-    # if at least pixel in the superpixel has a gt-label, assign this label (1-7)
-    if np.max(counts[1:])>0:
-        idx = np.argmax(counts[1:])+1
-    # if no gt-label is present in superpixel, assign background (0)
-    else:
-        idx = 0
+
+    weights = np.ones_like(counts)
+    weights[0] = 0.25
+
+   # if at least pixel in the superpixel has a gt-label, assign this label (1-7)
+    idx = np.argmax(counts * weights)
+
+    # if np.max(counts[1:])>0:
+    #     idx = np.argmax(counts[1:])+1
+
+    # # if no gt-label is present in superpixel, assign background (0)
+    # else:
+    #     idx = 0
     return idx
 
 def ground_truth_count(
@@ -144,4 +155,34 @@ def ground_truth_stats(
         "total",
     ]
 
-    return props_df
\ No newline at end of file
+    return props_df
+
+
+def copick_to_ground_truth_image(
+    root: CopickRoot,
+    run_name: str,
+) -> ArrayLike:
+
+    particles = {}
+
+    particles = dict()
+    for po in root.config.pickable_objects:
+        particles[po.name] = po.label
+
+    segmentations = root.get_run(run_name).get_segmentations()
+    ground_truth = None
+
+    for cls_seg in segmentations:
+        _, cls_mask = list(zarr.open(cls_seg.zarr()).arrays())[0]
+        cls_mask = cls_mask[...]
+
+        if ground_truth is None:
+            ground_truth = np.zeros(cls_mask.shape, dtype=int)
+
+        try:
+            ground_truth[cls_mask == 1] = particles[cls_seg.name]
+        except KeyError:
+            print(f"Class {cls_seg.name} not found in particles dictionary")
+            pass
+
+    return ground_truth