Add template for processing

schilling40 · schilling40 · commit 35014089fcab · 2025-09-19T15:18:26.000+02:00
diff --git a/flamingo_tools/segmentation/sgn_detection.py b/flamingo_tools/segmentation/sgn_detection.py
@@ -57,8 +57,7 @@ def extend_block(block_id):
         outer_block = block.outerBlock
         inner_block = block.innerBlock
 
-        # TODO get the indices and coordinates of the markers in the INNER block
-        # markers_in_block_ids = [int(i) for i in np.unique(inner_block)[1:]]
+        # get the indices and coordinates of the markers in the INNER block
         mask = (
             (inner_block.begin[0] <= markers[:, 0]) & (markers[:, 0] <= inner_block.end[0]) &
             (inner_block.begin[1] <= markers[:, 1]) & (markers[:, 1] <= inner_block.end[1]) &
@@ -67,33 +66,35 @@ def extend_block(block_id):
         markers_in_block_ids = np.where(mask)[0]
         markers_in_block_coords = markers[markers_in_block_ids]
 
-        # TODO offset the marker coordinates with respect to the OUTER block
-        markers_in_block_coords = [coord - outer_block.begin for coord in markers_in_block_coords]
-        markers_in_block_coords = [[round(c) for c in coord] for coord in markers_in_block_coords]
-        markers_in_block_coords = np.array(markers_in_block_coords, dtype=int)
-        z, y, x = markers_in_block_coords.T
+        # proceed if detections fall within inner block
+        if len(markers_in_block_coords) > 0:
+            markers_in_block_coords = [coord - outer_block.begin for coord in markers_in_block_coords]
+            markers_in_block_coords = [[round(c) for c in coord] for coord in markers_in_block_coords]
 
-        # Shift index by one so that zero is reserved for background id
-        markers_in_block_ids += 1
+            markers_in_block_coords = np.array(markers_in_block_coords, dtype=int)
+            z, y, x = markers_in_block_coords.T
 
-        # Create the seed volume.
-        outer_block_shape = tuple(end - begin for begin, end in zip(outer_block.begin, outer_block.end))
-        seeds = np.zeros(outer_block_shape, dtype="uint32")
-        seeds[z, y, x] = markers_in_block_ids
+            # Shift index by one so that zero is reserved for background id
+            markers_in_block_ids += 1
 
-        # Compute the distance map.
-        distance = distance_transform_edt(seeds == 0, sampling=sampling)
+            # Create the seed volume.
+            outer_block_shape = tuple(end - begin for begin, end in zip(outer_block.begin, outer_block.end))
+            seeds = np.zeros(outer_block_shape, dtype="uint32")
+            seeds[z, y, x] = markers_in_block_ids
 
-        # And extend the seeds
-        mask = distance < extension_distance
-        segmentation = watershed(distance.max() - distance, markers=seeds, mask=mask)
+            # Compute the distance map.
+            distance = distance_transform_edt(seeds == 0, sampling=sampling)
 
-        # Write the segmentation. Note: we need to lock here because we write outside of our inner block
-        bb = tuple(slice(begin, end) for begin, end in zip(outer_block.begin, outer_block.end))
-        with lock:
-            this_output = output[bb]
-            this_output[mask] = segmentation[mask]
-            output[bb] = this_output
+            # And extend the seeds
+            mask = distance < extension_distance
+            segmentation = watershed(distance.max() - distance, markers=seeds, mask=mask)
+
+            # Write the segmentation. Note: we need to lock here because we write outside of our inner block
+            bb = tuple(slice(begin, end) for begin, end in zip(outer_block.begin, outer_block.end))
+            with lock:
+                this_output = output[bb]
+                this_output[mask] = segmentation[mask]
+                output[bb] = this_output
 
     n_blocks = blocking.numberOfBlocks
     with futures.ThreadPoolExecutor(n_threads) as tp:
@@ -146,10 +147,18 @@ def sgn_detection(
     detection_path = os.path.join(output_folder, "SGN_detection.tsv")
     if not os.path.exists(detection_path):
         input_ = zarr.open(output_path, "r")[prediction_key]
-        detections = find_local_maxima(
+        detections_maxima = find_local_maxima(
             input_, block_shape=block_shape, min_distance=4, threshold_abs=0.5, verbose=True, n_threads=16,
         )
 
+        # Save the result in mobie compatible format.
+        detections = np.concatenate(
+            [np.arange(1, len(detections_maxima) + 1)[:, None], detections_maxima[:, ::-1]], axis=1
+        )
+        detections = pd.DataFrame(detections, columns=["spot_id", "x", "y", "z"])
+        detections.to_csv(detection_path, index=False, sep="\t")
+
+        # extend detection
         shape = input_.shape
         chunks = (128, 128, 128)
         segmentation_path = os.path.join(output_folder, "segmentation.zarr")
@@ -161,17 +170,11 @@ def sgn_detection(
         )
 
         distance_based_marker_extension(
-            markers=detections,
+            markers=detections_maxima,
             output=output_dataset,
             extension_distance=extension_distance,
             sampling=sampling,
             block_shape=(128, 128, 128),
             n_threads=n_threads,
         )
 
-        # Save the result in mobie compatible format.
-        detections = np.concatenate(
-            [np.arange(1, len(detections) + 1)[:, None], detections[:, ::-1]], axis=1
-        )
-        detections = pd.DataFrame(detections, columns=["spot_id", "x", "y", "z"])
-        detections.to_csv(detection_path, index=False, sep="\t")
diff --git a/reproducibility/templates_processing/detect_sgn_template.sbatch b/reproducibility/templates_processing/detect_sgn_template.sbatch
@@ -0,0 +1,60 @@
+#!/bin/bash
+#SBATCH --job-name=synapse-detect
+#SBATCH -t 03:00:00             # estimated time, adapt to your needs
+#SBATCH --mail-type=FAIL        # send mail when job begins and ends
+
+#SBATCH -p grete:shared         # the partition
+#SBATCH -G A100:1               # For requesting 1 A100 GPU.
+#SBATCH -A nim00007
+#SBATCH -c 4
+#SBATCH --mem 32G
+
+source ~/.bashrc
+# micromamba activate micro-sam_gpu
+micromamba activate sam
+
+# Print out some info.
+echo "Submitting job with sbatch from directory: ${SLURM_SUBMIT_DIR}"
+echo "Home directory: ${HOME}"
+echo "Working directory: $PWD"
+echo "Current node: ${SLURM_NODELIST}"
+
+# Run the script
+#python myprogram.py $SLURM_ARRAY_TASK_ID
+
+# SCRIPT_REPO=/user/schilling40/u15000/flamingo-tools
+SCRIPT_REPO=/user/pape41/u12086/Work/my_projects/flamingo-tools
+cd "$SCRIPT_REPO"/flamingo_tools/segmentation/ || exit
+
+export SCRIPT_DIR=$SCRIPT_REPO/scripts
+
+# name of cochlea, as it appears in MoBIE and the NHR
+COCHLEA=$1
+# model of SGN detection, e.g. v5b
+MODEL_VERSION=$2
+
+# data on NHR
+MOBIE_DIR=/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/mobie_project/cochlea-lightsheet/
+export INPUT_PATH="$MOBIE_DIR"/"$COCHLEA"/images/ome-zarr/PV.ome.zarr
+
+# data on MoBIE
+# export INPUT_PATH="$COCHLEA"/images/ome-zarr/PV.ome.zarr
+# use --s3 flag for script
+
+export OUTPUT_FOLDER=/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/predictions/"$COCHLEA"/SGN_detect-"$MODEL_VERSION"
+
+if ! [[ -f $OUTPUT_FOLDER ]] ; then
+	mkdir -p "$OUTPUT_FOLDER"
+fi
+
+export MODEL=/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/trained_models/SGN/sgn-detection-"$MODEL_VERSION".pt
+INPUT_KEY="s0"
+
+echo "OUTPUT_FOLDER $OUTPUT_FOLDER"
+echo "MODEL $MODEL"
+
+python ~/flamingo-tools/scripts/sgn_detection/sgn_detection.py \
+	--input "$INPUT_PATH" \
+	--input_key $INPUT_KEY \
+	--output_folder "$OUTPUT_FOLDER" \
+	--model "$MODEL"
