Update generalist experiments

Author: constantinpape

finetuning/generalists/cellpose_baseline.py

@@ -0,0 +1,124 @@
+import argparse
+import os
+from glob import glob
+from subprocess import run
+
+import imageio.v3 as imageio
+
+from tqdm import tqdm
+
+DATA_ROOT = "/scratch/projects/nim00007/sam/datasets"
+EXP_ROOT = "/scratch/projects/nim00007/sam/experiments/cellpose"
+
+DATASETS = (
+    "covid-if",
+    "deepbacs",
+    "hpa",
+    "livecell",
+    "lizard",
+    "mouse-embryo",
+    "plantseg-ovules",
+    "plantseg-root",
+    "tissuenet",
+)
+
+
+def load_cellpose_model():
+    from cellpose import models
+
+    device, gpu = models.assign_device(True, True)
+    model = models.Cellpose(gpu=gpu, model_type="cyto", device=device)
+    return model
+
+
+def run_cellpose_segmentation(datasets, job_id):
+    dataset = datasets[job_id]
+    experiment_folder = os.path.join(EXP_ROOT, dataset)
+
+    prediction_folder = os.path.join(experiment_folder, "predictions")
+    os.makedirs(prediction_folder, exist_ok=True)
+
+    image_paths = sorted(glob(os.path.join(DATA_ROOT, dataset, "test", "image*.tif")))
+    model = load_cellpose_model()
+
+    for path in tqdm(image_paths, desc=f"Segmenting {dataset} with cellpose"):
+        fname = os.path.basename(path)
+        out_path = os.path.join(prediction_folder, fname)
+        if os.path.exists(out_path):
+            continue
+        image = imageio.imread(path)
+        if image.ndim == 3:
+            assert image.shape[-1] == 3
+            image = image.mean(axis=-1)
+        assert image.ndim == 2
+        seg = model.eval(image, diameter=None, flow_threshold=None, channels=[0, 0])[0]
+        assert seg.shape == image.shape
+        imageio.imwrite(out_path, seg, compression=5)
+
+
+def submit_array_job(datasets):
+    n_datasets = len(datasets)
+    cmd = ["sbatch", "-a", f"0-{n_datasets-1}", "cellpose_baseline.sbatch"]
+    run(cmd)
+
+
+def evaluate_dataset(dataset):
+    from micro_sam.evaluation.evaluation import run_evaluation
+
+    gt_paths = sorted(glob(os.path.join(DATA_ROOT, dataset, "test", "label*.tif")))
+    experiment_folder = os.path.join(EXP_ROOT, dataset)
+    pred_paths = sorted(glob(os.path.join(experiment_folder, "predictions", "*.tif")))
+    assert len(gt_paths) == len(pred_paths), f"{len(gt_paths)}, {len(pred_paths)}"
+    result_path = os.path.join(experiment_folder, "cellpose.csv")
+    run_evaluation(gt_paths, pred_paths, result_path)
+
+
+def evaluate_segmentations(datasets):
+    for dataset in datasets:
+        # we skip livecell, which has already been processed by cellpose
+        if dataset == "livecell":
+            continue
+        evaluate_dataset(dataset)
+
+
+def check_results(datasets):
+    for ds in datasets:
+        # we skip livecell, which has already been processed by cellpose
+        if ds == "livecell":
+            continue
+        result_path = os.path.join(EXP_ROOT, ds, "cellpose.csv")
+        if not os.path.exists(result_path):
+            print("Cellpose results missing for", ds)
+    print("All checks passed")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--segment", "-s", action="store_true")
+    parser.add_argument("--evaluate", "-e", action="store_true")
+    parser.add_argument("--check", "-c", action="store_true")
+    parser.add_argument("--datasets", nargs="+")
+    args = parser.parse_args()
+
+    job_id = os.environ.get("SLURM_ARRAY_TASK_ID", None)
+
+    if args.datasets is None:
+        datasets = DATASETS
+    else:
+        datasets = args.datasets
+        assert all(ds in DATASETS for ds in datasets)
+
+    if job_id is not None:
+        run_cellpose_segmentation(datasets, int(job_id))
+    elif args.segment:
+        submit_array_job(datasets)
+    elif args.evaluate:
+        evaluate_segmentations(datasets)
+    elif args.check:
+        check_results(datasets)
+    else:
+        raise ValueError("Doing nothing")
+
+
+if __name__ == "__main__":
+    main()

finetuning/generalists/cellpose_baseline.sbatch

@@ -0,0 +1,10 @@
+#! /bin/bash
+#SBATCH -c 4
+#SBATCH --mem 48G
+#SBATCH -t 300
+#SBATCH -p grete:shared
+#SBATCH -G A100:1
+#SBATCH -A nim00007
+
+source activate cellpose
+python cellpose_baseline.py $@

finetuning/generalists/compile_results.py

@@ -54,7 +54,6 @@ def compile_em():
     )
 
 
-# TODO
 def compile_lm():
     compile_results(
         ["vit_h", "vit_h_lm", "vit_b", "vit_b_lm"],
@@ -64,7 +63,8 @@ def compile_lm():
 
 
 def main():
-    compile_em()
+    # compile_em()
+    compile_lm()
 
 
 if __name__ == "__main__":

finetuning/generalists/evaluate_generalist.py

@@ -16,6 +16,11 @@
     # Generalist EM models
     "vit_b_em": "/scratch/projects/nim00007/sam/models/EM/generalist/v2/vit_b/best.pt",
     "vit_h_em": "/scratch/projects/nim00007/sam/models/EM/generalist/v2/vit_h/best.pt",
+    # Specialist Models (we don't add livecell, because these results are all computed already)
+    "vit_b_tissuenet": "/scratch/projects/nim00007/sam/models/LM/TissueNet/vit_b/best.pt",
+    "vit_h_tissuenet": "/scratch/projects/nim00007/sam/models/LM/TissueNet/vit_h/best.pt",
+    "vit_b_deepbacs": "/scratch/projects/nim00007/sam/models/LM/DeepBacs/vit_b/best.pt",
+    "vit_h_deepbacs": "/scratch/projects/nim00007/sam/models/LM/DeepBacs/vit_h/best.pt",
 }
 
 

finetuning/generalists/generate_model_comparison.py

@@ -0,0 +1,59 @@
+import os
+
+import imageio.v3 as imageio
+import micro_sam.evaluation.model_comparison as comparison
+import torch_em
+
+from util import get_data_paths, EM_DATASETS
+
+OUTPUT_ROOT = "/scratch-grete/projects/nim00007/sam/experiments/model_comparison"
+
+
+def _get_patch_shape(path):
+    im_shape = imageio.imread(path).shape[:2]
+    patch_shape = tuple(min(sh, 512) for sh in im_shape)
+    return patch_shape
+
+
+def get_loader(dataset):
+    image_paths, gt_paths = get_data_paths(dataset, split="test")
+    image_paths, gt_paths = image_paths[:100], gt_paths[:100]
+
+    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,
+    )
+    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
+    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)
+
+
+# TODO
+def create_comparison_images():
+    pass
+
+
+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()
+
+
+def main():
+    # generate_comparison_lm()
+    generate_comparison_em()
+
+
+if __name__ == "__main__":
+    main()

micro_sam/evaluation/model_comparison.py

@@ -83,7 +83,7 @@ def _predict_models_with_loader(loader, n_samples, prompt_generator, predictor1,
 
 
 def generate_data_for_model_comparison(
-    loader: torch.utils.DataLoader,
+    loader: torch.utils.data.DataLoader,
     output_folder: Union[str, os.PathLike],
     model_type1: str,
     model_type2: str,