Improved functionality for exporting synapses

Author: schilling40

scripts/export_synapse_detections.py

@@ -1,6 +1,7 @@
 import argparse
 import json
 import os
+from typing import List
 
 import numpy as np
 import pandas as pd
@@ -12,7 +13,30 @@
 from tqdm import tqdm
 
 
-def export_synapse_detections(cochlea, scale, output_folder, synapse_name, reference_ihcs, max_dist, radius, id_offset):
+def export_synapse_detections(
+    cochlea: str,
+    scales: List[int],
+    output_folder: str,
+    synapse_name: str,
+    reference_ihcs: str,
+    max_dist: float,
+    radius: float,
+    id_offset: int,
+    filter_ihc_components: List[int],
+):
+    """Export synapse detections fro lower resolutions.
+
+    Args:
+        cochlea: Cochlea name on S3 bucket.
+        scales: Scale for export of lower resolution.
+        output_folder:
+        synapse_name:
+        reference_ihcs: Name of IHC segmentation.
+        max_dist: Maximal distance of synapse to IHC segmentation.
+        radius:
+        id_offset: Offset of label id of synapse output to have different colours for visualization.
+        filter_ihc_components: Component label(s) for filtering IHC segmentation.
+    """
     s3 = create_s3_target()
 
     content = s3.open(f"{BUCKET_NAME}/{cochlea}/dataset.json", mode="r", encoding="utf-8")
@@ -36,67 +60,69 @@ def export_synapse_detections(cochlea, scale, output_folder, synapse_name, refer
     seg_table = pd.read_csv(seg_table_content, sep="\t")
 
     # Only keep synapses that match to segmented IHCs of the main component.
-    valid_ihcs = seg_table[seg_table.component_labels == 1].label_id
+    valid_ihcs = seg_table[seg_table.component_labels.isin(filter_ihc_components)].label_id
     syn_table = syn_table[syn_table.matched_ihc.isin(valid_ihcs)]
 
-    # Get the reference shape at the given scale level.
-    seg_path = os.path.join(cochlea, reference_seg_info["imageData"]["ome.zarr"]["relativePath"])
-    s3_store, _ = get_s3_path(seg_path, bucket_name=BUCKET_NAME, service_endpoint=SERVICE_ENDPOINT)
-    input_key = f"s{scale}"
-    with zarr.open(s3_store, mode="r") as f:
-        shape = f[input_key].shape
-
-    # Scale the coordinates according to the scale level.
-    resolution = 0.38
-    coordinates = syn_table[["z", "y", "x"]].values
-    coordinates /= resolution
-    coordinates /= (2 ** scale)
-    coordinates = np.round(coordinates, 0).astype("int")
-
-    ihc_ids = syn_table["matched_ihc"].values
-
-    # Create the output.
-    output = np.zeros(shape, dtype="uint16")
-    mask = ball(radius).astype(bool)
-
-    for coord, matched_ihc in tqdm(
-        zip(coordinates, ihc_ids), total=len(coordinates), desc="Writing synapses to volume"
-    ):
-        bb = tuple(slice(c - radius, c + radius + 1) for c in coord)
-        try:
-            output[bb][mask] = matched_ihc + id_offset
-        except IndexError:
-            print("Index error for", coord)
-            continue
-
-    # Write the output.
-    out_folder = os.path.join(output_folder, cochlea, f"scale{scale}")
-    os.makedirs(out_folder, exist_ok=True)
-    if id_offset != 0:
-        out_path = os.path.join(out_folder, f"{synapse_name}_offset{id_offset}.tif")
-    else:
-        out_path = os.path.join(out_folder, f"{synapse_name}.tif")
-    print("Writing synapses to", out_path)
-    tifffile.imwrite(out_path, output, bigtiff=True, compression="zlib")
+    for scale in scales:
+        # Get the reference shape at the given scale level.
+        seg_path = os.path.join(cochlea, reference_seg_info["imageData"]["ome.zarr"]["relativePath"])
+        s3_store, _ = get_s3_path(seg_path, bucket_name=BUCKET_NAME, service_endpoint=SERVICE_ENDPOINT)
+        input_key = f"s{scale}"
+        with zarr.open(s3_store, mode="r") as f:
+            shape = f[input_key].shape
+
+        # Scale the coordinates according to the scale level.
+        resolution = 0.38
+        coordinates = syn_table[["z", "y", "x"]].values
+        coordinates /= resolution
+        coordinates /= (2 ** scale)
+        coordinates = np.round(coordinates, 0).astype("int")
+
+        ihc_ids = syn_table["matched_ihc"].values
+
+        # Create the output.
+        output = np.zeros(shape, dtype="uint16")
+        mask = ball(radius).astype(bool)
+
+        for coord, matched_ihc in tqdm(
+            zip(coordinates, ihc_ids), total=len(coordinates), desc="Writing synapses to volume"
+        ):
+            bb = tuple(slice(c - radius, c + radius + 1) for c in coord)
+            try:
+                output[bb][mask] = matched_ihc + id_offset
+            except IndexError:
+                print("Index error for", coord)
+                continue
+
+        # Write the output.
+        out_folder = os.path.join(output_folder, cochlea, f"scale{scale}")
+        os.makedirs(out_folder, exist_ok=True)
+        if id_offset != 0:
+            out_path = os.path.join(out_folder, f"{synapse_name}_offset{id_offset}.tif")
+        else:
+            out_path = os.path.join(out_folder, f"{synapse_name}.tif")
+        print("Writing synapses to", out_path)
+        tifffile.imwrite(out_path, output, bigtiff=True, compression="zlib")
 
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--cochlea", "-c", required=True)
-    parser.add_argument("--scale", "-s", type=int, required=True)
+    parser.add_argument("--scale", "-s", nargs="+", type=int, required=True)
     parser.add_argument("--output_folder", "-o", required=True)
     parser.add_argument("--synapse_name", default="synapse_v3_ihc_v4b")
     parser.add_argument("--reference_ihcs", default="IHC_v4b")
     parser.add_argument("--max_dist", type=float, default=3.0)
     parser.add_argument("--radius", type=int, default=3)
     parser.add_argument("--id_offset", type=int, default=0)
+    parser.add_argument("--filter_ihc_components", nargs="+", type=int, default=[1])
     args = parser.parse_args()
 
     export_synapse_detections(
         args.cochlea, args.scale, args.output_folder,
         args.synapse_name, args.reference_ihcs,
         args.max_dist, args.radius,
-        args.id_offset,
+        args.id_offset, args.filter_ihc_components,
     )