File size: 2,208 Bytes
f0e3abd a026586 f0e3abd a026586 f0e3abd a026586 f0e3abd |
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 |
import base64
import io
import json
from typing import Dict, Any
from PIL import Image
from transformers import pipeline
class EndpointHandler:
"""
Custom handler for the ZoeDepth model, fully compliant with the latest
Hugging Face Inference Endpoints documentation.
The final result is serialized into a single JSON string.
"""
def __init__(self, path=""):
# Initialize the pipeline for depth-estimation
self.pipe = pipeline(task="depth-estimation", model=path)
print("Depth estimation pipeline initialized successfully.")
def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
"""
This method is called for every API request.
Args:
data (Dict): The input data dictionary. Can be PIL Image or bytes.
Returns:
Dict[str, str]: A dictionary with a single key "generated_text",
containing a JSON string of the results.
"""
# Get image from the request
inputs = data.pop("inputs", data)
# Handle both PIL Image objects (from image content-type) and bytes (from JSON)
if isinstance(inputs, Image.Image):
image = inputs
else:
image = Image.open(io.BytesIO(inputs))
# Pass the image to the pipeline
prediction = self.pipe(image)
# Extract raw depth data and visual map
raw_depth_tensor = prediction["predicted_depth"]
raw_depth_data = raw_depth_tensor.cpu().tolist()
visual_map_image = prediction["depth"]
buffered = io.BytesIO()
visual_map_image.save(buffered, format="PNG")
visual_map_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
# Create a dictionary to hold all results
results = {
"raw_depth_data": raw_depth_data,
"visual_depth_map": f"data:image/png;base64,{visual_map_base64}"
}
# Serialize the entire results dictionary into a JSON string
json_output_string = json.dumps(results)
# Return the final dictionary in the required format
return {"generated_text": json_output_string}
|