core/heightmap.py

"""Detail-injection relief pipeline (§5.2) -- the core value of DepthForge.

A raw monocular depth map only gives the smooth *form*. The carveable detail
(wrinkles, hair, fabric folds, stitching) comes from blending the photo's own
high-frequency luminance texture onto that form. This module ports the proven
Colab-spike math VERBATIM. Do not "improve" the constants or reorder operations.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Tuple

import numpy as np
from PIL import Image
from scipy.ndimage import gaussian_filter

from .config import RELIEF


@dataclass
class ReliefParams:
    """Parameters for :func:`build_heightmap`.

    Mirrors the spike's ``params`` bag. ``invert`` orients the depth so near->high;
    its correct default depends on the depth backend (DA3: True, V2: False).
    """

    invert: bool = RELIEF.invert
    detail_strength: float = RELIEF.detail_strength
    detail_sigma: float = RELIEF.detail_sigma
    fine_strength: float = RELIEF.fine_strength
    contrast_gamma: float = RELIEF.contrast_gamma
    mask_thresh: float = RELIEF.mask_thresh
    bg_flatten: bool = RELIEF.bg_flatten
    smooth: float = RELIEF.smooth
    # Foreground depth expansion (the "Sculpt" fix): isolate the near subject
    # and re-stretch depth *within* it so the subject's own form fills the
    # relief, collapsing the far background into a shallow slab. Without this, a
    # scene with a distant background crushes the subject to a flat plateau.
    fg_expand: bool = True
    fg_bg_level: float = 0.15   # background slab height (0..1); subject sits above
    fg_local: float = 2.0       # depth-domain local-contrast (unsharp) on subject


def _otsu_threshold(x: np.ndarray) -> float:
    """Otsu threshold on values in ``[0, 1]`` (splits near subject from far bg)."""
    hist, _ = np.histogram(x.ravel(), bins=256, range=(0.0, 1.0))
    hist = hist.astype(np.float64)
    total = hist.sum()
    if total <= 0:
        return 0.5
    p = hist / total
    centers = (np.arange(256) + 0.5) / 256.0
    omega = np.cumsum(p)
    mu = np.cumsum(p * centers)
    mu_t = mu[-1]
    denom = omega * (1.0 - omega)
    denom[denom == 0] = 1e-12
    sigma_b2 = (mu_t * omega - mu) ** 2 / denom
    return float((np.argmax(sigma_b2) + 0.5) / 256.0)


def build_heightmap(
    depth_raw: np.ndarray,
    cleaned_rgb: Image.Image,
    params: ReliefParams,
) -> Tuple[np.ndarray, np.ndarray]:
    """Blend detail texture onto the depth form to make a carveable relief.

    Args:
        depth_raw: Raw depth from a backend, float array ``[H, W]`` (any scale;
            may contain NaN/inf). Resized to the cleaned image if needed.
        cleaned_rgb: The cleaned source image (supplies high-frequency detail).
        params: Relief parameters (see :class:`ReliefParams`).

    Returns:
        Tuple of ``(height_float, height_uint16)`` where ``height_float`` is in
        ``[0, 1]`` and ``height_uint16`` is the 16-bit heightmap for PNG export.
    """
    invert = params.invert
    detail_strength = params.detail_strength
    detail_sigma = params.detail_sigma
    fine_strength = params.fine_strength
    contrast_gamma = params.contrast_gamma
    mask_thresh = params.mask_thresh
    bg_flatten = params.bg_flatten
    smooth = params.smooth
    fg_expand = getattr(params, "fg_expand", False)
    fg_bg_level = getattr(params, "fg_bg_level", 0.15)
    fg_local = getattr(params, "fg_local", 0.0)

    # --- verbatim spike math -------------------------------------------------
    rgb = np.asarray(cleaned_rgb.convert("RGB"), dtype=np.float32) / 255.0
    H, W = rgb.shape[:2]
    L = 0.299 * rgb[..., 0] + 0.587 * rgb[..., 1] + 0.114 * rgb[..., 2]

    base = np.asarray(depth_raw, dtype=np.float32)
    fin = base[np.isfinite(base)]
    base = np.nan_to_num(
        base,
        nan=float(np.median(fin)),
        posinf=float(fin.max()),
        neginf=float(fin.min()),
    )
    if invert:
        base = base.max() - base
    if base.shape[:2] != (H, W):
        base = np.asarray(
            Image.fromarray(base, mode="F").resize((W, H), Image.BICUBIC),
            np.float32,
        )
    # Robust global normalize to [0,1]. Wide percentiles keep the subject's own
    # (small) depth range intact rather than clipping it flat.
    blo, bhi = np.percentile(base, [1, 99])
    d = np.clip((base - blo) / (bhi - blo + 1e-8), 0, 1)

    if fg_expand:
        # Split near subject (high depth) from far background via Otsu, then
        # re-stretch depth *inside* the subject so its face/folds fill the
        # relief. The far background is compressed into a shallow slab below.
        t = _otsu_threshold(d)
        fg = d >= t
        frac = float(fg.mean())
        if frac < 0.03 or frac > 0.9:  # degenerate split -> percentile fallback
            t = float(np.percentile(d, 70))
            fg = d >= t
            frac = float(fg.mean())
        if frac >= 0.005:
            flo, fhi = np.percentile(d[fg], [2, 98])
        else:
            flo, fhi = t, 1.0
        subj = np.clip((d - flo) / (fhi - flo + 1e-8), 0, 1)  # expanded subject
        if fg_local > 0:
            # Depth-domain unsharp: amplify the subject's OWN local relief
            # (nose, eyes, folds) that global expansion leaves low-contrast.
            sig = max(6.0, min(H, W) / 60.0)
            subj = np.clip(subj + fg_local * (subj - gaussian_filter(subj, sigma=sig)), 0, 1)
        soft = np.clip(gaussian_filter(fg.astype(np.float32), sigma=6), 0, 1)
        bg = np.clip(d / (t + 1e-8), 0, 1) * fg_bg_level       # shallow far slab
        base = soft * (fg_bg_level + subj * (1.0 - fg_bg_level)) + (1.0 - soft) * bg
    else:
        base = d

    mask = gaussian_filter((base > mask_thresh).astype(np.float32), sigma=4)

    detail_coarse = L - gaussian_filter(L, sigma=detail_sigma)
    detail_fine = L - gaussian_filter(L, sigma=max(1, detail_sigma / 3))
    detail = (detail_strength * detail_coarse + fine_strength * detail_fine) * mask

    h = base + detail
    if bg_flatten:
        h = h * mask
    if smooth > 0:
        h = gaussian_filter(h, sigma=smooth)

    h = np.clip(h, 0, 1)
    h = (h - h.min()) / (h.max() - h.min() + 1e-8)
    h = np.power(h, contrast_gamma)
    height16 = (np.clip(h, 0, 1) * 65535).astype(np.uint16)
    return h, height16