visual_prediction.py

"""Public-safe contract for task-aware computer-vision prediction previews.

The private backend renders PNG bytes.  This module validates their bounded
wire representation and materializes them for Gradio without importing any ML
backend or accepting backend filesystem paths.
"""

from __future__ import annotations

import base64
import binascii
import hashlib
import json
import math
import re
import struct
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Mapping, Sequence


VISUAL_PREDICTION_SCHEMA_VERSION = 1
MAX_VISUAL_DIMENSION = 1_600
MAX_VISUAL_ARTIFACT_BYTES = 8 * 1024 * 1024
MAX_VISUAL_ARTIFACTS = 3
MAX_VISUAL_RECORDS = 1_000
MAX_VISUAL_WARNINGS = 8
MAX_VISUAL_WARNING_CHARS = 320

_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
_TASKS = frozenset(
    {"object_detection", "instance_segmentation", "semantic_segmentation"}
)
_ROLES_BY_TASK = {
    "object_detection": frozenset({"annotated"}),
    "instance_segmentation": frozenset({"source", "overlay"}),
    "semantic_segmentation": frozenset({"source", "mask", "overlay"}),
}
_FILENAMES = {
    "source": "source.png",
    "annotated": "annotated.png",
    "mask": "predicted-mask.png",
    "overlay": "overlay.png",
}
_SENSITIVE_TEXT = re.compile(
    r"(?:hf_[A-Za-z0-9]{6,}|/(?:Users|home|tmp|var|opt)/|[A-Za-z]:[\\/])",
    re.IGNORECASE,
)


class VisualPredictionError(ValueError):
    """Raised when visual-result metadata crosses the public boundary unsafely."""


@dataclass(frozen=True)
class MaterializedVisualPrediction:
    """Validated request-scoped files ready for non-interactive Gradio images."""

    task_type: str = ""
    artifacts: Mapping[str, str] | None = None
    records: tuple[Mapping[str, Any], ...] = ()
    warnings: tuple[str, ...] = ()
    request_dir: str = ""

    def path(self, role: str) -> str | None:
        value = dict(self.artifacts or {}).get(role)
        return str(value) if value else None


def _safe_warning(value: Any) -> str:
    text = " ".join(str(value or "").split())[:MAX_VISUAL_WARNING_CHARS]
    if not text or _SENSITIVE_TEXT.search(text):
        return "Visual preview could not be generated; structured predictions remain available."
    return text


def _safe_source_name(value: Any) -> str:
    text = str(value or "prediction-image").strip()
    if not text:
        return "prediction-image"
    if "/" in text or "\\" in text or text in {".", ".."}:
        text = Path(text).name
    text = re.sub(r"[^A-Za-z0-9._ -]", "_", text)[:160]
    return text or "prediction-image"


def _finite_number(value: Any, field: str) -> float:
    try:
        result = float(value)
    except (TypeError, ValueError) as exc:
        raise VisualPredictionError(f"{field} must be numeric") from exc
    if not math.isfinite(result):
        raise VisualPredictionError(f"{field} must be finite")
    return result


def _integer(value: Any, field: str, *, minimum: int | None = None) -> int:
    try:
        result = int(value)
    except (TypeError, ValueError) as exc:
        raise VisualPredictionError(f"{field} must be an integer") from exc
    if minimum is not None and result < minimum:
        raise VisualPredictionError(f"{field} is below its supported minimum")
    return result


def _png_dimensions(data: bytes) -> tuple[int, int]:
    if len(data) < 24 or not data.startswith(_PNG_SIGNATURE):
        raise VisualPredictionError("visual artifact must be a PNG")
    if data[12:16] != b"IHDR":
        raise VisualPredictionError("visual artifact PNG is missing IHDR")
    width, height = struct.unpack(">II", data[16:24])
    if not width or not height:
        raise VisualPredictionError("visual artifact dimensions must be positive")
    return int(width), int(height)


