accessibilityamodal/visual_completion.py

"""Mask policy, prompts, and quality gates for visual accessibility completion.

The geometry target and the visual removal region have deliberately different
roles.  ``hidden`` remains the geometry target.  For visual inpainting, an
entire foreground obstacle instance is removed when any pixel in its
8-connected component intersects ``hidden``.  Nearby people or objects that do
not intersect ``hidden`` remain untouched.
"""

from __future__ import annotations

import math
from typing import Any

import cv2
import numpy as np
from PIL import Image


PROMPTS = {
    "stairs": (
        "photorealistic continuation of the same staircase behind the removed foreground "
        "occluder, continuous stair treads aligned with the visible steps, same perspective, "
        "same material and texture, same lighting and exposure, empty completed surface"
    ),
    "ramp": (
        "photorealistic continuation of the same accessible ramp behind the removed foreground "
        "occluder, continuous sloped walking surface, same perspective, same material and "
        "texture, same lighting and exposure, empty completed surface"
    ),
    "curb_cut": (
        "photorealistic continuation of the same curb cut and pavement transition behind the "
        "removed foreground occluder, same perspective, same material and texture, same lighting "
        "and exposure, continuous empty completed surface"
    ),
    "raised_curb": (
        "photorealistic continuation of the same raised curb behind the removed foreground "
        "occluder, one continuous level curb top and straight curb face, same perspective, "
        "same material and texture, same lighting and exposure, no added steps"
    ),
    "tactile_paving": (
        "photorealistic continuation of the same tactile paving path behind the removed "
        "foreground occluder, regularly aligned tactile pattern, same perspective, same material "
        "and texture, same lighting and exposure, continuous empty completed surface"
    ),
    "walkway": (
        "photorealistic continuation of the same accessible pedestrian walkway behind the "
        "removed foreground occluder, continuous walking surface, same perspective, same material "
        "and texture, same lighting and exposure, empty completed surface"
    ),
}

NEGATIVE_PROMPT = (
    "person, human, legs, pedestrian, bicycle, wheel, wheelchair, stroller, walker, cart, "
    "luggage, bag, backpack, cane, obstacle, vehicle, animal, fog, haze, blur, melted "
    "geometry, warped stairs, broken steps, duplicate steps, misaligned edges, text, "
    "watermark, illustration"
)

FOG_HAZE_MAX_SHARPNESS_RATIO = 0.65
FOG_HAZE_MAX_CONTRAST_RATIO = 0.85
OVER_SHARP_MIN_SHARPNESS_RATIO = 1.9
OVER_SHARP_MIN_SEAM_PENALTY = 1.45


def _binary_array(mask: np.ndarray, name: str) -> np.ndarray:
    array = np.asarray(mask)
    if array.ndim != 2:
        raise ValueError(f"{name} must be a 2D mask, got shape={array.shape}")
    return array.astype(bool, copy=False)


def derive_hidden_mask(target_amodal: np.ndarray, target_visible: np.ndarray) -> np.ndarray:
    """Derive the geometry hidden target without changing either input mask."""

    amodal = _binary_array(target_amodal, "target_amodal")
    visible = _binary_array(target_visible, "target_visible")
    if amodal.shape != visible.shape:
        raise ValueError("Target amodal and visible masks must have the same shape")
    return amodal & ~visible


def mask_statistics(mask: np.ndarray) -> dict[str, int | float]:
    """Return path-free mask statistics suitable for a portable manifest."""

    binary = _binary_array(mask, "mask")
    height, width = binary.shape
    pixels = int(binary.sum())
    image_pixels = int(binary.size)
    return {
        "width": int(width),
        "height": int(height),
        "image_pixels": image_pixels,
        "pixel_count": pixels,
        "image_fraction": round(pixels / image_pixels, 8) if image_pixels else 0.0,
    }


def _retain_components_occluding_hidden(
    detected_obstacles: np.ndarray,
    hidden: np.ndarray,
) -> tuple[np.ndarray, int, int]:
    """Equivalent policy to filter_accessibility_occluding_obstacles."""

    if detected_obstacles.shape != hidden.shape:
        raise ValueError("Obstacle and hidden masks must have the same shape")
    count, labels = cv2.connectedComponents(
        detected_obstacles.astype(np.uint8),
        connectivity=8,
    )
    keep_labels = np.unique(labels[hidden & detected_obstacles])
    keep_labels = keep_labels[keep_labels != 0]
    return np.isin(labels, keep_labels), int(count - 1), int(len(keep_labels))


