pixal3d_mrp_mlx/pixal3d_mlx/cond.py

"""Image -> conditioning dicts, the MLX side of upstream's `encode_image_proj`.

Produces exactly the two-key contract the DiT blocks expect:

    cond   = {'global': [B, 1+regs, D], 'proj': [B, R^3, C]}
    uncond = zeros of the same shapes        (classifier-free guidance negative)

`global` is cross-attended; `proj` is projected and ADDED, one token per voxel. The
token counts differ (5 vs R^3) and that asymmetry IS the architecture — see proj.py.
"""

from __future__ import annotations

from typing import Dict, Tuple

import mlx.core as mx
import numpy as np
from PIL import Image

from .dino import DinoV3Encoder
from .proj import ProjGrid, distance_from_fov

# The four extractor configurations upstream builds in inference.py. Only `ss` is wired
# so far; the other three additionally need the NAF upsampler, which doubles their
# proj_channels to embed_dim*2 by concatenating a high-res branch.
CONFIGS = {
    "ss":          dict(image_size=512,  grid_resolution=16, use_naf_upsample=False),
    "shape_512":   dict(image_size=512,  grid_resolution=32, use_naf_upsample=True,  naf_target_size=512),
    "shape_1024":  dict(image_size=1024, grid_resolution=64, use_naf_upsample=True,  naf_target_size=512),
    "tex_1024":    dict(image_size=1024, grid_resolution=64, use_naf_upsample=True,  naf_target_size=1024),
}


def preprocess_image(img: Image.Image, bg_color=(0, 0, 0)) -> Image.Image:
    """Upstream's `preprocess_image`, alpha path only.

    Crops to the subject's bounding box with 1.1x headroom and composites onto a flat
    background. This is not cosmetic: the camera solve assumes the object fills the
    frame (`distance_from_fov` places the camera so a unit mesh exactly spans it), so
    an uncropped image silently mis-scales the whole reconstruction.

    Background REMOVAL is not implemented — upstream calls a rembg model for images
    with no usable alpha. Such an image is passed through unchanged here, which will
    reconstruct the background along with the subject. Feed RGBA with a real matte.
    """
    has_alpha = False
    if img.mode == "RGBA":
        alpha = np.array(img)[:, :, 3]
        has_alpha = not np.all(alpha == 255)

    scale = min(1, 1024 / max(img.size))
    if scale < 1:
        img = img.resize((int(img.width * scale), int(img.height * scale)), Image.LANCZOS)
    if not has_alpha:
        return img.convert("RGB")

    arr = np.array(img)
    ys, xs = np.nonzero(arr[:, :, 3] > 0.8 * 255)
    cx, cy = (xs.min() + xs.max()) / 2, (ys.min() + ys.max()) / 2
    size = int(max(xs.max() - xs.min(), ys.max() - ys.min()) * 1.1)
    img = img.crop((cx - size // 2, cy - size // 2, cx + size // 2, cy + size // 2))

    out = np.asarray(img, dtype=np.float32) / 255.0
    rgb, a = out[:, :, :3], out[:, :, 3:4]
    bg = np.array(bg_color, dtype=np.float32) / 255.0
    return Image.fromarray((np.clip(rgb * a + bg * (1 - a), 0, 1) * 255).astype(np.uint8))


def load_image(path: str, image_size: int, preprocess: bool = True) -> np.ndarray:
    """Path -> [1,3,S,S] float32 in [0,1], preprocessed and LANCZOS-resized."""
    img = Image.open(path)
    if preprocess:
        img = preprocess_image(img)
    img = img.convert("RGB").resize((image_size, image_size), Image.LANCZOS)
    return (np.asarray(img, dtype=np.float32) / 255.0).transpose(2, 0, 1)[None]


class ProjConditioner:
    """One extractor stage: DINOv3 + a ProjGrid at that stage's voxel resolution."""

    def __init__(self, stage: str = "ss", device: str | None = None):
        cfg = CONFIGS[stage]
        self.stage = stage
        self.image_size = cfg["image_size"]
        self.grid_resolution = cfg["grid_resolution"]
        self.use_naf = bool(cfg.get("use_naf_upsample"))
        self.naf_target = cfg.get("naf_target_size")
        self.device = device
        self.encoder = DinoV3Encoder(image_size=self.image_size, device=device)
        # ONE grid for both branches, at image_size — upstream reuses `self.proj_grid`
        # for the high-res sample too. That is not interchangeable with a grid built at
        # naf_target_size: the normalised coordinate carries a 1/resolution term, so
        # projecting at 1024 and normalising by 1024 lands ~0.001 off in [-1,1] — a
        # sub-pixel shift on every voxel, in a model whose whole point is pixel
        # alignment. Sampling is resolution-agnostic anyway; grid_sample maps [-1,1]
        # onto whatever feature map it is handed.
        self.proj_grid = ProjGrid(grid_resolution=self.grid_resolution,
                                  image_resolution=self.image_size)
        self.proj_channels = self.encoder.embed_dim * (2 if self.use_naf else 1)

    def _grid(self, grid_resolution: int | None):
        """ProjGrid at `grid_resolution`, cached. Always at image_size (see __init__)."""
        r = grid_resolution or self.grid_resolution
        if r == self.grid_resolution:
            return self.proj_grid
        cache = getattr(self, "_grid_cache", None) or {}
        if r not in cache:
            cache[r] = ProjGrid(grid_resolution=r, image_resolution=self.image_size)
            self._grid_cache = cache
        return cache[r]

    def __call__(self, image_chw01: np.ndarray, camera_angle_x: float,
                 distance: float | None = None, mesh_scale: float = 1.0,
                 grid_resolution: int | None = None,
                 ) -> Tuple[Dict[str, mx.array], Dict[str, mx.array]]:
        """`grid_resolution` overrides this stage's lattice — the cascade backs the
        high-res grid off in steps when an object is too dense for max_num_tokens."""
        if distance is None:
            distance = distance_from_fov(camera_angle_x, mesh_scale, self.image_size)

        grid = self._grid(grid_resolution)
        z_global, z_patch = self.encoder(image_chw01)
        z_proj = grid(z_patch, camera_angle_x, distance, mesh_scale)

        if self.use_naf:
            from .naf import upsample
            import numpy as _np
            from PIL import Image as _Image

            # NAF's guide is the UNNORMALISED image at the target size
            guide = image_chw01
            if guide.shape[-1] != self.naf_target:
                arr = (guide[0].transpose(1, 2, 0) * 255).astype(_np.uint8)
                arr = _Image.fromarray(arr).resize((self.naf_target, self.naf_target),
                                                   _Image.LANCZOS)
                guide = (_np.asarray(arr, dtype=_np.float32) / 255.0).transpose(2, 0, 1)[None]

            hr = upsample(z_patch, guide, self.naf_target,
                          device=self.device or "mps")
            z_hr = grid(hr, camera_angle_x, distance, mesh_scale)
            z_proj = mx.concatenate([z_proj, z_hr], axis=-1)

        cond = {"global": z_global, "proj": z_proj}
        uncond = {"global": mx.zeros_like(z_global), "proj": mx.zeros_like(z_proj)}
        return cond, uncond