feat: GPU-side ImageNet normalize + concat for full-frame path
PIL bicubic vs MLX cubic diff too large (0.22-0.59) — resize stays on CPU. Only normalize+concat moved to MLX. Tiled path unchanged. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
f80bec92b6
commit
31d83e218b
@ -26,6 +26,7 @@ from corridorkey_mlx.io.image import (
|
|||||||
postprocess_foreground,
|
postprocess_foreground,
|
||||||
preprocess,
|
preprocess,
|
||||||
)
|
)
|
||||||
|
from corridorkey_mlx.io.preprocess_mlx import preprocess_mlx
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
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
|
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)
|
outputs = self._model(x)
|
||||||
# NOTE: mx.eval is MLX array materialization, not Python eval()
|
# NOTE: mx.eval is MLX array materialization, not Python eval()
|
||||||
mx.eval(outputs) # noqa: S307
|
mx.eval(outputs) # noqa: S307
|
||||||
|
|||||||
41
src/corridorkey_mlx/io/preprocess_mlx.py
Normal file
41
src/corridorkey_mlx/io/preprocess_mlx.py
Normal file
@ -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)
|
||||||
Loading…
Reference in New Issue
Block a user