def build_visual_removal_mask(
    hidden: np.ndarray,
    obstacle: np.ndarray | None = None,
) -> tuple[np.ndarray, dict[str, Any]]:
    """Build the visual inpaint mask while preserving ``hidden`` for geometry.

    When ``obstacle`` is omitted, the returned visual mask is exactly ``hidden``
    for backward compatibility.
    """

    geometry_hidden = _binary_array(hidden, "hidden")
    if obstacle is None:
        detected = np.zeros_like(geometry_hidden)
        retained = np.zeros_like(geometry_hidden)
        before_components = 0
        retained_components = 0
    else:
        detected = _binary_array(obstacle, "obstacle")
        if detected.shape != geometry_hidden.shape:
            raise ValueError("Obstacle and hidden masks must have the same shape")
        retained, before_components, retained_components = (
            _retain_components_occluding_hidden(detected, geometry_hidden)
        )

    visual_removal = geometry_hidden | retained
    stats = {
        "policy": (
            "hidden_union_full_8_connected_obstacle_components_intersecting_hidden"
            if obstacle is not None
            else "legacy_hidden_only"
        ),
        "obstacle_mask_provided": obstacle is not None,
        "geometry_hidden_unchanged": True,
        "geometry_hidden": mask_statistics(geometry_hidden),
        "obstacle_input": mask_statistics(detected),
        "obstacle_components_input": before_components,
        "obstacle_retained": mask_statistics(retained),
        "obstacle_components_retained": retained_components,
        "non_occluding_obstacle_pixels_excluded": int((detected & ~retained).sum()),
        "visual_removal": mask_statistics(visual_removal),
        "visual_extra_pixels_beyond_hidden": int((visual_removal & ~geometry_hidden).sum()),
    }
    return visual_removal, stats


def build_completion_envelope(
    visual_removal: np.ndarray,
    obstacle: np.ndarray | None,
    *,
    margin_fraction: float = 0.022,
) -> tuple[np.ndarray, dict[str, Any]]:
    """Fill retained foreground-instance boxes before generative completion.

    A person mask often excludes a carried bag, walker, bicycle frame, or the
    small gaps between limbs.  Inpainting only the segmentation silhouette can
    therefore preserve or regenerate those objects.  This appearance-only mask
    fills the bounding box of each retained obstacle component and adds a small
    image-relative margin.  Geometry continues to use the unchanged hidden
    target; non-occluding obstacle pixels are protected by the caller.
    """

    removal = _binary_array(visual_removal, "visual_removal")
    if margin_fraction < 0:
        raise ValueError("margin_fraction must be non-negative")
    if obstacle is None:
        return removal.copy(), {
            "policy": "visual_removal_without_obstacle_envelope",
            "component_count": 0,
            "margin_pixels": 0,
            "extra_pixels": 0,
            "completion_envelope": mask_statistics(removal),
        }

    detected = _binary_array(obstacle, "obstacle")
    if detected.shape != removal.shape:
        raise ValueError("Obstacle and visual removal masks must have the same shape")
    retained_obstacle = detected & removal
    count, labels, stats, _ = cv2.connectedComponentsWithStats(
        retained_obstacle.astype(np.uint8),
        connectivity=8,
    )
    height, width = removal.shape
    margin = int(round(min(height, width) * margin_fraction))
    envelope = removal.copy()
    boxes: list[dict[str, int]] = []
    for label in range(1, count):
        x, y, box_width, box_height, area = (
            int(value) for value in stats[label]
        )
        if area <= 0:
            continue
        x1 = max(0, x - margin)
        y1 = max(0, y - margin)
        x2 = min(width, x + box_width + margin)
        y2 = min(height, y + box_height + margin)
        envelope[y1:y2, x1:x2] = True
        boxes.append(
            {
                "x1": x1,
                "y1": y1,
                "x2": x2,
                "y2": y2,
                "source_component_pixels": area,
            }
        )
    return envelope, {
        "policy": "retained_obstacle_component_bounding_envelopes",
        "component_count": len(boxes),
        "margin_pixels": margin,
        "boxes": boxes,
        "extra_pixels": int((envelope & ~removal).sum()),
        "completion_envelope": mask_statistics(envelope),
    }


def quality_flags_for_metrics(
    *,
    sharpness_ratio: float,
    contrast_ratio: float,
    seam_penalty: float,
) -> list[str]:
    """Return deterministic visual-risk flags for candidate metrics."""

    flags: list[str] = []
    if (
        sharpness_ratio < FOG_HAZE_MAX_SHARPNESS_RATIO
        and contrast_ratio < FOG_HAZE_MAX_CONTRAST_RATIO
    ):
        flags.append("fog_haze")
    if (
        sharpness_ratio > OVER_SHARP_MIN_SHARPNESS_RATIO
        and seam_penalty > OVER_SHARP_MIN_SEAM_PENALTY
    ):
        flags.append("over_sharp_foreground_artifact")
    return flags


