Added option to save extracted block as TIF

scripts/extract_block.py

@@ -5,6 +5,7 @@
 import os
 from typing import Optional, List
 
+import imageio.v3 as imageio
 import numpy as np
 import zarr
 
@@ -16,8 +17,10 @@ def main(
     coords: List[int],
     output_dir: str = None,
     input_key: str = "setup0/timepoint0/s0",
-    resolution: float = 0.38,
+    output_key: Optional[str] = None,
+    resolution: Optional[float] = 0.38,
     roi_halo: List[int] = [128, 128, 64],
+    tif: Optional[bool] = False,
     s3: Optional[bool] = False,
     s3_credentials: Optional[str] = None,
     s3_bucket_name: Optional[str] = None,
@@ -30,14 +33,17 @@ def main(
         input_path: Input folder in n5 / ome-zarr format.
         coords: Center coordinates of extracted 3D volume.
         output_dir: Output directory for saving output as <basename>_crop.n5. Default: input directory.
+        input_key: Input key for data in input file.
+        output_key: Output key for data in n5 output or used as suffix for tif output.
         roi_halo: ROI halo of extracted 3D volume.
+        tif: Flag for tif output
         s3: Flag for considering input_path for S3 bucket.
         s3_bucket_name: S3 bucket name.
         s3_service_endpoint: S3 service endpoint.
         s3_credentials: File path to credentials for S3 bucket.
     """
-
-    coord_string = "-".join([str(c) for c in coords])
+    coords = [int(round(c)) for c in coords]
+    coord_string = "-".join([str(c).zfill(4) for c in coords])
 
     # Dimensions are inversed to view in MoBIE (x y z) -> (z y x)
     coords.reverse()
@@ -56,7 +62,14 @@ def main(
     if output_dir == "":
         output_dir = input_dir
 
-    output_file = os.path.join(output_dir, basename + "_crop_" + coord_string + ".n5")
+    if tif:
+        if output_key is None:
+            output_file = os.path.join(output_dir, basename + "_crop_" + coord_string + ".tif")
+        else:
+            output_file = os.path.join(output_dir, basename + "_crop_" + coord_string + "_" + output_key + ".tif")
+    else:
+        output_key = "raw" if output_key is None else output_key
+        output_file = os.path.join(output_dir, basename + "_crop_" + coord_string + ".n5")
 
     coords = np.array(coords)
     coords = coords / resolution
@@ -75,26 +88,31 @@ def main(
         with zarr.open(input_path, mode="r") as f:
             raw = f[input_key][roi]
 
-    with zarr.open(output_file, mode="w") as f_out:
-        f_out.create_dataset("raw", data=raw, compression="gzip")
+    if tif:
+        imageio.imwrite(output_file, raw)
+    else:
+        with zarr.open(output_file, mode="w") as f_out:
+            f_out.create_dataset(output_key, data=raw, compression="gzip")
 
 
 if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(
         description="Script to extract region of interest (ROI) block around center coordinate.")
 
-    parser.add_argument('-i', '--input', type=str, help="Input file in n5 / ome-zarr format.")
+    parser.add_argument('-i', '--input', type=str, required=True, help="Input file in n5 / ome-zarr format.")
     parser.add_argument('-o', "--output", type=str, default="", help="Output directory.")
     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",
                         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.")
     parser.add_argument('-r', "--resolution", type=float, default=0.38, help="Resolution of input in micrometer.")
-
     parser.add_argument("--roi_halo", type=str, default="[128,128,64]",
                         help="ROI halo around center coordinate, json-encoded.")
+    parser.add_argument("--tif", action="store_true", help="Store output as tif file.")
 
     parser.add_argument("--s3", action="store_true", help="Use S3 bucket.")
     parser.add_argument("--s3_credentials", type=str, default=None,
@@ -111,6 +129,6 @@ def main(
     roi_halo = json.loads(args.roi_halo)
 
     main(
-        args.input, coords, args.output, args.input_key, args.resolution, roi_halo,
-        args.s3, args.s3_credentials, args.s3_bucket_name, args.s3_service_endpoint,
+        args.input, coords, args.output, args.input_key, args.output_key, args.resolution, roi_halo,
+        args.tif, args.s3, args.s3_credentials, args.s3_bucket_name, args.s3_service_endpoint,
     )