app.py

import io
import time
import numpy as np
import cv2
import torch
from transformers import DPTFeatureExtractor, DPTForDepthEstimation
from fastapi import FastAPI, File, UploadFile, Response
from fastapi.responses import FileResponse
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-Hybrid để tăng tốc
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-swinv2-tiny-256")
model = DPTForDepthEstimation.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)

    # 🟢 Chuẩn bị ảnh cho mô hình
    inputs = feature_extractor(images=flipped_image, return_tensors="pt").to(device)

    # 🟢 Dự đoán Depth Map với DPT-Hybrid
    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 * 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"⏳ DPT xử lý trong {end_time - start_time:.4f} giây")

    return {"message": "Depth Map processed successfully. View at /depth_map/"}

@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")

# 🟢 Chạy server FastAPI trên Hugging Face Spaces
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)