rezzzq

Add custom inference handler for DA3METRIC-LARGE depth estimation

e561725 10 months ago

3.37 kB

"""
Custom handler for Hugging Face Inference Endpoints.
Serves the Depth Anything V3 Metric Large model for depth estimation.
"""

import base64
import io
from typing import Any

import numpy as np
import torch
from PIL import Image


class EndpointHandler:
    def __init__(self, path: str = ""):
        """
        Initialize the depth estimation model.
        Args:
            path: Path to the model directory (provided by HF Inference Endpoints)
        """
        from depth_anything_3.api import DepthAnything3

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = DepthAnything3.from_pretrained("depth-anything/da3metric-large")
        self.model = self.model.to(device=self.device)

    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
        """
        Process incoming requests for depth estimation.
        Args:
            data: Request payload with 'inputs' containing base64 image(s)
        Returns:
            Dictionary with depth map, confidence, intrinsics, extrinsics
        """
        inputs = data.get("inputs")

        # Handle base64 encoded image input
        if isinstance(inputs, str):
            # Single base64 image
            image_data = base64.b64decode(inputs)
            image = Image.open(io.BytesIO(image_data)).convert("RGB")
            images = [image]
        elif isinstance(inputs, dict) and "image" in inputs:
            # Dict with image key
            image_data = base64.b64decode(inputs["image"])
            image = Image.open(io.BytesIO(image_data)).convert("RGB")
            images = [image]
        elif isinstance(inputs, list):
            # List of base64 images
            images = []
            for img_b64 in inputs:
                image_data = base64.b64decode(img_b64)
                image = Image.open(io.BytesIO(image_data)).convert("RGB")
                images.append(image)
        else:
            return {"error": "Invalid input format. Expected base64 encoded image(s)."}

        # Run inference
        with torch.inference_mode():
            prediction = self.model.inference(images)

        # Extract results
        depth = prediction.depth.cpu().numpy()  # [N, H, W]
        conf = prediction.conf.cpu().numpy()    # [N, H, W]
        intrinsics = prediction.intrinsics.cpu().numpy()  # [N, 3, 3]
        extrinsics = prediction.extrinsics.cpu().numpy()  # [N, 3, 4]

        # Return base64-encoded numpy arrays
        response = {
            "depth": self._encode_array(depth),
            "confidence": self._encode_array(conf),
            "intrinsics": self._encode_array(intrinsics),
            "extrinsics": self._encode_array(extrinsics),
            "shape": {
                "depth": list(depth.shape),
                "confidence": list(conf.shape),
                "intrinsics": list(intrinsics.shape),
                "extrinsics": list(extrinsics.shape),
            },
            "depth_range": {
                "min": float(depth.min()),
                "max": float(depth.max()),
            },
        }

        return response

    def _encode_array(self, arr: np.ndarray) -> str:
        """Encode numpy array as base64 string."""
        buffer = io.BytesIO()
        np.save(buffer, arr.astype(np.float32))
        return base64.b64encode(buffer.getvalue()).decode("utf-8")

	"""
	Custom handler for Hugging Face Inference Endpoints.
	Serves the Depth Anything V3 Metric Large model for depth estimation.
	"""

	import base64
	import io
	from typing import Any

	import numpy as np
	import torch
	from PIL import Image


	class EndpointHandler:
	def __init__(self, path: str = ""):
	"""
	Initialize the depth estimation model.

	Args:
	path: Path to the model directory (provided by HF Inference Endpoints)
	"""
	from depth_anything_3.api import DepthAnything3

	self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	self.model = DepthAnything3.from_pretrained("depth-anything/da3metric-large")
	self.model = self.model.to(device=self.device)

	def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
	"""
	Process incoming requests for depth estimation.

	Args:
	data: Request payload with 'inputs' containing base64 image(s)

	Returns:
	Dictionary with depth map, confidence, intrinsics, extrinsics
	"""
	inputs = data.get("inputs")

	# Handle base64 encoded image input
	if isinstance(inputs, str):
	# Single base64 image
	image_data = base64.b64decode(inputs)
	image = Image.open(io.BytesIO(image_data)).convert("RGB")
	images = [image]
	elif isinstance(inputs, dict) and "image" in inputs:
	# Dict with image key
	image_data = base64.b64decode(inputs["image"])
	image = Image.open(io.BytesIO(image_data)).convert("RGB")
	images = [image]
	elif isinstance(inputs, list):
	# List of base64 images
	images = []
	for img_b64 in inputs:
	image_data = base64.b64decode(img_b64)
	image = Image.open(io.BytesIO(image_data)).convert("RGB")
	images.append(image)
	else:
	return {"error": "Invalid input format. Expected base64 encoded image(s)."}

	# Run inference
	with torch.inference_mode():
	prediction = self.model.inference(images)

	# Extract results
	depth = prediction.depth.cpu().numpy() # [N, H, W]
	conf = prediction.conf.cpu().numpy() # [N, H, W]
	intrinsics = prediction.intrinsics.cpu().numpy() # [N, 3, 3]
	extrinsics = prediction.extrinsics.cpu().numpy() # [N, 3, 4]

	# Return base64-encoded numpy arrays
	response = {
	"depth": self._encode_array(depth),
	"confidence": self._encode_array(conf),
	"intrinsics": self._encode_array(intrinsics),
	"extrinsics": self._encode_array(extrinsics),
	"shape": {
	"depth": list(depth.shape),
	"confidence": list(conf.shape),
	"intrinsics": list(intrinsics.shape),
	"extrinsics": list(extrinsics.shape),
	},
	"depth_range": {
	"min": float(depth.min()),
	"max": float(depth.max()),
	},
	}

	return response

	def _encode_array(self, arr: np.ndarray) -> str:
	"""Encode numpy array as base64 string."""
	buffer = io.BytesIO()
	np.save(buffer, arr.astype(np.float32))
	return base64.b64encode(buffer.getvalue()).decode("utf-8")