From 329db3fc43e379a26f4453d5adca0dbc4769b038 Mon Sep 17 00:00:00 2001
From: cmoyates <cmoyates@gmail.com>
Date: Sun, 1 Mar 2026 06:11:57 -0330
Subject: [PATCH] fix: resize inputs to model img_size, use infer_and_save in
 CLI

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 scripts/infer.py                          | 52 +++++------------------
 src/corridorkey_mlx/inference/pipeline.py |  7 ++-
 src/corridorkey_mlx/io/image.py           | 22 ++++++++--
 3 files changed, 32 insertions(+), 49 deletions(-)

diff --git a/scripts/infer.py b/scripts/infer.py
index 30a68bf..ec82519 100644
--- a/scripts/infer.py
+++ b/scripts/infer.py
@@ -12,17 +12,11 @@ import argparse
 import time
 from pathlib import Path
 
-import mlx.core as mx
-
-from corridorkey_mlx.inference.pipeline import DEFAULT_CHECKPOINT, DEFAULT_IMG_SIZE, load_model
-from corridorkey_mlx.io.image import (
-    load_alpha_hint,
-    load_image,
-    postprocess_alpha,
-    postprocess_foreground,
-    preprocess,
-    save_alpha,
-    save_foreground,
+from corridorkey_mlx.inference.pipeline import (
+    DEFAULT_CHECKPOINT,
+    DEFAULT_IMG_SIZE,
+    infer_and_save,
+    load_model,
 )
 
 
@@ -56,38 +50,14 @@ def main() -> None:
     model = load_model(args.checkpoint, args.img_size)
     print(f"  Model loaded in {time.perf_counter() - t0:.2f}s")
 
-    print(f"Preprocessing {args.image}...")
-    rgb = load_image(args.image)
-    alpha_hint = load_alpha_hint(args.hint)
-    x = preprocess(rgb, alpha_hint)
-    # materialize input — mx.eval is MLX lazy graph eval, not Python eval
-    mx.eval(x)  # noqa: S307
-
-    print("Running inference...")
+    print(f"Running inference on {args.image}...")
     t0 = time.perf_counter()
-    outputs = model(x)
-    # materialize outputs — mx.eval is MLX lazy graph eval, not Python eval
-    mx.eval(outputs)  # noqa: S307
-    print(f"  Inference in {time.perf_counter() - t0:.2f}s")
+    results = infer_and_save(model, args.image, args.hint, args.output_dir)
+    print(f"  Inference + save in {time.perf_counter() - t0:.2f}s")
 
-    # Print summary
-    print("\nOutput tensors:")
-    for key, arr in outputs.items():
-        print(
-            f"  {key:<20s} shape={arr.shape} "
-            f"min={float(mx.min(arr)):.4f} max={float(mx.max(arr)):.4f}"
-        )
-
-    # Save results
-    args.output_dir.mkdir(parents=True, exist_ok=True)
-    alpha_arr = postprocess_alpha(outputs["alpha_final"])
-    fg_arr = postprocess_foreground(outputs["fg_final"])
-
-    alpha_path = args.output_dir / "alpha.png"
-    fg_path = args.output_dir / "foreground.png"
-    save_alpha(alpha_arr, alpha_path)
-    save_foreground(fg_arr, fg_path)
-    print(f"\nSaved: {alpha_path}, {fg_path}")
+    print(f"\nSaved: {args.output_dir / 'alpha.png'}, {args.output_dir / 'foreground.png'}")
+    print(f"  Alpha shape: {results['alpha'].shape}")
+    print(f"  Foreground shape: {results['foreground'].shape}")
 
 
 if __name__ == "__main__":
diff --git a/src/corridorkey_mlx/inference/pipeline.py b/src/corridorkey_mlx/inference/pipeline.py
index e57b35e..a8947c2 100644
--- a/src/corridorkey_mlx/inference/pipeline.py
+++ b/src/corridorkey_mlx/inference/pipeline.py
@@ -53,11 +53,10 @@ def infer(
     Returns:
         Model output dict with all intermediate and final tensors.
     """
-    rgb = load_image(image_path)
-    alpha_hint = load_alpha_hint(alpha_hint_path)
+    img_size = model.backbone.img_size
+    rgb = load_image(image_path, img_size=img_size)
+    alpha_hint = load_alpha_hint(alpha_hint_path, img_size=img_size)
     x = preprocess(rgb, alpha_hint)
-    # materialize input
-    mx.eval(x)  # noqa: S307
     outputs = model(x)
     # materialize all outputs
     mx.eval(outputs)  # noqa: S307
diff --git a/src/corridorkey_mlx/io/image.py b/src/corridorkey_mlx/io/image.py
index 74d3f05..a1e8517 100644
--- a/src/corridorkey_mlx/io/image.py
+++ b/src/corridorkey_mlx/io/image.py
@@ -19,15 +19,29 @@ IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
 IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
 
 
-def load_image(path: str | Path) -> np.ndarray:
-    """Load image as RGB float32 array in [0, 1] range, shape (H, W, 3)."""
+def load_image(path: str | Path, img_size: int | None = None) -> np.ndarray:
+    """Load image as RGB float32 array in [0, 1] range, shape (H, W, 3).
+
+    Args:
+        path: Path to image file.
+        img_size: If provided, resize to (img_size, img_size) using bicubic interpolation.
+    """
     img = Image.open(path).convert("RGB")
+    if img_size is not None:
+        img = img.resize((img_size, img_size), Image.BICUBIC)
     return np.asarray(img, dtype=np.float32) / 255.0
 
 
-def load_alpha_hint(path: str | Path) -> np.ndarray:
-    """Load alpha hint as grayscale float32 array in [0, 1], shape (H, W, 1)."""
+def load_alpha_hint(path: str | Path, img_size: int | None = None) -> np.ndarray:
+    """Load alpha hint as grayscale float32 array in [0, 1], shape (H, W, 1).
+
+    Args:
+        path: Path to alpha hint image file.
+        img_size: If provided, resize to (img_size, img_size) using bicubic interpolation.
+    """
     img = Image.open(path).convert("L")
+    if img_size is not None:
+        img = img.resize((img_size, img_size), Image.BICUBIC)
     return np.asarray(img, dtype=np.float32)[:, :, np.newaxis] / 255.0