def candidate_quality(
    image: Image.Image,
    mask: Image.Image,
    category: str,
) -> dict[str, Any]:
    """Score a visual candidate and attach conservative review-gate signals."""

    rgb = np.asarray(image.convert("RGB"), dtype=np.uint8)
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY).astype(np.float32)
    binary = (np.asarray(mask.convert("L")) > 127).astype(np.uint8)
    if binary.shape != gray.shape:
        raise ValueError(
            f"Candidate/mask raster mismatch: image={gray.shape}, mask={binary.shape}"
        )
    if not binary.any():
        raise ValueError("Candidate quality requires a nonempty visual removal mask")

    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (17, 17))
    outer = (cv2.dilate(binary, kernel) > 0) & ~(binary > 0)
    inner = (binary > 0) & ~(cv2.erode(binary, kernel) > 0)
    interior = cv2.erode(binary, np.ones((5, 5), np.uint8)) > 0
    if not interior.any():
        interior = binary > 0
    if not outer.any():
        outer = ~(binary > 0)
    if not outer.any():
        outer = np.ones_like(binary, dtype=bool)

    gx = np.abs(cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3))
    gy = np.abs(cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3))
    laplacian = np.abs(cv2.Laplacian(gray, cv2.CV_32F, ksize=3))

    eps = 1e-6
    reference_sharpness = float(np.mean(laplacian[outer])) + eps
    sharpness_ratio = float(np.mean(laplacian[interior])) / reference_sharpness
    reference_contrast = float(np.std(gray[outer])) + eps
    contrast_ratio = float(np.std(gray[interior])) / reference_contrast
    seam_energy = float(np.mean((gx + gy)[inner])) if inner.any() else 0.0
    reference_edge = float(np.mean((gx + gy)[outer])) + eps
    seam_penalty = seam_energy / reference_edge

    horizontal_fraction = float(np.mean(gy[interior])) / (
        float(np.mean(gx[interior]) + np.mean(gy[interior])) + eps
    )
    structure_bonus = horizontal_fraction if category == "stairs" else 0.5
    sharp_term = min(sharpness_ratio, 1.8) / 1.8
    contrast_term = min(contrast_ratio, 1.5) / 1.5
    seam_term = math.exp(-max(0.0, seam_penalty - 1.0))
    score = (
        0.38 * sharp_term
        + 0.24 * contrast_term
        + 0.23 * structure_bonus
        + 0.15 * seam_term
    )
    quality_flags = quality_flags_for_metrics(
        sharpness_ratio=sharpness_ratio,
        contrast_ratio=contrast_ratio,
        seam_penalty=seam_penalty,
    )

    # A flagged candidate must rank below every unflagged candidate.  Keeping
    # the unbounded negative value preserves useful ordering when all
    # candidates are risky and must be withheld for review.
    gate_score = float(score) - float(len(quality_flags))
    return {
        "score": round(float(score), 6),
        "gate_score": round(gate_score, 6),
        "quality_flags": quality_flags,
        "review_required": bool(quality_flags),
        "sharpness_ratio": round(sharpness_ratio, 6),
        "contrast_ratio": round(contrast_ratio, 6),
        "horizontal_edge_fraction": round(horizontal_fraction, 6),
        "seam_penalty": round(seam_penalty, 6),
    }


def candidate_clutter_metrics(
    image: Image.Image,
    original: Image.Image,
    completion_envelope: np.ndarray,
    target_amodal: np.ndarray | None,
) -> dict[str, Any]:
    """Measure new line/edge clutter outside the modeled support target."""

    envelope = _binary_array(completion_envelope, "completion_envelope")
    if target_amodal is None:
        return {
            "available": False,
            "reason": "target_amodal_unavailable",
        }
    target = _binary_array(target_amodal, "target_amodal")
    if target.shape != envelope.shape:
        raise ValueError("Target amodal and completion envelope must have the same shape")

    candidate_rgb = np.asarray(image.convert("RGB"), dtype=np.uint8)
    original_rgb = np.asarray(original.convert("RGB"), dtype=np.uint8)
    if candidate_rgb.shape[:2] != envelope.shape:
        raise ValueError("Candidate and completion envelope must have the same shape")
    if original_rgb.shape != candidate_rgb.shape:
        raise ValueError("Original and candidate RGB rasters must have the same shape")

    outside_target = envelope & ~target
    minimum_pixels = max(256, int(round(0.001 * outside_target.size)))
    outside_pixels = int(outside_target.sum())
    if outside_pixels < minimum_pixels:
        return {
            "available": False,
            "reason": "outside_target_clutter_unavailable",
            "outside_target_pixels": outside_pixels,
            "minimum_pixels": minimum_pixels,
        }

    original_gray = cv2.cvtColor(original_rgb, cv2.COLOR_RGB2GRAY)
    candidate_gray = cv2.cvtColor(candidate_rgb, cv2.COLOR_RGB2GRAY)
    median_gray = float(np.median(original_gray))
    canny_low = max(20, int(0.5 * median_gray))
    canny_high = max(canny_low + 1, min(220, int(1.2 * median_gray)))
    edges = cv2.Canny(
        candidate_gray,
        canny_low,
        canny_high,
        L2gradient=True,
    ) > 0
    outside_edges = edges & outside_target
    edge_density = float(outside_edges.sum()) / outside_pixels

    height, width = outside_target.shape
    minimum_dimension = min(height, width)
    minimum_line_length = max(12, int(round(0.015 * minimum_dimension)))
    maximum_line_gap = max(4, int(round(0.005 * minimum_dimension)))
    lines = cv2.HoughLinesP(
        outside_edges.astype(np.uint8) * 255,
        1,
        np.pi / 180.0,
        threshold=minimum_line_length,
        minLineLength=minimum_line_length,
        maxLineGap=maximum_line_gap,
    )
    total_line_length = 0.0
    line_count = 0
    if lines is not None:
        line_count = int(len(lines))
        for line in lines[:, 0, :]:
            x1, y1, x2, y2 = (int(value) for value in line)
            total_line_length += math.hypot(x2 - x1, y2 - y1)
    line_density_per_1000 = 1000.0 * total_line_length / outside_pixels
    return {
        "available": True,
        "outside_target_pixels": outside_pixels,
        "minimum_pixels": minimum_pixels,
        "canny_low": canny_low,
        "canny_high": canny_high,
        "edge_density": round(edge_density, 8),
        "hough_line_count": line_count,
        "hough_minimum_line_length": minimum_line_length,
        "hough_maximum_line_gap": maximum_line_gap,
        "line_density_per_1000_pixels": round(line_density_per_1000, 8),
    }


