Small fixes

Author: schilling40

flamingo_tools/segmentation/postprocessing.py

@@ -325,12 +325,12 @@ def graph_connected_components(coords, min_edge_distance, min_component_length):
         graph.add_node(num, pos=pos)
 
     # create edges between points whose distance is less than threshold min_edge_distance
-    for i in coords:
-        for j in coords:
-            if i < j:
-                dist = math.dist(coords[i], coords[j])
+    for num_i, pos_i in coords.items():
+        for num_j, pos_j in coords.items():
+            if num_i < num_j:
+                dist = math.dist(pos_i, pos_j)
                 if dist <= min_edge_distance:
-                    graph.add_edge(i, j, weight=dist)
+                    graph.add_edge(num_i, num_j, weight=dist)
 
     components = list(nx.connected_components(graph))
 

flamingo_tools/segmentation/unet_prediction.py

@@ -432,7 +432,8 @@ def run_unet_prediction_preprocess_slurm(
     if not os.path.isdir(os.path.join(output_folder, "mask.zarr")):
         find_mask(input_path, input_key, output_folder)
 
-    calc_mean_and_std(input_path, input_key, output_folder)
+    if not os.path.isfile(os.path.join(output_folder, "mean_std.json")):
+        calc_mean_and_std(input_path, input_key, output_folder)
 
 
 def run_unet_prediction_slurm(
@@ -504,15 +505,23 @@ def run_unet_prediction_slurm(
 def run_unet_segmentation_slurm(
     output_folder: str,
     min_size: int,
+    center_distance_threshold: float = 0.4,
     boundary_distance_threshold: float = 0.5,
+    fg_threshold: float = 0.5,
 ) -> None:
     """Create segmentation from prediction.
 
     Args:
         output_folder: The output folder for storing the segmentation related data.
         min_size: The minimal size of segmented objects in the output.
+        center_distance_threshold: The threshold applied to the distance center predictions to derive seeds.
+        boundary_distance_threshold: The threshold applied to the boundary predictions to derive seeds.
+            By default this is set to 'None', in which case the boundary distances are not used for the seeds.
+        fg_threshold: The threshold applied to the foreground prediction for deriving the watershed mask.
     """
     min_size = int(min_size)
     pmap_out = os.path.join(output_folder, "predictions.zarr")
-    distance_watershed_implementation(pmap_out, output_folder, boundary_distance_threshold=boundary_distance_threshold,
+    distance_watershed_implementation(pmap_out, output_folder, center_distance_threshold=center_distance_threshold,
+                                      boundary_distance_threshold=boundary_distance_threshold,
+                                      fg_threshold=fg_threshold,
                                       min_size=min_size)

scripts/extract_block.py

@@ -10,13 +10,14 @@
 import zarr
 
 import flamingo_tools.s3_utils as s3_utils
+from flamingo_tools.file_utils import read_image_data
 
 
 def main(
     input_path: str,
     coords: List[int],
-    output_dir: str = None,
-    input_key: str = "setup0/timepoint0/s0",
+    output_dir: Optional[str] = None,
+    input_key: Optional[str] = None,
     output_key: Optional[str] = None,
     resolution: float = 0.38,
     roi_halo: List[int] = [128, 128, 64],
@@ -88,24 +89,17 @@ def main(
     roi = tuple(slice(co - rh, co + rh) for co, rh in zip(coords, roi_halo))
 
     if s3:
-        s3_path, fs = s3_utils.get_s3_path(input_path, bucket_name=s3_bucket_name,
-                                           service_endpoint=s3_service_endpoint, credential_file=s3_credentials)
+        input_path, fs = s3_utils.get_s3_path(input_path, bucket_name=s3_bucket_name,
+                                              service_endpoint=s3_service_endpoint, credential_file=s3_credentials)
 
-        with zarr.open(s3_path, mode="r") as f:
-            raw = f[input_key][roi]
-
-    elif ".tif" in input_path:
-        raw = read_tif(input_path)[roi]
-
-    else:
-        with zarr.open(input_path, mode="r") as f:
-            raw = f[input_key][roi]
+    data_ = read_image_data(input_path, input_key)
+    data_roi = data_[roi]
 
     if tif:
-        imageio.imwrite(output_file, raw, compression="zlib")
+        imageio.imwrite(output_file, data_roi, compression="zlib")
     else:
         with zarr.open(output_file, mode="w") as f_out:
-            f_out.create_dataset(output_key, data=raw, compression="gzip")
+            f_out.create_dataset(output_key, data=data_roi, compression="gzip")
 
 
 if __name__ == "__main__":
@@ -118,7 +112,7 @@ def main(
     parser.add_argument('-c', "--coord", type=str, required=True,
                         help="3D coordinate as center of extracted block, json-encoded.")
 
-    parser.add_argument('-k', "--input_key", type=str, default="setup0/timepoint0/s0",
+    parser.add_argument('-k', "--input_key", type=str, default=None,
                         help="Input key for data in input file.")
     parser.add_argument("--output_key", type=str, default=None,
                         help="Output key for data in output file.")