diff --git a/src/corridorkey_mlx/engine.py b/src/corridorkey_mlx/engine.py index d2127f4..5170bd6 100644 --- a/src/corridorkey_mlx/engine.py +++ b/src/corridorkey_mlx/engine.py @@ -26,6 +26,7 @@ from corridorkey_mlx.io.image import ( postprocess_foreground, preprocess, ) +from corridorkey_mlx.io.preprocess_mlx import preprocess_mlx if TYPE_CHECKING: from corridorkey_mlx.model.corridorkey import GreenFormer @@ -178,7 +179,7 @@ class CorridorKeyMLXEngine: ) mask_f32 = np.asarray(mask_pil, dtype=np.float32)[:, :, np.newaxis] / 255.0 - x = preprocess(rgb_f32, mask_f32) + x = preprocess_mlx(rgb_f32, mask_f32) outputs = self._model(x) # NOTE: mx.eval is MLX array materialization, not Python eval() mx.eval(outputs) # noqa: S307 diff --git a/src/corridorkey_mlx/io/preprocess_mlx.py b/src/corridorkey_mlx/io/preprocess_mlx.py new file mode 100644 index 0000000..f47b2f3 --- /dev/null +++ b/src/corridorkey_mlx/io/preprocess_mlx.py @@ -0,0 +1,41 @@ +"""GPU-side preprocessing for full-frame inference. + +Moves ImageNet normalization and channel concatenation to MLX while +keeping PIL bicubic resize on CPU (MLX cubic kernel differs too much). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import mlx.core as mx + +if TYPE_CHECKING: + import numpy as np + +# ImageNet normalization constants (GPU-resident) +_IMAGENET_MEAN = mx.array([0.485, 0.456, 0.406]) +_IMAGENET_STD = mx.array([0.229, 0.224, 0.225]) + + +def preprocess_mlx(rgb_f32: np.ndarray, mask_f32: np.ndarray) -> mx.array: + """GPU-side normalize + concat for full-frame inference. + + Assumes resize already happened on CPU via PIL bicubic. + + Args: + rgb_f32: (H, W, 3) float32 in [0, 1], already resized. + mask_f32: (H, W, 1) float32 in [0, 1], already resized. + + Returns: + (1, H, W, 4) mx.array — ImageNet-normalized RGB + raw alpha hint. + """ + img = mx.array(rgb_f32) + mask = mx.array(mask_f32) + + # ImageNet normalize on GPU + img = (img - _IMAGENET_MEAN) / _IMAGENET_STD + + # Concat and add batch dim + combined = mx.concatenate([img, mask], axis=-1) # (H, W, 4) + return mx.expand_dims(combined, axis=0) # (1, H, W, 4)