from fastapi import FastAPI
import gradio as gr
from PIL import Image
import numpy as np
import torch
from transformers import pipeline
import cv2

app = FastAPI()

# 딥러닝 모델 로드 (Depth Anything)
# 모델 로딩은 앱 시작 시 한 번만 하도록 글로벌 변수로 설정
print("Loading Depth Anything model...")
try:
    depth_estimator = pipeline(task="depth-estimation", model="LiangNX/depth-anything-hf")
    print("Depth Anything model loaded successfully.")
except Exception as e:
    print(f"Error loading Depth Anything model: {e}")
    depth_estimator = None # 모델 로드 실패 시 None으로 설정

def process_image_for_depth(image_path_or_pil_image):
    if depth_estimator is None:
        return None, "Error: Depth Anything model not loaded."

    # Gradio는 PIL Image 객체로 이미지를 전달합니다.
    if isinstance(image_path_or_pil_image, str):
        image = Image.open(image_path_or_pil_image).convert("RGB")
    else:
        image = image_path_or_pil_image.convert("RGB")

    try:
        # Depth Anything 모델 추론
        # result는 딕셔너리로, 'depth' (PIL Image)와 'depth_npy' (numpy array)를 포함
        result = depth_estimator(image)
        
        # 뎁스 맵 (PIL Image)
        depth_image_pil = result["depth"]
        
        # 뎁스 맵 (Numpy Array) - 시각화를 위해 정규화
        depth_np = result["depth_npy"]
        normalized_depth_np = (depth_np - depth_np.min()) / (depth_np.max() - depth_np.min()) * 255
        normalized_depth_np = normalized_depth_np.astype(np.uint8)
        
        # 흑백 이미지로 변환 (PIL Image)
        depth_grayscale_pil = Image.fromarray(normalized_depth_np)

        return depth_grayscale_pil, None
    except Exception as e:
        return None, f"Error processing image for depth: {e}"

# Gradio 인터페이스 정의
with gr.Blocks() as demo:
    gr.Markdown("# 🧑‍💻 얼굴 뎁스 맵 추출기")
    gr.Markdown("여러 장의 얼굴 사진을 업로드하면 각 사진에서 딥러닝을 통해 뎁스 맵(깊이 정보)을 추출합니다.")
    
    with gr.Row():
        input_images = gr.File(label="얼굴 사진 업로드 (최대 10장 권장)", file_count="multiple", type="filepath")
        
    output_gallery = gr.Gallery(label="원본 이미지 및 뎁스 맵", columns=[2], rows=[1], object_fit="contain", height="auto")
    
    process_button = gr.Button("뎁스 맵 추출 시작")
    
    def process_all_images(image_paths):
        if not image_paths:
            return [(None, "이미지를 업로드해주세요.")]
        
        results = []
        for i, path in enumerate(image_paths):
            original_image = Image.open(path).convert("RGB")
            depth_map_pil, error = process_image_for_depth(original_image)
            
            if error:
                print(f"Error processing image {i+1}: {error}")
                results.append((original_image, f"Error: {error}"))
            else:
                results.append((original_image, f"원본 이미지 {i+1}"))
                results.append((depth_map_pil, f"뎁스 맵 {i+1}"))
        return results

    process_button.click(
        fn=process_all_images,
        inputs=input_images,
        outputs=output_gallery
    )
    
# Gradio 앱을 FastAPI에 마운트
app = gr.mount_gradio_app(app, demo, path="/")

# FastAPI 기본 엔드포인트 (선택 사항, Gradio 앱이 기본 경로를 점유함)
@app.get("/api")
def read_root():
    return {"message": "Welcome to the Face Depth Map Extractor! Visit / for the UI."}