Update scripts for processing portal segmentations

Author: constantinpape

scripts/cryo/cryo-et-portal/download_tomogram_lists.py

@@ -27,12 +27,13 @@ def download_tomograms_for_da():
 def download_tomograms_for_eval():
     with open("./list_for_eval.json") as f:
         run_ids = json.load(f)
-    download_tomogram_list(run_ids)
+    output_root = "/scratch-grete/projects/nim00007/cryo-et/from_portal/for_eval"
+    download_tomogram_list(run_ids, output_root)
 
 
 def main():
-    # download_tomograms_for_eval()
-    download_tomograms_for_da()
+    download_tomograms_for_eval()
+    # download_tomograms_for_da()
 
 
 if __name__ == "__main__":

scripts/cryo/cryo-et-portal/process_results.py

@@ -0,0 +1,69 @@
+import os
+from glob import glob
+from pathlib import Path
+
+import h5py
+import numpy as np
+import zarr
+
+from synapse_net.file_utils import read_mrc
+from tqdm import tqdm
+
+from ome_zarr.writer import write_image
+from ome_zarr.io import parse_url
+
+
+IN_ROOT = "/scratch-grete/projects/nim00007/cryo-et/from_portal/for_eval"
+OUT_ROOT = "/scratch-grete/projects/nim00007/cryo-et/from_portal/segmentations/DA_with_new_portalData_origDim"  # noqa
+
+IN_ROOT = "/scratch-grete/projects/nim00007/cryo-et/from_portal/for_domain_adaptation"
+OUT_ROOT = "/scratch-grete/projects/nim00007/cryo-et/from_portal/segmentations/DA_with_new_portalData_forDAdata"
+
+
+def export_to_ome_zarr(export_file, seg, voxel_size):
+    store = parse_url(export_file, mode="w").store
+    root = zarr.group(store=store)
+
+    scale = list(voxel_size.values())
+    trafo = [
+        [{"scale": scale, "type": "scale"}]
+    ]
+    write_image(seg, root, axes="zyx", coordinate_transformations=trafo, scaler=None)
+
+
+def export_segmentation(export_folder, segmentation_file):
+    fname = Path(segmentation_file).stem
+    key = "/vesicles/segment_from_vesicle_DA_portal_v3"
+    export_file = os.path.join(export_folder, f"{fname}.ome.zarr")
+
+    if os.path.exists(export_file):
+        return
+
+    input_file = os.path.join(IN_ROOT, f"{fname}.mrc")
+    raw, voxel_size = read_mrc(input_file)
+    voxel_size = {k: v * 10 for k, v in voxel_size.items()}
+
+    try:
+        with h5py.File(segmentation_file, "r") as f:
+            seg = f[key][:]
+    except OSError as e:
+        print(e)
+        return
+
+    seg = np.flip(seg, axis=1)
+    assert seg.shape == raw.shape
+
+    assert seg.max() < 128, f"{seg.max()}"
+    seg = seg.astype("int8")
+    export_to_ome_zarr(export_file, seg, voxel_size)
+
+
+def main():
+    export_folder = "./for_portal2"
+    os.makedirs(export_folder, exist_ok=True)
+    files = glob(os.path.join(OUT_ROOT, "*.h5"))
+    for file in tqdm(files):
+        export_segmentation(export_folder, file)
+
+
+main()

synapse_net/file_utils.py

@@ -170,13 +170,15 @@ def read_data_from_cryo_et_portal_run(
         output_path: The path for saving the data. The data will be streamed if the path is not given.
         use_zarr_format: Whether to use the data in zarr format instead of mrc.
         processing_type: The processing type of the tomogram to download.
-        id_field: The name of the id field.
+        id_field: The name of the id field. One of 'id' or 'run_id'.
+            The 'id' references specific tomograms, whereas 'run_id' references a collection of experimental data.
         scale_level: The scale level to read from the data. Only valid for zarr data.
 
     Returns:
         The data read from the run.
         The voxel size read from the run.
     """
+    assert id_field in ("id", "run_id")
     if output_path is not None and os.path.exists(output_path):
         return read_ome_zarr(output_path) if use_zarr_format else read_mrc(output_path)