DepthFlowTest / handler.py
selectmixer's picture
headless
327c44c
# handler.py
import os
import io
import tempfile
import numpy as np
from PIL import Image
from DepthFlow import DepthScene
from DepthFlow.Motion import Presets
class EndpointHandler():
def __init__(self, path=""):
"""
Initialize the handler and load necessary resources.
This method is called once when the service starts.
"""
# Initialize DepthFlow once for efficiency
self.depthflow = DepthScene(backend='headless')
def preprocess(self, inputs):
"""
Preprocess the input data.
Args:
inputs (dict): The input payload containing the image data.
Returns:
str: Path to the preprocessed image file.
"""
if 'image' not in inputs:
raise ValueError("Missing 'image' in inputs")
image_bytes = inputs['image'].read()
image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
# Save image to a temporary file
temp_image_file = tempfile.NamedTemporaryFile(suffix='.jpg', delete=False)
image.save(temp_image_file.name)
return temp_image_file.name
def inference(self, image_path):
"""
Perform the main inference logic.
Args:
image_path (str): Path to the preprocessed image file.
Returns:
str: Path to the output video file.
"""
# Load the image into DepthFlow
self.depthflow.input(image=image_path)
# Set custom parameters (modify as needed)
self.depthflow.state.height = 1
self.depthflow.state.zoom = 1.1
self.depthflow.state.dolly = 1
self.depthflow.state.dof_enable = True
self.depthflow.state.dof_intensity = 1.2
self.depthflow.state.vignette_intensity = 40
# Apply the animation preset
self.depthflow.add_animation(Presets.Dolly())
# Generate the output video
temp_video_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
self.depthflow.update()
self.depthflow.main(output=temp_video_file.name)
return temp_video_file.name
def postprocess(self, video_path):
"""
Postprocess the output and prepare the response.
Args:
video_path (str): Path to the generated video file.
Returns:
dict: Response containing the video file.
"""
with open(video_path, 'rb') as f:
video_bytes = f.read()
# Clean up temporary files
os.remove(video_path)
return {
'video': video_bytes
}
def __call__(self, data):
"""
Handle the incoming request.
Args:
inputs (dict): The input payload.
Returns:
dict: The response payload.
"""
try:
# Preprocess
image_path = self.preprocess(data)
# Inference
video_path = self.inference(image_path)
# Postprocess
result = self.postprocess(video_path)
# Clean up image file
os.remove(image_path)
return result
except Exception as e:
return {'error': str(e)}