File size: 2,578 Bytes
02a2bdf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | 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,
}
|