handler.py

import base64
import datetime as dt
import io
import os
import time
import uuid
from typing import Any, Dict, List

from PIL import Image


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.path = path
        self.provider = "hf_endpoint"
        self.hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
        self.sam_model_id = os.getenv("HF_SAM_MODEL_ID", "facebook/sam3")
        self.depth_model_id = os.getenv("HF_DEPTH_MODEL_ID", "depth-anything/Depth-Anything-V2-Small-hf")
        self._device = None
        self._torch = None
        self._np = None
        self._sam_model = None
        self._sam_processor = None
        self._depth_model = None
        self._depth_processor = None

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        batch_id = str(uuid.uuid4())
        created_at = dt.datetime.now(dt.timezone.utc).isoformat()
        batch_start = time.perf_counter()

        inputs = data.get("inputs")
        task = data.get("task")
        parameters = data.get("parameters", {}) or {}

        if not isinstance(inputs, list) or not inputs:
            return {
                "batch_id": batch_id,
                "status": "failed_validation",
                "provider": self.provider,
                "task": task,
                "error": {"message": "Payload must include a non-empty 'inputs' list."},
                "submitted_count": 0,
                "succeeded_count": 0,
                "failed_count": 0,
                "created_at": created_at,
                "results": [],
                "timing_ms": self._timing_block(total_ms=self._elapsed_ms(batch_start))
            }

        if task not in {"object_segmentation", "depth_estimation"}:
            return {
                "batch_id": batch_id,
                "status": "failed_validation",
                "provider": self.provider,
                "task": task,
                "error": {"message": "Unsupported task. Use 'object_segmentation' or 'depth_estimation'."},
                "submitted_count": len(inputs),
                "succeeded_count": 0,
                "failed_count": len(inputs),
                "created_at": created_at,
                "results": [],
                "timing_ms": self._timing_block(total_ms=self._elapsed_ms(batch_start))
            }

        results: List[Dict[str, Any]] = []
        succeeded_count = 0
        failed_count = 0
        aggregate_decode = 0.0
        aggregate_preprocess = 0.0
        aggregate_inference = 0.0
        aggregate_postprocess = 0.0

        for index, item in enumerate(inputs):
            image_id = self._image_id(item, index)
            item_start = time.perf_counter()

            try:
                decode_start = time.perf_counter()
                image = self._decode_image(item)
                decode_ms = self._elapsed_ms(decode_start)

                preprocess_start = time.perf_counter()
                prepared = self._prepare_image(image)
                preprocess_ms = self._elapsed_ms(preprocess_start)

                inference_start = time.perf_counter()
                model_output = self._infer_single_image(prepared, task, parameters)
                inference_ms = self._elapsed_ms(inference_start)

                postprocess_start = time.perf_counter()
                normalized = self._normalize_output(image, task, model_output)
                postprocess_ms = self._elapsed_ms(postprocess_start)

                item_timing = self._timing_block(
                    decode_ms,
                    preprocess_ms,
                    inference_ms,
                    postprocess_ms,
                    self._elapsed_ms(item_start),
                )

                result = {
                    "image_id": image_id,
                    "index": index,
                    "status": "succeeded",
                    "retryable": False,
                    "timing_ms": item_timing,
                }
                result.update(normalized)
                results.append(result)
                succeeded_count += 1

                aggregate_decode += decode_ms
                aggregate_preprocess += preprocess_ms
                aggregate_inference += inference_ms
                aggregate_postprocess += postprocess_ms

            except Exception as error:
                status, retryable, message = self._classify_error(error)
                results.append(
                    {
                        "image_id": image_id,
                        "index": index,
                        "status": status,
                        "retryable": retryable,
                        "error": {"message": message, "type": type(error).__name__},
                        "timing_ms": self._timing_block(total_ms=self._elapsed_ms(item_start)),
                    }
                )
                failed_count += 1

        batch_status = "succeeded" if failed_count == 0 else ("failed" if succeeded_count == 0 else "partially_succeeded")

        return {
            "batch_id": batch_id,
            "status": batch_status,
            "provider": self.provider,
            "task": task,
            "submitted_count": len(inputs),
            "succeeded_count": succeeded_count,
            "failed_count": failed_count,
            "created_at": created_at,
            "results": results,
            "timing_ms": self._timing_block(
                aggregate_decode,
                aggregate_preprocess,
                aggregate_inference,
                aggregate_postprocess,
                self._elapsed_ms(batch_start),
            ),
        }

    def _decode_image(self, item: Dict[str, Any]) -> Image.Image:
        if not isinstance(item, dict):
            raise ValueError("Each input item must be an object.")

        image_base64 = item.get("image_base64")
        if not image_base64:
            raise ValueError("Each input item must include 'image_base64'.")

        try:
            payload = base64.b64decode(image_base64)
            return Image.open(io.BytesIO(payload)).convert("RGB")
        except Exception as error:
            raise ValueError("Invalid base64 image payload.") from error

    def _prepare_image(self, image: Image.Image) -> Image.Image:
        return image

    def _infer_single_image(self, image: Image.Image, task: str, parameters: Dict[str, Any]) -> Dict[str, Any]:
        if task == "object_segmentation":
            return self._infer_sam(image, parameters)

        if task == "depth_estimation":
            return self._infer_depth(image)

        raise ValueError(f"Unsupported task: {task}")

    def _infer_sam(self, image: Image.Image, parameters: Dict[str, Any]) -> Dict[str, Any]:
        self._ensure_sam_loaded()

        text_query = (parameters.get("text_query") or "").strip()
        if not text_query:
            raise ValueError("'text_query' is required for object_segmentation.")

        confidence_threshold = float(parameters.get("confidence_threshold", 0.5))
        width, height = image.size
        predictions: List[Dict[str, Any]] = []
        queries = [query.strip() for query in text_query.split(",") if query.strip()]

        for class_id, query in enumerate(queries):
            model_inputs = self._sam_processor(images=image, text=query, return_tensors="pt")
            model_inputs = {key: value.to(self._device) for key, value in model_inputs.items()}

            with self._torch.inference_mode():
                inference_output = self._sam_model(**model_inputs)

            processed_results = self._sam_processor.post_process_instance_segmentation(
                inference_output,
                threshold=confidence_threshold,
                mask_threshold=0.5,
                target_sizes=model_inputs.get("original_sizes").tolist(),
            )[0]

            raw_masks = processed_results["masks"].detach().cpu().numpy()
            raw_scores = processed_results["scores"].detach().cpu().numpy()

            for mask_index, mask_array in enumerate(raw_masks):
                predictions.append(
                    self._build_sam_prediction(
                        mask_array=mask_array,
                        score=float(raw_scores[mask_index]),
                        class_id=class_id,
                        class_name=query,
                    )
                )

        return {"image": {"width": width, "height": height}, "predictions": predictions}

    def _infer_depth(self, image: Image.Image) -> Dict[str, Any]:
        self._ensure_depth_loaded()

        width, height = image.size
        model_inputs = self._depth_processor(images=image, return_tensors="pt")
        model_inputs = {key: value.to(self._device) for key, value in model_inputs.items()}

        with self._torch.inference_mode():
            inference_output = self._depth_model(**model_inputs)

        depth_tensor = inference_output.predicted_depth
        resized_depth = self._torch.nn.functional.interpolate(
            depth_tensor.unsqueeze(1),
            size=(height, width),
            mode="bicubic",
            align_corners=False,
        ).squeeze()

        depth_cpu = resized_depth.detach().cpu()
        min_depth = float(depth_cpu.min().item())
        max_depth = float(depth_cpu.max().item())

        if max_depth > min_depth:
            normalized_depth = ((depth_cpu - min_depth) / (max_depth - min_depth) * 255.0).clamp(0, 255)
        else:
            normalized_depth = self._torch.zeros_like(depth_cpu)

        depth_array = normalized_depth.to(self._torch.uint8).numpy()
        depth_image = Image.fromarray(depth_array, mode="L")
        buffer = io.BytesIO()
        depth_image.save(buffer, format="PNG")

        return {
            "image": {"width": width, "height": height},
            "depth_map": {
                "encoding": "png_base64",
                "image_base64": base64.b64encode(buffer.getvalue()).decode("utf-8"),
                "min_depth": round(min_depth, 8),
                "max_depth": round(max_depth, 8),
            },
        }

    def _normalize_output(self, image: Image.Image, task: str, model_output: Dict[str, Any]) -> Dict[str, Any]:
        if task == "object_segmentation":
            return {"sam": model_output}
        if task == "depth_estimation":
            return {"depth": model_output}
        raise ValueError(f"Unsupported task: {task}")

    def _ensure_sam_loaded(self) -> None:
        if self._sam_model is not None and self._sam_processor is not None:
            return

        self._ensure_runtime_loaded()

        try:
            from transformers import Sam3Model, Sam3Processor
        except ImportError as error:
            raise RuntimeError("transformers with SAM3 support is required for object_segmentation.") from error

        self._sam_model = Sam3Model.from_pretrained(self.sam_model_id, token=self.hf_token).to(self._device)
        self._sam_processor = Sam3Processor.from_pretrained(self.sam_model_id, token=self.hf_token)

    def _ensure_depth_loaded(self) -> None:
        if self._depth_model is not None and self._depth_processor is not None:
            return

        self._ensure_runtime_loaded()

        try:
            from transformers import AutoImageProcessor, AutoModelForDepthEstimation
        except ImportError as error:
            raise RuntimeError("transformers with depth-estimation support is required for depth_estimation.") from error

        self._depth_processor = AutoImageProcessor.from_pretrained(self.depth_model_id, token=self.hf_token)
        self._depth_model = AutoModelForDepthEstimation.from_pretrained(self.depth_model_id, token=self.hf_token).to(self._device)

    def _ensure_runtime_loaded(self) -> None:
        if self._torch is not None and self._np is not None and self._device is not None:
            return

        try:
            import numpy as np
            import torch
        except ImportError as error:
            raise RuntimeError("numpy and torch are required for model inference.") from error

        self._np = np
        self._torch = torch
        self._device = "cuda" if torch.cuda.is_available() else "cpu"

    def _build_sam_prediction(
        self,
        mask_array: Any,
        score: float,
        class_id: int,
        class_name: str,
    ) -> Dict[str, Any]:
        rows = self._np.any(mask_array, axis=1)
        cols = self._np.any(mask_array, axis=0)

        if rows.any() and cols.any():
            y_indices = self._np.where(rows)[0]
            x_indices = self._np.where(cols)[0]
            y1, y2 = int(y_indices[0]), int(y_indices[-1])
            x1, x2 = int(x_indices[0]), int(x_indices[-1])
        else:
            x1 = y1 = x2 = y2 = 0

        bbox_width = x2 - x1
        bbox_height = y2 - y1
        center_x = x1 + bbox_width / 2
        center_y = y1 + bbox_height / 2

        mask_image = Image.fromarray((mask_array * 255).astype("uint8"), mode="L")
        buffer = io.BytesIO()
        mask_image.save(buffer, format="PNG")

        return {
            "width": bbox_width,
            "height": bbox_height,
            "x": round(center_x, 1),
            "y": round(center_y, 1),
            "confidence": round(score, 8),
            "class_id": class_id,
            "class": class_name,
            "detection_id": str(uuid.uuid4()),
            "parent_id": "image",
            "mask_base64": base64.b64encode(buffer.getvalue()).decode("utf-8"),
        }

    def _classify_error(self, error: Exception) -> tuple[str, bool, str]:
        message = str(error)
        lowered = message.lower()

        if isinstance(error, ValueError):
            return "failed_validation", False, message
        if "out of memory" in lowered or "cuda" in lowered:
            return "failed_model", True, "Resource pressure or GPU error"
        if "timed out" in lowered or "timeout" in lowered:
            return "failed_model", True, "Inference timed out"
        if isinstance(error, RuntimeError):
            return "failed_model", True, message
        return "failed_internal", True, message

    def _image_id(self, item: Dict[str, Any], index: int) -> str:
        if isinstance(item, dict) and item.get("image_id"):
            return str(item["image_id"])
        return f"image-{index}"

    def _timing_block(
        self,
        decode_ms: float = 0.0,
        preprocess_ms: float = 0.0,
        inference_ms: float = 0.0,
        postprocess_ms: float = 0.0,
        total_ms: float = 0.0,
    ) -> Dict[str, float]:
        return {
            "decode": round(decode_ms, 3),
            "preprocess": round(preprocess_ms, 3),
            "inference": round(inference_ms, 3),
            "postprocess": round(postprocess_ms, 3),
            "total": round(total_ms, 3),
        }

    def _elapsed_ms(self, start: float) -> float:
        return (time.perf_counter() - start) * 1000.0