def _average_tie_percentile_ranks(values: list[float]) -> list[float]:
    if len(values) <= 1:
        return [0.0] * len(values)
    denominator = len(values) - 1
    ranks: list[float] = []
    for value in values:
        lower = sum(other < value for other in values)
        equal_other = sum(other == value for other in values) - 1
        ranks.append((lower + 0.5 * equal_other) / denominator)
    return ranks


def apply_selection_clutter_penalty(
    candidates: list[dict[str, Any]],
    *,
    maximum_penalty: float = 0.15,
) -> bool:
    """Add a cohort-relative clutter term used only to rank candidates."""

    if maximum_penalty < 0:
        raise ValueError("maximum_penalty must be non-negative")
    if not candidates:
        return False
    metrics = [row["quality"].get("clutter_metrics", {}) for row in candidates]
    if not all(metric.get("available") is True for metric in metrics):
        for row in candidates:
            quality = row["quality"]
            base = float(quality.get("gate_score", quality.get("score", 0.0)))
            quality["selection_score"] = round(base, 6)
            quality["clutter_selection_penalty"] = None
            quality["clutter_selection_note"] = (
                "unavailable_fallback_to_absolute_quality_score"
            )
        return False

    edge_ranks = _average_tie_percentile_ranks(
        [float(metric["edge_density"]) for metric in metrics]
    )
    line_ranks = _average_tie_percentile_ranks(
        [float(metric["line_density_per_1000_pixels"]) for metric in metrics]
    )
    for row, edge_rank, line_rank in zip(candidates, edge_ranks, line_ranks):
        quality = row["quality"]
        base = float(quality.get("gate_score", quality.get("score", 0.0)))
        clutter_rank = 0.5 * (edge_rank + line_rank)
        penalty = maximum_penalty * clutter_rank
        quality["clutter_edge_percentile_rank"] = round(edge_rank, 6)
        quality["clutter_line_percentile_rank"] = round(line_rank, 6)
        quality["clutter_rank"] = round(clutter_rank, 6)
        quality["clutter_selection_penalty"] = round(penalty, 6)
        quality["selection_score"] = round(base - penalty, 6)
        quality["clutter_selection_note"] = (
            "cohort_relative_ranking_only_not_an_acceptance_or_passability_gate"
        )
    return True


def select_candidate(candidates: list[dict[str, Any]]) -> tuple[dict[str, Any], str]:
    """Select by gate score and withhold publication when every row is risky."""

    if not candidates:
        raise ValueError("At least one completion candidate is required")

    def rank_key(row: dict[str, Any]) -> tuple[float, float, float, int]:
        quality = row["quality"]
        return (
            float(
                quality.get(
                    "selection_score",
                    quality.get("gate_score", quality.get("score", 0.0)),
                )
            ),
            float(quality.get("gate_score", quality.get("score", 0.0))),
            float(quality.get("score", 0.0)),
            -int(row.get("index", 0)),
        )

    selected = max(candidates, key=rank_key)
    all_risky = all(bool(row["quality"].get("quality_flags", [])) for row in candidates)
    status = "withheld_needs_review" if all_risky else "candidate_selected_for_review"
    return selected, status