helpers/geometry.py

"""
Step 2 — keypoint selection.

`organoid` asks the vision model to invent coordinates from a lookup table
("centred -> (w/2, h/2); upper-left -> (w/4, h/4)"). That is a nine-cell
quantisation of the image and it breaks on any concave object: the centroid of
a C-shaped or L-shaped fossil can sit in the background, so SAM3 gets a positive
point that is not on the specimen at all.

This module proposes points *geometrically* from the alpha matte instead, and
then hands them to the vision model to review. Two guarantees fall out of that:

  1. Every positive point is inside the mask by construction — the anchor is the
     maximum of the Euclidean distance transform, i.e. the point furthest from
     any boundary. For a crescent-shaped fossil that lands in the thickest part
     of the crescent, never in the hollow.
  2. Negative points come for free. SAM3 accepts label-0 prompts, and the single
     biggest failure mode when photographing a specimen on a desk is the model
     annexing the cast shadow. Seeding negatives into the shadow band suppresses
     that directly.

The vision model's job becomes *review* (accept / move / add / drop), which is
what current models are actually reliable at, rather than *estimation*, which
they are not.
"""

from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np

from helpers.photometry import shadow_band

from config.settings import CORE_DT_FRACTION, N_NEGATIVE_POINTS, N_POSITIVE_POINTS


def _farthest_point_sample(candidates: np.ndarray, seed: tuple[int, int], k: int) -> list[tuple[int, int]]:
    """Greedy farthest-point sampling — spreads points across the specimen."""
    chosen = [np.array(seed, dtype=np.float32)]
    if len(candidates) == 0:
        return [tuple(int(v) for v in seed)]

    for _ in range(max(0, k - 1)):
        d = np.min(
            np.linalg.norm(candidates[:, None, :] - np.stack(chosen)[None], axis=2),
            axis=1,
        )
        chosen.append(candidates[int(np.argmax(d))].astype(np.float32))

    return [(int(p[0]), int(p[1])) for p in chosen]


def propose_keypoints(
    image_path: str | Path,
    alpha_mask_path: str | Path,
    n_positive: int = N_POSITIVE_POINTS,
    n_negative: int = N_NEGATIVE_POINTS,
) -> dict:
    """
    Derive positive (inside-fossil) and negative (background/shadow) prompt points.

    Returns a dict with `positive_points` / `negative_points` as lists of
    `{"x": int, "y": int, "label": str, "source": "geometric"}`.
    """
    bgr = cv2.imread(str(image_path))
    if bgr is None:
        raise FileNotFoundError(f"Could not read image: {image_path}")

    alpha = cv2.imread(str(alpha_mask_path), cv2.IMREAD_GRAYSCALE)
    if alpha is None:
        raise FileNotFoundError(f"Could not read alpha matte: {alpha_mask_path}")

    binary = (alpha > 127).astype(np.uint8) * 255

    # --- positives: distance-transform anchor + spread ---------------------
    dt = cv2.distanceTransform(binary, cv2.DIST_L2, 5)
    if dt.max() <= 0:
        raise RuntimeError("Alpha matte is empty — nothing to place keypoints on.")

    anchor_y, anchor_x = np.unravel_index(int(np.argmax(dt)), dt.shape)

    core = (dt > CORE_DT_FRACTION * dt.max()).astype(np.uint8)
    ys, xs = np.nonzero(core)
    candidates = np.stack([xs, ys], axis=1)

    positives = _farthest_point_sample(candidates, (int(anchor_x), int(anchor_y)), n_positive)

    # --- negatives: shadow band first, then plain background ---------------
    h, w = binary.shape
    ring_kernel = np.ones((61, 61), np.uint8)
    ring = cv2.bitwise_and(cv2.dilate(binary, ring_kernel), cv2.bitwise_not(binary))

    # One shared definition of "shadow" — see helpers/photometry.py.
    shadow = shadow_band(bgr, binary).astype(np.uint8) * 255

    n_shadow = min(n_negative // 2, 3)
    negatives: list[tuple[int, int]] = []
    for mask, want in ((shadow, n_shadow), (ring, n_negative - n_shadow)):
        yy, xx = np.nonzero(mask)
        if len(yy) == 0 or want <= 0:
            continue
        idx = np.linspace(0, len(yy) - 1, want).astype(int)
        negatives += [(int(xx[i]), int(yy[i])) for i in idx]

    # Deduplicate while preserving order.
    negatives = list(dict.fromkeys(negatives))

    shadow_px = int((shadow > 0).sum())
    notes = (
        f"Anchor placed at the distance-transform maximum "
        f"(depth {dt.max():.0f}px from the nearest boundary). "
        f"{len(positives)} positive, {len(negatives)} negative points. "
        + (
            f"A shadow band of {shadow_px} px was detected and seeded with "
            f"{n_shadow} negative points."
            if shadow_px > 0
            else "No cast shadow detected around the specimen."
        )
    )

    return {
        "status": "success",
        "image_path": str(image_path),
        "dimensions": {"width": int(w), "height": int(h)},
        "positive_points": [
            {"x": x, "y": y, "label": "positive", "source": "geometric"} for x, y in positives
        ],
        "negative_points": [
            {"x": x, "y": y, "label": "negative", "source": "geometric"} for x, y in negatives
        ],
        "notes": notes,
    }