Merge pull request #209 from computational-cell-analytics/fix-precompute

Fix CLI for embedding precomputation

micro_sam/precompute_state.py

@@ -101,7 +101,7 @@ def precompute_state(
     model_type: str = util._DEFAULT_MODEL,
     checkpoint_path: Optional[Union[os.PathLike, str]] = None,
     key: Optional[str] = None,
-    ndim: Union[int] = None,
+    ndim: Optional[int] = None,
     tile_shape: Optional[Tuple[int, int]] = None,
     halo: Optional[Tuple[int, int]] = None,
     precompute_amg_state: bool = False,
@@ -158,8 +158,8 @@ def main():
     parser.add_argument(
         "--halo", nargs="+", type=int, help="The halo for using tiled prediction", default=None
     )
-    parser.add_argument("-n", "--ndim")
-    parser.add_argument("-p", "--precompute_amg_state")
+    parser.add_argument("-n", "--ndim", type=int)
+    parser.add_argument("-p", "--precompute_amg_state", action="store_true")
 
     args = parser.parse_args()
     precompute_state(

setup.cfg

@@ -48,7 +48,7 @@ console_scripts =
     micro_sam.annotator_3d = micro_sam.sam_annotator.annotator_3d:main
     micro_sam.annotator_tracking = micro_sam.sam_annotator.annotator_tracking:main
     micro_sam.image_series_annotator = micro_sam.sam_annotator.image_series_annotator:main
-    micro_sam.precompute_embeddings = micro_sam.util:main
+    micro_sam.precompute_embeddings = micro_sam.precompute_state:main
 
 # make sure it gets included in your package
 [options.package_data]