import io
import time
import numpy as np
import cv2
import torch
from transformers import AutoImageProcessor, AutoModelForDepthEstimation
from fastapi import FastAPI, File, UploadFile, Response
from fastapi.responses import FileResponse, HTMLResponse
from PIL import Image
import uvicorn
app = FastAPI()
# 🟢 Chọn thiết bị xử lý (GPU nếu có)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 🟢 Tải model DPT-Swinv2 Tiny để tăng tốc
processor = AutoImageProcessor.from_pretrained("Intel/dpt-swinv2-tiny-256")
model = AutoModelForDepthEstimation.from_pretrained("Intel/dpt-swinv2-tiny-256").to(device)
model.eval()
@app.post("/analyze_path/")
async def analyze_path(file: UploadFile = File(...)):
"""Xử lý ảnh Depth Map và lưu ảnh để hiển thị"""
start_time = time.time()
# 🟢 Đọc file ảnh từ ESP32
image_bytes = await file.read()
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
# 🔵 Resize ảnh để xử lý nhanh hơn
image = image.resize((256, 256))
image_np = np.array(image)
# flipped_image = cv2.flip(image_np, -1) # Lật ảnh
# 🟢 Chuẩn bị ảnh cho mô hình
inputs = processor(images=image_np, return_tensors="pt").to(device)
# 🟢 Dự đoán Depth Map với model mới
with torch.no_grad():
outputs = model(**inputs)
# 🟢 Xử lý ảnh sau khi dự đoán
predicted_depth = outputs.predicted_depth.squeeze().cpu().numpy()
# depth_map = (predicted_depth - predicted_depth.min()) / (predicted_depth.max() - predicted_depth.min())
depth_map = (predicted_depth * 255 / predicted_depth.max()).astype("uint8")
# 🔵 Chuyển depth_map thành ảnh có thể hiển thị
depth_colored = cv2.applyColorMap(depth_map, cv2.COLORMAP_INFERNO)
depth_pil = Image.fromarray(depth_colored)
# 🟢 Lưu ảnh Depth Map để hiển thị trên trình duyệt
depth_pil.save("depth_map.png")
end_time = time.time()
print(f"⏳ Model xử lý trong {end_time - start_time:.4f} giây")
return {"message": "Depth Map processed successfully. View at /view/"}
@app.get("/depth_map/")
async def get_depth_map():
"""Trả về ảnh Depth Map để hiển thị trên trình duyệt"""
return FileResponse("depth_map.png", media_type="image/png")
@app.get("/view/", response_class=HTMLResponse)
async def view_depth_map():
"""Trả về trang HTML để hiển thị ảnh Depth Map với auto-refresh"""
html_content = """
Depth Map Viewer
Depth Map (Tự động cập nhật mỗi 5 giây)
"""
return HTMLResponse(content=html_content)
# 🟢 Chạy server FastAPI trên Hugging Face Spaces
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)