def _artifact(value: Any, index: int) -> tuple[dict[str, Any], bytes]:
    if not isinstance(value, Mapping):
        raise VisualPredictionError("visual artifacts must be mappings")
    allowed = {
        "artifact_id",
        "role",
        "mime_type",
        "width",
        "height",
        "sha256",
        "data_base64",
    }
    if set(value) != allowed:
        raise VisualPredictionError("visual artifact contains unsupported fields")
    role = str(value.get("role") or "").strip().lower()
    if role not in _FILENAMES:
        raise VisualPredictionError("visual artifact role is unsupported")
    artifact_id = str(value.get("artifact_id") or "").strip()
    if not re.fullmatch(r"[a-z0-9][a-z0-9._-]{0,79}", artifact_id):
        raise VisualPredictionError("visual artifact ID is invalid")
    if value.get("mime_type") != "image/png":
        raise VisualPredictionError("visual artifacts must use image/png")
    encoded = value.get("data_base64")
    if not isinstance(encoded, str) or not encoded:
        raise VisualPredictionError("visual artifact data must be base64 text")
    if len(encoded) > ((MAX_VISUAL_ARTIFACT_BYTES + 2) // 3) * 4 + 4:
        raise VisualPredictionError("visual artifact exceeds the encoded size limit")
    try:
        data = base64.b64decode(encoded, validate=True)
    except (binascii.Error, ValueError) as exc:
        raise VisualPredictionError("visual artifact base64 is invalid") from exc
    if len(data) > MAX_VISUAL_ARTIFACT_BYTES:
        raise VisualPredictionError("visual artifact exceeds the decoded size limit")
    width, height = _png_dimensions(data)
    declared_width = _integer(value.get("width"), "visual artifact width", minimum=1)
    declared_height = _integer(value.get("height"), "visual artifact height", minimum=1)
    if (width, height) != (declared_width, declared_height):
        raise VisualPredictionError("visual artifact dimensions do not match the PNG")
    if max(width, height) > MAX_VISUAL_DIMENSION:
        raise VisualPredictionError("visual artifact dimensions exceed the preview limit")
    digest = hashlib.sha256(data).hexdigest()
    if not isinstance(value.get("sha256"), str) or value["sha256"] != digest:
        raise VisualPredictionError("visual artifact hash does not match its data")
    return (
        {
            "artifact_id": artifact_id,
            "role": role,
            "mime_type": "image/png",
            "width": width,
            "height": height,
            "sha256": digest,
            "data_base64": encoded,
        },
        data,
    )


def _point_array(value: Any, field: str, *, normalized: bool) -> list[list[float]]:
    if not isinstance(value, (list, tuple)):
        raise VisualPredictionError(f"{field} must be an array")
    points: list[list[float]] = []
    if value and all(isinstance(item, (int, float)) for item in value):
        if len(value) % 2:
            raise VisualPredictionError(f"{field} must contain x/y pairs")
        value = [value[index : index + 2] for index in range(0, len(value), 2)]
    for index, point in enumerate(value):
        if not isinstance(point, (list, tuple)) or len(point) != 2:
            raise VisualPredictionError(f"{field}[{index}] must be an x/y pair")
        x = _finite_number(point[0], f"{field}[{index}].x")
        y = _finite_number(point[1], f"{field}[{index}].y")
        if normalized and not (0.0 <= x <= 1.0 and 0.0 <= y <= 1.0):
            raise VisualPredictionError(f"{field} normalized coordinates must be in [0, 1]")
        points.append([x, y])
    if points and len(points) < 3:
        raise VisualPredictionError(f"{field} must contain at least three points")
    return points


def _record(value: Any, task_type: str, index: int) -> dict[str, Any]:
    if not isinstance(value, Mapping):
        raise VisualPredictionError("visual prediction records must be mappings")
    common = {
        "class_id",
        "class_name",
        "confidence",
        "x1",
        "y1",
        "x2",
        "y2",
    }
    allowed = set(common)
    if task_type == "instance_segmentation":
        allowed.update({"polygon", "polygon_normalized"})
    elif task_type == "semantic_segmentation":
        allowed = {
            "class_id",
            "class_name",
            "mask_value",
            "pixel_count",
            "pixel_percentage",
        }
    if set(value) - allowed:
        raise VisualPredictionError("visual prediction record contains unsupported fields")

    class_id = _integer(value.get("class_id"), f"records[{index}].class_id")
    class_name = " ".join(str(value.get("class_name") or class_id).split())[:160]
    record: dict[str, Any] = {"class_id": class_id, "class_name": class_name}
    if task_type == "semantic_segmentation":
        mask_value = _integer(value.get("mask_value"), f"records[{index}].mask_value")
        if mask_value not in {0, 255}:
            raise VisualPredictionError("semantic mask values must be 0 or 255")
        pixel_count = _integer(
            value.get("pixel_count"), f"records[{index}].pixel_count", minimum=0
        )
        percentage = _finite_number(
            value.get("pixel_percentage"), f"records[{index}].pixel_percentage"
        )
        if not 0.0 <= percentage <= 100.0:
            raise VisualPredictionError("semantic pixel percentage must be in [0, 100]")
        record.update(
            mask_value=mask_value,
            pixel_count=pixel_count,
            pixel_percentage=percentage,
        )
        return record

    confidence = value.get("confidence")
    if confidence is not None:
        confidence = _finite_number(confidence, f"records[{index}].confidence")
    coords = {
        key: _finite_number(value.get(key), f"records[{index}].{key}")
        for key in ("x1", "y1", "x2", "y2")
    }
    if coords["x2"] < coords["x1"] or coords["y2"] < coords["y1"]:
        raise VisualPredictionError("visual prediction bounding box is inverted")
    record.update(confidence=confidence, **coords)
    if task_type == "instance_segmentation":
        record["polygon"] = _point_array(
            value.get("polygon") or [], f"records[{index}].polygon", normalized=False
        )
        record["polygon_normalized"] = _point_array(
            value.get("polygon_normalized") or [],
            f"records[{index}].polygon_normalized",
            normalized=True,
        )
        if not record["polygon"] and not record["polygon_normalized"]:
            raise VisualPredictionError(
                "instance segmentation records require a polygon mask"
            )
    return record


def normalize_visual_prediction(value: Mapping[str, Any] | None) -> dict[str, Any]:
    """Validate and JSON-copy one optional visual prediction sidecar."""

    if value in (None, {}):
        return {}
    if not isinstance(value, Mapping):
        raise VisualPredictionError("visual prediction must be a mapping")
    allowed = {
        "schema_version",
        "task_type",
        "source_name",
        "artifacts",
        "records",
        "warnings",
    }
    if set(value) != allowed:
        raise VisualPredictionError("visual prediction contains unsupported fields")
    if value.get("schema_version") != VISUAL_PREDICTION_SCHEMA_VERSION:
        raise VisualPredictionError("visual prediction schema version is unsupported")
    task_type = str(value.get("task_type") or "").strip().lower()
    if task_type not in _TASKS:
        raise VisualPredictionError("visual prediction task is unsupported")
    raw_artifacts = value.get("artifacts")
    if not isinstance(raw_artifacts, (list, tuple)) or len(raw_artifacts) > MAX_VISUAL_ARTIFACTS:
        raise VisualPredictionError("visual artifacts exceed the supported count")
    artifacts: list[dict[str, Any]] = []
    roles: set[str] = set()
    for index, raw in enumerate(raw_artifacts):
        artifact, _ = _artifact(raw, index)
        if artifact["role"] in roles:
            raise VisualPredictionError("visual artifact roles must be unique")
        roles.add(artifact["role"])
        artifacts.append(artifact)
    raw_warnings = value.get("warnings")
    if not isinstance(raw_warnings, (list, tuple)) or len(raw_warnings) > MAX_VISUAL_WARNINGS:
        raise VisualPredictionError("visual warnings exceed the supported count")
    warnings = [_safe_warning(item) for item in raw_warnings]
    if artifacts and not _ROLES_BY_TASK[task_type].issubset(roles):
        raise VisualPredictionError("visual prediction is missing a required image role")
    if not artifacts and not warnings:
        raise VisualPredictionError("visual prediction requires images or a warning")

    raw_records = value.get("records")
    if not isinstance(raw_records, (list, tuple)) or len(raw_records) > MAX_VISUAL_RECORDS:
        raise VisualPredictionError("visual prediction records exceed the supported count")
    records = [_record(item, task_type, index) for index, item in enumerate(raw_records)]
    result = {
        "schema_version": VISUAL_PREDICTION_SCHEMA_VERSION,
        "task_type": task_type,
        "source_name": _safe_source_name(value.get("source_name")),
        "artifacts": artifacts,
        "records": records,
        "warnings": warnings,
    }
    return json.loads(json.dumps(result, allow_nan=False))


def attach_visual_prediction(frame: Any, value: Mapping[str, Any] | None) -> Any:
    """Attach a validated sidecar without changing the DataFrame's rows/columns."""

    normalized = normalize_visual_prediction(value)
    if hasattr(frame, "attrs"):
        attrs = dict(getattr(frame, "attrs", {}) or {})
        if normalized:
            attrs["visual_prediction"] = normalized
        else:
            attrs.pop("visual_prediction", None)
        frame.attrs = attrs
    return frame


def visual_prediction_from_frame(frame: Any) -> dict[str, Any]:
    attrs = getattr(frame, "attrs", {}) or {}
    if not isinstance(attrs, Mapping):
        return {}
    return normalize_visual_prediction(attrs.get("visual_prediction"))


def materialize_visual_prediction(
    value: Mapping[str, Any] | None,
    root: str | Path,
    *,
    request_id: str | None = None,
) -> MaterializedVisualPrediction:
    """Decode validated PNGs into one unique request-scoped directory."""

    normalized = normalize_visual_prediction(value)
    if not normalized:
        return MaterializedVisualPrediction(artifacts={})
    base = Path(root).expanduser().resolve()
    base.mkdir(parents=True, exist_ok=True)
    prefix = re.sub(r"[^a-zA-Z0-9_-]", "", str(request_id or ""))[:48]
    request_dir = Path(tempfile.mkdtemp(prefix=f"visual-{prefix}-", dir=base)).resolve()
    try:
        request_dir.relative_to(base)
    except ValueError as exc:  # pragma: no cover - tempfile containment guard
        raise VisualPredictionError("visual request directory escaped its root") from exc
    paths: dict[str, str] = {}
    try:
        for raw in normalized["artifacts"]:
            artifact, data = _artifact(raw, 0)
            output = (request_dir / _FILENAMES[artifact["role"]]).resolve()
            try:
                output.relative_to(request_dir)
            except ValueError as exc:  # pragma: no cover - fixed names guard
                raise VisualPredictionError("visual artifact path escaped its request") from exc
            output.write_bytes(data)
            paths[artifact["role"]] = str(output)
    except Exception:
        for child in request_dir.glob("*"):
            child.unlink(missing_ok=True)
        request_dir.rmdir()
        raise
    return MaterializedVisualPrediction(
        task_type=normalized["task_type"],
        artifacts=paths,
        records=tuple(normalized["records"]),
        warnings=tuple(normalized["warnings"]),
        request_dir=str(request_dir),
    )


__all__ = [
    "MAX_VISUAL_ARTIFACTS",
    "MAX_VISUAL_ARTIFACT_BYTES",
    "MAX_VISUAL_DIMENSION",
    "MaterializedVisualPrediction",
    "VISUAL_PREDICTION_SCHEMA_VERSION",
    "VisualPredictionError",
    "attach_visual_prediction",
    "materialize_visual_prediction",
    "normalize_visual_prediction",
    "visual_prediction_from_frame",
]