Update CLI

trackastra/cli.py

@@ -84,6 +84,18 @@ def cli():
             " falling back to cpu."
         ),
     )
+    p_track.add_argument(
+        "--batch-size",
+        type=int,
+        default=None,
+        help="Batch size for model inference. If not set, uses device-dependent default.",
+    )
+    p_track.add_argument(
+        "--n-workers",
+        type=int,
+        default=0,
+        help="Number of workers for feature extraction.",
+    )
     p_track.set_defaults(cmd=_track_from_disk)
 
     if len(sys.argv) == 1:
@@ -114,7 +126,12 @@ def _track_from_disk(args):
         )
 
     track_graph, masks = model.track_from_disk(
-        args.imgs, args.masks, mode=args.mode, max_distance=args.max_distance
+        args.imgs,
+        args.masks,
+        mode=args.mode,
+        max_distance=args.max_distance,
+        batch_size=args.batch_size,
+        n_workers=args.n_workers,
     )
 
     if args.output_ctc:

trackastra/model/model_api.py

@@ -132,6 +132,7 @@ def from_folder(
         dir: Path | str,
         device: str | None = None,
         checkpoint_path: str | None = None,
+        **kwargs,
     ):
         """Load a Trackastra model from a local folder.
 
@@ -152,11 +153,17 @@ def from_folder(
             Path(dir).expanduser(), map_location="cpu", checkpoint_path=checkpoint_path
         )
         train_args = yaml.load(open(dir / "train_config.yaml"), Loader=yaml.FullLoader)
-        return cls(transformer=transformer, train_args=train_args, device=device)
+        return cls(
+            transformer=transformer, train_args=train_args, device=device, **kwargs
+        )
 
     @classmethod
     def from_pretrained(
-        cls, name: str, device: str | None = None, download_dir: Path | None = None
+        cls,
+        name: str,
+        device: str | None = None,
+        download_dir: Path | None = None,
+        **kwargs,
     ):
         """Load a pretrained Trackastra model.
 
@@ -172,7 +179,7 @@ def from_pretrained(
         """
         folder = download_pretrained(name, download_dir)
         # download zip from github to location/name, then unzip
-        return cls.from_folder(folder, device=device)
+        return cls.from_folder(folder, device=device, **kwargs)
 
     def _predict(
         self,
@@ -210,15 +217,16 @@ def _predict(
             as_torch=True,
         )
 
-        logger.info("Predicting windows")
+        batch_size = batch_size or self.batch_size
+        logger.info(f"Predicting windows with batch size {batch_size}")
         predictions = predict_windows(
             windows=windows,
             features=features,
             model=self.transformer,
             edge_threshold=edge_threshold,
             spatial_dim=masks.ndim - 1,
             progbar_class=progbar_class,
-            batch_size=batch_size or self.batch_size,
+            batch_size=batch_size,
         )
 
         return predictions