| import base64 |
| import io |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from PIL import Image, ImageFilter |
| from transformers import AutoImageProcessor, AutoModelForDepthEstimation |
|
|
|
|
| MODEL_ID = "depth-anything/Depth-Anything-V2-Small-hf" |
| MAX_PIXELS = 12_000_000 |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path: str = "") -> None: |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model_path = path or MODEL_ID |
| self.processor = AutoImageProcessor.from_pretrained(model_path) |
| self.model = AutoModelForDepthEstimation.from_pretrained(model_path).to(self.device) |
| self.model.eval() |
|
|
| def __call__(self, data: dict[str, Any]) -> dict[str, Any]: |
| encoded = data.get("inputs") |
| if not isinstance(encoded, str) or not encoded: |
| raise ValueError("inputs must be a base64-encoded image") |
| try: |
| image_bytes = base64.b64decode(encoded, validate=True) |
| image = Image.open(io.BytesIO(image_bytes)).convert("RGB") |
| except Exception as error: |
| raise ValueError("inputs must contain a valid base64-encoded image") from error |
|
|
| width, height = image.size |
| if width < 1 or height < 1 or width * height > MAX_PIXELS: |
| raise ValueError("image dimensions exceed the endpoint limit") |
|
|
| inputs = self.processor(images=image, return_tensors="pt") |
| inputs = {key: value.to(self.device) for key, value in inputs.items()} |
| with torch.inference_mode(): |
| predicted = self.model(**inputs).predicted_depth |
| predicted = F.interpolate( |
| predicted.unsqueeze(1), |
| size=(height, width), |
| mode="bicubic", |
| align_corners=False, |
| )[0, 0] |
|
|
| values = predicted.float().cpu().numpy() |
| values = np.nan_to_num(values, nan=0.0, posinf=0.0, neginf=0.0) |
| low = float(values.min()) |
| high = float(values.max()) |
| if high > low: |
| normalized = np.clip((values - low) / (high - low), 0.0, 1.0) |
| else: |
| normalized = np.zeros_like(values) |
| depth = Image.fromarray(np.round(normalized * 255).astype(np.uint8), mode="L") |
| depth = depth.filter(ImageFilter.GaussianBlur(radius=0.8)) |
|
|
| output = io.BytesIO() |
| depth.save(output, format="PNG", optimize=True) |
| return { |
| "depth_map_png": base64.b64encode(output.getvalue()).decode("ascii"), |
| "width": width, |
| "height": height, |
| } |
|
|
