helpers/validate.py

"""
Point validation — the agent's feedback loop.

The agent chooses the keypoints. This module does not choose any. It only
*measures* what the agent chose and reports back, so the agent can revise before
paying for a SAM3 pass:

  * is the point inside the image at all?
  * what is under it — specimen, backdrop, or cast shadow? (measured from the
    step-1 alpha matte and from helpers/photometry.py, not guessed)
  * how far is it from the nearest boundary? A positive point 3 px from the edge
    is technically inside and practically useless.

The verdicts are advisory. `agrees_with_step1: false` is a disagreement between
the agent and the background remover, and the agent is allowed to win — it can
see the photograph and the matte cannot. What it is not allowed to do is emit a
point off-canvas, or claim a positive point on a region it also called negative.
"""

from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np
from PIL import Image

from helpers.images import load_rgb
from helpers.photometry import shadow_band

# A positive point closer than this (as a fraction of the image's shorter side)
# to the mask boundary is flagged as precarious.
_EDGE_MARGIN_FRAC = 0.01


def validate_points(
    image_path: str | Path,
    alpha_mask_path: str | Path | None,
    positive_points: list[dict],
    negative_points: list[dict],
) -> dict:
    """
    Measure the agent's chosen points against the image. Chooses nothing.

    Returns {ok, dimensions, positives: [...verdicts], negatives: [...],
             errors: [...], warnings: [...], summary}.
    Only `errors` are blocking. Warnings are for the agent to weigh.
    """
    rgb_img = load_rgb(image_path)
    w, h = rgb_img.size
    rgb = np.array(rgb_img)

    if alpha_mask_path and Path(alpha_mask_path).exists():
        alpha = np.array(Image.open(alpha_mask_path).convert("L").resize((w, h)))
        fg = alpha > 127
    else:
        fg = np.zeros((h, w), bool)

    shadow = shadow_band(cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR), fg)
    # Distance from every pixel to the nearest non-specimen pixel.
    depth = cv2.distanceTransform(fg.astype(np.uint8), cv2.DIST_L2, 5)
    margin = max(2.0, _EDGE_MARGIN_FRAC * min(w, h))

    errors: list[str] = []
    warnings: list[str] = []

    def _verdict(p: dict, want_inside: bool, tag: str) -> dict:
        x, y = int(p["x"]), int(p["y"])
        if not (0 <= x < w and 0 <= y < h):
            errors.append(f"{tag} ({x},{y}) is outside the image (bounds 0..{w-1}, 0..{h-1}).")
            return {"x": x, "y": y, "in_bounds": False}

        inside = bool(fg[y, x])
        in_shadow = bool(shadow[y, x])
        d = float(depth[y, x])
        region = "specimen" if inside else ("cast_shadow" if in_shadow else "backdrop")

        v = {
            "x": x,
            "y": y,
            "in_bounds": True,
            "region_per_step1": region,
            "px_from_boundary": round(d, 1),
            "rgb_under_point": [int(c) for c in rgb[y, x]],
            "agrees_with_step1": inside == want_inside,
        }

        if want_inside and not inside:
            warnings.append(
                f"{tag} ({x},{y}) is a POSITIVE point, but step 1 calls that pixel "
                f"{region} (RGB {v['rgb_under_point']}). Three possibilities, and only "
                f"you can tell which: (a) step 1 clipped the specimen there — common, "
                f"keep the point; (b) this is ROCK MATRIX the fossil is embedded in and "
                f"you are correctly telling SAM3 the whole piece is one object — keep "
                f"the point; or (c) the point is genuinely off in the backdrop and "
                f"should move. Do NOT reflexively delete it: on an embedded specimen, "
                f"points out on the matrix are how the matrix stays in the mask."
            )
        if want_inside and inside and d < margin:
            warnings.append(
                f"{tag} ({x},{y}) is only {d:.0f}px from the specimen edge. SAM3 "
                f"prompts near a boundary are ambiguous; move it toward the interior."
            )
        if (not want_inside) and inside:
            warnings.append(
                f"{tag} ({x},{y}) is a NEGATIVE point but sits {d:.0f}px INSIDE the "
                f"specimen per step 1. A negative on the fossil punches a hole in the "
                f"mask. Move it or drop it."
            )
        return v

    pos = [_verdict(p, True, f"P{i+1}") for i, p in enumerate(positive_points)]
    neg = [_verdict(p, False, f"N{i+1}") for i, p in enumerate(negative_points)]

    if not positive_points:
        errors.append("No positive points. SAM3 needs at least one point on the specimen.")

    seen: dict[tuple[int, int], str] = {}
    for label, pts in (("positive", positive_points), ("negative", negative_points)):
        for p in pts:
            key = (int(p["x"]), int(p["y"]))
            if key in seen and seen[key] != label:
                errors.append(f"({key[0]},{key[1]}) is listed as both positive and negative.")
            seen[key] = label

    shadow_px = int(shadow.sum())
    if shadow_px > 0.005 * w * h and not any(
        n.get("region_per_step1") == "cast_shadow" for n in neg
    ):
        warnings.append(
            f"A cast shadow of ~{shadow_px} px was detected and NONE of your negative "
            f"points are in it. The shadow is the single most common thing SAM3 "
            f"wrongly annexes into the specimen. Put a negative point in it."
        )

    return {
        "ok": not errors,
        "dimensions": {"width": w, "height": h},
        "positives": pos,
        "negatives": neg,
        "errors": errors,
        "warnings": warnings,
        "shadow_pixels_detected": shadow_px,
        "summary": (
            f"{len(pos)} positive, {len(neg)} negative. "
            f"{len(errors)} blocking error(s), {len(warnings)} warning(s)."
        ),
    }