From 7245e30236d3c67b31fd74e877f9dfaa98a71313 Mon Sep 17 00:00:00 2001
From: brianreicher <reicher.b@northeastern.edu>
Date: Thu, 30 Nov 2023 14:30:18 -0500
Subject: [PATCH] tif extensions

---
 src/autoseg/train_job.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/autoseg/train_job.py b/src/autoseg/train_job.py
index a42bdf6..33baafe 100644
--- a/src/autoseg/train_job.py
+++ b/src/autoseg/train_job.py
@@ -8,20 +8,21 @@ def train_model(
     warmup: int = 100000,
     raw_file: str = "path/to/.zarr/or/.n5/or/.tiff",
     rewrite_file: str = "./rewritten.zarr",
-    rewrite_ds: str = "rewritten_volume",
+    rewrite_ds: str = "training_raw",
     out_file: str = "./raw_predictions.zarr",
     voxel_size: int = 33,
     save_every=2500,
 ) -> None:
     
-    # TODO: add util funcs for generating masks 
-    if raw_file.endswith(".tiff"):
+    # TODO: add util funcs for generating masks, pulling paintings
+    if raw_file.endswith(".tiff") or raw_file.endswith(".tif"):
         tiff_to_zarr(tiff_file=raw_file,
                      out_file=rewrite_file,
                      out_ds=rewrite_ds)
+        raw_file: str = rewrite_file
     
 
-    model_type = model_type.lower()
+    model_type: str = model_type.lower()
     if model_type == "mtlsd":
         mtlsd_train(
             iterations=iterations,