"""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", ]