import io
import time
import numpy as np
import cv2
import torch
import gradio as gr
from transformers import DPTFeatureExtractor, DPTForDepthEstimation
from fastapi import FastAPI, File, UploadFile
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
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-swinv2-tiny-256")
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-swinv2-tiny-256").to(device)
model.eval()

def process_depth_map(image_pil):
    """Xử lý Depth Map và hiển thị trên Hugging Face"""
    start_time = time.time()

    # 🔵 Resize ảnh để tăng tốc độ xử lý
    image = image_pil.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=image_np, return_tensors="pt").to(device)

    # 🟢 Dự đoán Depth Map với DPT-Swinv2 Tiny
    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)

    end_time = time.time()
    print(f"⏳ DPT xử lý trong {end_time - start_time:.4f} giây")

    return depth_pil  # Trả về ảnh để hiển thị trên Gradio

# 🟢 Tạo UI trên Hugging Face Spaces với Gradio
gr.Interface(
    fn=process_depth_map,
    inputs=gr.Image(type="pil"),  # 🔥 Sửa lỗi: dùng "pil" thay vì "bytes"
    outputs=gr.Image(type="pil"),
    title="🔍 Depth Map Estimation",
    description="Tải ảnh lên để xem Depth Map",
).launch(share=True)  # `share=True` để tạo link truy cập từ bên ngoài

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)