Merge pull request #54 from computational-cell-analytics/export-updates

Update the data export scripts to preserve the label IDs of the IHC segmentations and their corresponding synapses.
The segmentation must be exported as "uint16," which limits the number of SGN instances to the data type's maximum of 65,535. A warning will be displayed if the label is higher than this value.
An offset option for the synapse IDs in the data export has been added to allow for different coloring in Imaris.

Author: schilling40

scripts/export_lower_resolution.py

@@ -1,13 +1,13 @@
 import argparse
 import os
+import warnings
 
 import numpy as np
 import pandas as pd
 import tifffile
 import zarr
 
 from flamingo_tools.s3_utils import get_s3_path, BUCKET_NAME, SERVICE_ENDPOINT
-# from skimage.segmentation import relabel_sequential
 
 
 def filter_component(fs, segmentation, cochlea, seg_name, components):
@@ -19,10 +19,12 @@ def filter_component(fs, segmentation, cochlea, seg_name, components):
     # Then we get the ids for the components and us them to filter the segmentation.
     component_mask = np.isin(table.component_labels.values, components)
     keep_label_ids = table.label_id.values[component_mask].astype("int64")
+    if max(keep_label_ids) > np.iinfo("uint16").max:
+        warnings.warn(f"Label ID exceeds maximum of data type 'uint16': {np.iinfo('uint16').max}.")
+
     filter_mask = ~np.isin(segmentation, keep_label_ids)
     segmentation[filter_mask] = 0
-
-    # segmentation, _, _ = relabel_sequential(segmentation)
+    segmentation = segmentation.astype("uint16")
     return segmentation
 
 
@@ -41,7 +43,7 @@ def export_lower_resolution(args):
         s3_store, fs = get_s3_path(internal_path, bucket_name=BUCKET_NAME, service_endpoint=SERVICE_ENDPOINT)
         with zarr.open(s3_store, mode="r") as f:
             data = f[input_key][:]
-        print(data.shape)
+
         if args.filter_by_components is not None:
             data = filter_component(fs, data, args.cochlea, channel, args.filter_by_components)
         if args.binarize:

scripts/export_synapse_detections.py

@@ -12,7 +12,7 @@
 from tqdm import tqdm
 
 
-def export_synapse_detections(cochlea, scale, output_folder, synapse_name, reference_ihcs, max_dist, radius):
+def export_synapse_detections(cochlea, scale, output_folder, synapse_name, reference_ihcs, max_dist, radius, id_offset):
     s3 = create_s3_target()
 
     content = s3.open(f"{BUCKET_NAME}/{cochlea}/dataset.json", mode="r", encoding="utf-8")
@@ -53,22 +53,29 @@ def export_synapse_detections(cochlea, scale, output_folder, synapse_name, refer
     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 in tqdm(coordinates, desc="Writing synapses to volume"):
+    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] = 1
+            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)
-    out_path = os.path.join(out_folder, f"{synapse_name}.tif")
+    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")
 
@@ -78,16 +85,18 @@ def main():
     parser.add_argument("--cochlea", "-c", required=True)
     parser.add_argument("--scale", "-s", type=int, required=True)
     parser.add_argument("--output_folder", "-o", required=True)
-    parser.add_argument("--synapse_name", default="synapse_v3_ihc_v4")
-    parser.add_argument("--reference_ihcs", default="IHC_v4")
+    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)
     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.max_dist, args.radius,
+        args.id_offset,
     )