From 175aeea206176ea924691c9c9f431107cde2a8f3 Mon Sep 17 00:00:00 2001
From: brianreicher <reicher.b@northeastern.edu>
Date: Wed, 29 Nov 2023 16:00:05 -0500
Subject: [PATCH] tiff to zarr util

---
 src/autoseg/train_job.py |  4 +++-
 src/autoseg/utils.py     | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/autoseg/train_job.py b/src/autoseg/train_job.py
index 1dd11b3..5eab6fc 100644
--- a/src/autoseg/train_job.py
+++ b/src/autoseg/train_job.py
@@ -5,11 +5,13 @@ def train_model(
     model_type: str = "MTLSD",
     iterations: int = 100000,
     warmup: int = 100000,
-    raw_file: str = "path/to/.zarr/or/.n5",
+    raw_file: str = "path/to/.zarr/or/.n5/or/.tiff",
     out_file: str = "./raw_predictions.zarr",
     voxel_size: int = 33,
     save_every=2500,
 ) -> None:
+    
+    # TODO: call ztools to rewrite .tiff file to zarr format
     model_type = model_type.lower()
     if model_type == "mtlsd":
         mtlsd_train(
diff --git a/src/autoseg/utils.py b/src/autoseg/utils.py
index 367fc9c..fbdf8c6 100644
--- a/src/autoseg/utils.py
+++ b/src/autoseg/utils.py
@@ -1,3 +1,9 @@
+import numpy as np
+from funlib.persistence import prepare_ds
+from funlib.geometry import Coordinate, Roi
+import tifffile
+
+
 neighborhood: list[list[int]] = [
     [1, 0, 0],
     [0, 1, 0],
@@ -15,3 +21,33 @@
     [0, 10, 0],
     [0, 0, 10],
 ]
+
+def tiff_to_zarr(tiff_file:str="path/to/.tiff",
+                out_file:str="path/to/.zarr",
+                out_ds:str="path/to/volume",
+                voxel_size: int = 33,
+                offset: int = 0,
+                dtype=np.uint8,
+                transpose:bool=False) -> None:
+    tiff_stack: np.ndarray = tifffile.imread(tiff_file)
+    if transpose:
+        tiff_stack = np.transpose(tiff_stack, (2, 1, 0))
+
+    voxel_size: Coordinate = Coordinate((voxel_size)*3)
+    roi: Roi = Roi(offset=(offset)*3, shape=tiff_stack.shape * np.array(voxel_size))
+
+    print("Roi: ", roi)
+    voxel_size: Coordinate = Coordinate(100, 100, 100)
+
+    ds = prepare_ds(
+        filename=out_file,
+        ds_name=out_ds,
+        total_roi=roi,
+        voxel_size=voxel_size,
+        dtype=dtype,
+        delete=True,
+    )
+
+    ds[roi] = tiff_stack
+
+    print("TIFF Image stack saved as Zarr dataset.")