heatmap_nextgen.py

"""Decision-linked localization heatmap for the A-EYE backend (model 54+ family).

The heatmap is the PatchGuard patch-probability map, refined for localization:
  * computed at higher resolution (24x24) for a finer, less blocky map,
  * sharpened by a noise-residual map (insertions disturb camera noise),
  * per-image normalized and focused to the single strongest connected region
    (kills scattered noise so only the suspected insert lights up),
  * intensity-gated by the calibrated image confidence, so weak/uncertain maps
    stay faint instead of painting the photo with noise.
Pure numpy / PIL / scipy. New file; nothing existing is modified.
"""
from __future__ import annotations

import numpy as np
from PIL import Image
from scipy.ndimage import gaussian_filter, label, uniform_filter


def jet(values: np.ndarray) -> np.ndarray:
    v = np.clip(values, 0.0, 1.0)
    r = np.clip(1.5 - np.abs(4.0 * v - 3.0), 0.0, 1.0)
    g = np.clip(1.5 - np.abs(4.0 * v - 2.0), 0.0, 1.0)
    b = np.clip(1.5 - np.abs(4.0 * v - 1.0), 0.0, 1.0)
    return np.stack([r, g, b], axis=-1)


def residual_var(image: Image.Image, grid: int, win: int = 10) -> np.ndarray:
    """Local noise-residual variance, block-reduced to grid x grid. Camera regions
    carry consistent sensor noise; AI-inserted regions usually break it, so this
    helps pin the patch map onto the real seam."""
    g = np.asarray(image.convert("L").resize((288, 288)), np.float32)
    res = g - gaussian_filter(g, 2)
    m2 = uniform_filter(res * res, win)
    m1 = uniform_filter(res, win)
    var = np.maximum(m2 - m1 * m1, 0.0)
    h, w = var.shape
    bh, bw = h // grid, w // grid
    return var[: bh * grid, : bw * grid].reshape(grid, bh, grid, bw).mean(axis=(1, 3))


def _focus(loc: np.ndarray) -> np.ndarray:
    """Per-image normalize, then keep only the strongest connected blob."""
    rng = float(np.ptp(loc))
    if rng < 1e-6:
        return np.zeros_like(loc)
    n = (loc - loc.min()) / (rng + 1e-6)
    binary = n > 0.55
    lab, k = label(binary)
    if k > 1:
        sums = [float((loc * (lab == i)).sum()) for i in range(1, k + 1)]
        keep = 1 + int(np.argmax(sums))
        n = n * (lab == keep)
    elif k == 0:
        n = n * 0.0
    return n


def _big(focus_map: np.ndarray, size: tuple[int, int]) -> np.ndarray:
    """Upsample the small focus map to image size, feathered for clean edges."""
    img = Image.fromarray((np.clip(focus_map, 0.0, 1.0) * 255).astype(np.uint8))
    big = np.asarray(img.resize(size, Image.Resampling.BICUBIC), np.float32) / 255.0
    return gaussian_filter(big, max(1.0, size[0] / 130.0))


def pure_heatmap(loc: np.ndarray, size: tuple[int, int], conf: float = 1.0, blanket: bool = False) -> Image.Image:
    """Standalone jet heatmap (no original image)."""
    if blanket:
        big = np.full((size[1], size[0]), 0.9, np.float32)
    else:
        big = _big(_focus(loc), size) * float(np.clip(conf, 0.0, 1.0))
    return Image.fromarray((jet(big) * 255).astype(np.uint8))


def overlay(
    image: Image.Image,
    loc: np.ndarray,
    conf: float = 1.0,
    blanket: bool = False,
    floor: float = 0.30,
    max_alpha: float = 0.92,
    gamma: float = 0.55,
) -> Image.Image:
    """Jet overlay on the photo. The hot region is rendered VIVID (the per-image
    map is normalized so its peak is full red at max_alpha). `blanket=True` paints
    the whole image red. `conf` is kept for API compatibility but the caller passes
    1.0 so the color is always strong, per product preference."""
    rgb = np.asarray(image.convert("RGB"), np.float32) / 255.0
    h, w = rgb.shape[:2]
    if blanket:
        big = np.full((h, w), 0.9, np.float32)
        c = 1.0
    else:
        big = _big(_focus(loc), (w, h))
        c = float(np.clip(conf, 0.0, 1.0))
    norm = np.clip((big - floor) / (1.0 - floor), 0.0, 1.0)
    alpha = max_alpha * c * (norm ** gamma)
    blended = rgb * (1.0 - alpha[..., None]) + jet(big) * alpha[..., None]
    return Image.fromarray((np.clip(blended, 0.0, 1.0) * 255).astype(np.uint8))