Spaces:
Sleeping
Sleeping
| """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.""" | |
| 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", | |
| ] | |