predict solutions now accept input folders

src/io.github.betaseg/csbdeep_unet_predict/solution.py

@@ -55,18 +55,25 @@ def apply(model, x0):
 
         return y
 
-    # load file
-    x0 = imread(args.fname_input)
-
     model = UNet(None, "unet", basedir=args.root)
 
-    y = apply(model, x0)
-
-    # save output
     out = Path(args.outdir)
-
     out.mkdir(exist_ok=True, parents=True)
-    imwrite(out / f"{Path(args.fname_input).stem}.unet.tif", y.astype(np.uint16))
+
+    # load file(s)
+    input_path = Path(args.fname_input)
+    if input_path.is_dir():
+        for f in input_path.iterdir():
+            if f.suffix == ".tif":
+                x0 = imread(f)
+                y = apply(model, x0)
+
+                imwrite(out / f"{f.stem}.unet.tif", y.astype(np.uint16))
+    else:
+        x0 = imread(input_path)
+        y = apply(model, x0)
+
+        imwrite(out / f"{input_path.stem}.unet.tif", y.astype(np.uint16))
 
 
 setup(
@@ -97,7 +104,7 @@ def apply(model, x0):
         },
         {
             "name": "fname_input",
-            "description": "file path to your input image. Should be a tif file.",
+            "description": "file path to your input image or folder. Should be tif file(s).",
             "type": "string",
             "required": True
         },

src/io.github.betaseg/stardist_predict/solution.py

@@ -1,5 +1,6 @@
 from album.runner.api import get_args, setup, get_environment_path
 import sys
+
 env_file = """name:  stardist_train
 channels:
   - conda-forge
@@ -34,6 +35,7 @@
     - gputools
 """
 
+
 def run():
     from csbdeep.utils.tf import limit_gpu_memory
     limit_gpu_memory(fraction=0.8, total_memory=12000)
@@ -42,40 +44,46 @@ def run():
     import numpy as np
     from stardist.models import StarDist3D
 
+    def predict_tif(f, model, outdir, x0):
+        # compute n_tiles
+        n_tiles = tuple(int(np.ceil(s / 160)) for s in x0.shape)
+        print(f"using {n_tiles} tiles")
+
+        # normalize
+        print("normalizing...")
+        x = x0.astype(np.float32) / 255
+        # predict
+        y, polys = model.predict_instances(x, n_tiles=n_tiles)
+        rays = polys["rays"]
+        polys["rays_vertices"] = rays.vertices
+        polys["rays_faces"] = rays.faces
+        # save output
+        out = Path(outdir)
+        out.mkdir(exist_ok=True, parents=True)
+        imsave(out / f"{Path(f).stem}.stardist.tif", y)
+        np.savez(out / f"{Path(f).stem}.stardist.npz", **polys)
+
     # parse arguments
     args = get_args()
 
     # options
-    fname = args.fname_input
     model_name = args.model_name
     outdir = args.output_dir
     basedir = args.model_basedir
 
-    # load file
-    x0 = imread(fname)
-
-    # compute n_tiles
-    n_tiles = tuple(int(np.ceil(s / 160)) for s in x0.shape)
-    print(f"using {n_tiles} tiles")
-
     # load model and apply it to the stack
     model = StarDist3D(None, name=model_name, basedir=basedir)
 
-    # normalize
-    print("normalizing...")
-    x = x0.astype(np.float32) / 255
-
-    # predict
-    y, polys = model.predict_instances(x, n_tiles=n_tiles)
-    rays = polys["rays"]
-    polys["rays_vertices"] = rays.vertices
-    polys["rays_faces"] = rays.faces
-
-    # save output
-    out = Path(outdir)
-    out.mkdir(exist_ok=True, parents=True)
-    imsave(out / f"{Path(fname).stem}.stardist.tif", y)
-    np.savez(out / f"{Path(fname).stem}.stardist.npz", **polys)
+    # load files
+    input_path = Path(args.fname_input)
+    if input_path.is_dir():
+        for f in input_path.iterdir():
+            if f.suffix == ".tif":
+                x0 = imread(f)
+                predict_tif(f, model, outdir, x0)
+    else:
+        x0 = imread(input_path)
+        predict_tif(input_path, model, outdir, x0)
 
 
 setup(
@@ -103,7 +111,7 @@ def run():
     args=[
         {
             "name": "fname_input",
-            "description": "Path to the input file.",
+            "description": "Path to the input tif file or folder of tif images.",
             "required": True
         },
         {