Implement more debugging for detection model

Author: constantinpape

scripts/la-vision/check_detections.py

@@ -0,0 +1,61 @@
+import napari
+import zarr
+
+
+resolution = [3.0, 1.887779, 1.887779]
+positions = [
+    [2002.95539395823, 1899.9032205156411, 264.7747008147759]
+]
+
+
+def _load_from_mobie(bb):
+    path = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/mobie_project/cochlea-lightsheet/LaVision-M04/images/ome-zarr/PV.ome.zarr"
+    f = zarr.open(path, mode="r")
+    data = f["s0"][bb]
+    print(bb)
+
+    path = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/mobie_project/cochlea-lightsheet/LaVision-M04/images/ome-zarr/SGN_detect-v1.ome.zarr"
+    f = zarr.open(path, mode="r")
+    seg = f["s0"][bb]
+
+    return data, seg
+
+
+def _load_prediction(bb):
+    path = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/predictions/LaVision-M04/SGN_detect-v1/predictions.zarr"
+    f = zarr.open(path, mode="r")
+    data = f["prediction"][bb]
+    return data
+
+
+def _load_prediction_debug():
+    path = "./debug-pred/pred-v5.h5"
+    with zarr.open(path, "r") as f:
+        pred = f["pred"][:]
+    return pred
+
+
+def check_detection(position, halo=[32, 384, 384]):
+
+    bb = tuple(
+        slice(int(pos / re) - ha, int(pos / re) + ha) for pos, re, ha in zip(position[::-1], resolution, halo)
+    )
+
+    pv, detections_mobie = _load_from_mobie(bb)
+    # pred = _load_prediction(bb)
+    pred = _load_prediction_debug()
+
+    v = napari.Viewer()
+    v.add_image(pv)
+    v.add_image(pred)
+    v.add_labels(detections_mobie)
+    napari.run()
+
+
+def main():
+    position = positions[0]
+    check_detection(position)
+
+
+if __name__ == "__main__":
+    main()

scripts/la-vision/debug_prediction.py

@@ -0,0 +1,62 @@
+import os
+from functools import partial
+
+import numpy as np
+import torch
+import zarr
+from torch_em.transform.raw import standardize
+from torch_em.util.prediction import predict_with_halo
+
+
+resolution = [3.0, 1.887779, 1.887779]
+positions = [
+    [2002.95539395823, 1899.9032205156411, 264.7747008147759]
+]
+
+
+def _load_from_mobie(bb):
+    path = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/mobie_project/cochlea-lightsheet/LaVision-M04/images/ome-zarr/PV.ome.zarr"
+    f = zarr.open(path, mode="r")
+    data = f["s0"][bb]
+    return data
+
+
+def run_prediction(position, halo=[32, 384, 384]):
+    bb = tuple(
+        slice(int(pos / re) - ha, int(pos / re) + ha) for pos, re, ha in zip(position[::-1], resolution, halo)
+    )
+    pv = _load_from_mobie(bb)
+    mean, std = np.mean(pv), np.std(pv)
+    print(mean, std)
+    preproc = partial(standardize, mean=mean, std=std)
+
+    block_shape = (24, 256, 256)
+    halo = (8, 64, 64)
+
+    model_path = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/trained_models/SGN/sgn-detection-v1.pt"
+    model = torch.load(model_path, weights_only=False)
+
+    def postproc(x):
+        x = np.clip(x, 0, 1)
+        max_ = np.percentile(x, 99)
+        x = x / max_
+        return x
+
+    pred = predict_with_halo(pv, model, [0], block_shape, halo, preprocess=preproc, postprocess=postproc).squeeze()
+
+    pred_name = "pred-v5"
+    out_folder = "./debug-pred"
+    os.makedirs(out_folder, exist_ok=True)
+
+    out_path = os.path.join(out_folder, f"{pred_name}.h5")
+    with zarr.open(out_path, "w") as f:
+        f.create_dataset("pred", data=pred)
+
+
+def main():
+    position = positions[0]
+    run_prediction(position)
+
+
+if __name__ == "__main__":
+    main()

scripts/la-vision/detect_blocks.py

@@ -0,0 +1,34 @@
+import os
+import imageio.v3 as imageio
+from pathlib import Path
+
+import pandas as pd
+import torch
+from skimage.feature import peak_local_max
+from torch_em.util.prediction import predict_with_halo
+
+ims = [
+    "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/SGN/sgn-detection/images/LaVision-M04_crop_2580-2266-0533_PV.tif",
+    "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/SGN/sgn-detection/empty_images/LaVision-M04_crop_0400-2500-0840_PV_empty.tif"
+]
+
+model_path = "checkpoints/sgn-detection.pt"
+model = torch.load(model_path, weights_only=False)
+
+block_shape = [24, 256, 256]
+halo = (8, 64, 64)
+
+out = "./detections-v1"
+os.makedirs(out, exist_ok=True)
+for im in ims:
+    data = imageio.imread(im)
+    pred = predict_with_halo(data, model, [0], block_shape, halo).squeeze()
+
+    coords = peak_local_max(pred, min_distance=4, threshold_abs=0.5)
+
+    # coords = np.concatenate([np.arange(0, len(coords))[:, None], coords], axis=1)
+    coords = pd.DataFrame(coords, columns=["axis-0", "axis-1", "axis-2"])
+
+    name = Path(im).stem
+    imageio.imwrite(os.path.join(out, f"{name}.tif"), pred)
+    coords.to_csv(os.path.join(out, f"{name}.csv"), index=False)

scripts/la-vision/export_model.py

@@ -0,0 +1,26 @@
+import argparse
+import sys
+
+import torch
+from torch_em.util import load_model
+
+sys.path.append("/home/pape/Work/my_projects/czii-protein-challenge")
+sys.path.append("/user/pape41/u12086/Work/my_projects/czii-protein-challenge")
+sys.path.append("../synapse_marker_detection")
+
+
+def export_model(input_, output):
+    model = load_model(input_, device="cpu")
+    torch.save(model, output)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input", required=True)
+    parser.add_argument("-o", "--output", required=True)
+    args = parser.parse_args()
+    export_model(args.input, args.output)
+
+
+if __name__ == "__main__":
+    main()