Merge pull request #149 from computational-cell-analytics/generalist-experiments

Support RGB data in micro_sam.evaluation.model_comparison

Author: constantinpape

finetuning/generalists/generate_model_comparison.py

@@ -4,7 +4,7 @@
 import micro_sam.evaluation.model_comparison as comparison
 import torch_em
 
-from util import get_data_paths, EM_DATASETS
+from util import get_data_paths, EM_DATASETS, LM_DATASETS
 
 OUTPUT_ROOT = "/scratch-grete/projects/nim00007/sam/experiments/model_comparison"
 
@@ -15,44 +15,67 @@ def _get_patch_shape(path):
     return patch_shape
 
 
+def raw_trafo(raw):
+    raw = raw.transpose((2, 0, 1))
+    print(raw.shape)
+    return raw
+
+
 def get_loader(dataset):
     image_paths, gt_paths = get_data_paths(dataset, split="test")
     image_paths, gt_paths = image_paths[:100], gt_paths[:100]
 
+    with_channels = dataset in ("hpa", "lizard")
+
     label_transform = torch_em.transform.label.connected_components
     loader = torch_em.default_segmentation_loader(
         image_paths, None, gt_paths, None,
         batch_size=1, patch_shape=_get_patch_shape(image_paths[0]),
         shuffle=True, n_samples=25, label_transform=label_transform,
+        with_channels=with_channels, is_seg_dataset=not with_channels
     )
     return loader
 
 
 def generate_comparison_for_dataset(dataset, model1, model2):
     output_folder = os.path.join(OUTPUT_ROOT, dataset)
     if os.path.exists(output_folder):
-        return
+        return output_folder
     print("Generate model comparison data for", dataset)
     loader = get_loader(dataset)
     comparison.generate_data_for_model_comparison(loader, output_folder, model1, model2, n_samples=25)
+    return output_folder
 
 
-# TODO
-def create_comparison_images():
-    pass
+def create_comparison_images(output_folder, dataset):
+    plot_folder = os.path.join(OUTPUT_ROOT, "images", dataset)
+    if os.path.exists(plot_folder):
+        return
+    comparison.model_comparison(
+        output_folder, n_images_per_sample=25, min_size=100,
+        plot_folder=plot_folder, outline_dilation=1
+    )
 
 
 def generate_comparison_em():
     model1 = "vit_h"
     model2 = "vit_h_em"
     for dataset in EM_DATASETS:
-        generate_comparison_for_dataset(dataset, model1, model2)
-        create_comparison_images()
+        folder = generate_comparison_for_dataset(dataset, model1, model2)
+        create_comparison_images(folder, dataset)
+
+
+def generate_comparison_lm():
+    model1 = "vit_h"
+    model2 = "vit_h_lm"
+    for dataset in LM_DATASETS:
+        folder = generate_comparison_for_dataset(dataset, model1, model2)
+        create_comparison_images(folder, dataset)
 
 
 def main():
-    # generate_comparison_lm()
-    generate_comparison_em()
+    generate_comparison_lm()
+    # generate_comparison_em()
 
 
 if __name__ == "__main__":

micro_sam/evaluation/model_comparison.py

@@ -35,6 +35,9 @@ def _predict_models_with_loader(loader, n_samples, prompt_generator, predictor1,
         out_path = os.path.join(output_folder, f"sample_{i}.h5")
 
         im = x.numpy().squeeze()
+        if im.ndim == 3 and im.shape[0] == 3:
+            im = im.transpose((1, 2, 0))
+
         gt = y.numpy().squeeze().astype("uint32")
         gt = relabel_sequential(gt)[0]
 
@@ -152,12 +155,13 @@ def _evaluate_samples(f, prefix, min_size):
 
 
 def _overlay_mask(image, mask):
-    # TODO add support for RGB inptus
-    assert image.ndim == 2
+    assert image.ndim in (2, 3)
     # overlay the mask
-    overlay = np.stack(
-        [image, image, image]
-    ).transpose((1, 2, 0))
+    if image.ndim == 2:
+        overlay = np.stack([image, image, image]).transpose((1, 2, 0))
+    else:
+        overlay = image
+    assert overlay.shape[-1] == 3
     mask_overlay = np.zeros_like(overlay)
     mask_overlay[mask == 1] = [255, 0, 0]
     alpha = 0.6