From 6470984279f838419cf9c0bc3e367b115a8676a9 Mon Sep 17 00:00:00 2001
From: Anwai Archit <anwai.archit@gmail.com>
Date: Tue, 24 Sep 2024 17:58:47 +0200
Subject: [PATCH] Support saving data to a target path

---
 micro_sam/instance_segmentation.py | 18 +++++++++---------
 micro_sam/sam_annotator/util.py    |  2 +-
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/micro_sam/instance_segmentation.py b/micro_sam/instance_segmentation.py
index f79f7b46..9fc5d33a 100644
--- a/micro_sam/instance_segmentation.py
+++ b/micro_sam/instance_segmentation.py
@@ -6,12 +6,14 @@
 
 import os
 from abc import ABC
+from pathlib import Path
 from copy import deepcopy
 from collections import OrderedDict
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import vigra
 import numpy as np
+import imageio.v3 as imageio
 from skimage.measure import regionprops
 
 import torch
@@ -1145,6 +1147,7 @@ def get_amg(
 
 def automatic_instance_segmentation(
     input_path: Union[os.PathLike, str],
+    output_path: Union[os.PathLike, str],
     embedding_path: Optional[Union[os.PathLike, str]] = None,
     model_type: str = util._DEFAULT_MODEL,
     checkpoint_path: Optional[Union[os.PathLike, str]] = None,
@@ -1191,13 +1194,8 @@ def automatic_instance_segmentation(
     else:
         instances = mask_data_to_segmentation(masks, with_background=True, min_object_size=0)
 
-    import napari
-    v = napari.Viewer()
-    v.add_image(image_data)
-    v.add_labels(instances)
-    napari.run()
-
-    breakpoint()
+    output_path = Path(output_path).with_suffix(".tif")
+    imageio.imwrite(output_path, instances, compression="zlib")
 
 
 def main():
@@ -1213,10 +1211,13 @@ def main():
         help="The filepath to the image data. Supports all data types that can be read by imageio (e.g. tif, png, ...) "
         "or elf.io.open_file (e.g. hdf5, zarr, mrc). For the latter you also need to pass the 'key' parameter."
     )
+    parser.add_argument(
+        "-o", "--output_path", required=True,
+        help="The filepath to store the instance segmentation. The current support stores segmentation in a 'tif' file."
+    )
     parser.add_argument(
         "-e", "--embedding_path", default=None, type=str, help="The path where the embeddings will be saved."
     )
-
     parser.add_argument(
         "--pattern", help="Pattern / wildcard for selecting files in a folder. To select all files use '*'."
     )
@@ -1225,7 +1226,6 @@ def main():
         help="The key for opening data with elf.io.open_file. This is the internal path for a hdf5 or zarr container, "
         "for an image stack it is a wild-card, e.g. '*.png' and for mrc it is 'data'."
     )
-
     parser.add_argument(
         "-m", "--model_type", default=util._DEFAULT_MODEL,
         help=f"The segment anything model that will be used, one of {available_models}."
diff --git a/micro_sam/sam_annotator/util.py b/micro_sam/sam_annotator/util.py
index aae5a75e..29c5b5d4 100644
--- a/micro_sam/sam_annotator/util.py
+++ b/micro_sam/sam_annotator/util.py
@@ -679,7 +679,7 @@ def _sync_embedding_widget(widget, model_type, save_path, checkpoint_path, devic
     if index > 0:
         widget.model_dropdown.setCurrentIndex(index)
 
-    if save_path is not None:
+    if save_path is not None and isinstance(save_path, str):
         widget.embeddings_save_path_param.setText(save_path)
 
     if checkpoint_path is not None: