from fastapi import FastAPI
import gradio as gr
from PIL import Image
import numpy as np
import torch
from transformers import pipeline
import cv2 # cv2는 이미지 로드/저장/전처리 등에 사용될 수 있지만, 현재 코드에서는 PIL만으로도 충분합니다.

app = FastAPI()

# 딥러닝 모델 로드 (Depth Anything)
# 모델 로딩은 앱 시작 시 한 번만 하도록 글로벌 변수로 설정
print("Loading Depth Anything model...")
try:
    # 🌟🌟🌟 모델 이름 수정: 정확한 Depth Anything v2 Large 모델 ID 🌟🌟🌟
    # 더 작은 모델을 원하시면 "LiangNX/depth-anything-v2-base-nyu" 로 변경하세요.
    depth_estimator = pipeline(task="depth-estimation", model="LiangNX/depth-anything-v2-large-nyu", device="cpu") # GPU 사용 시 device="cuda"
    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. Check server logs."

    # Gradio는 PIL Image 객체로 이미지를 전달합니다.
    if isinstance(image_path_or_pil_image, str):
        # 파일 경로로 이미지가 전달된 경우 (예: Gradio의 `filepath` 타입)
        image = Image.open(image_path_or_pil_image).convert("RGB")
    else:
        # PIL Image 객체가 직접 전달된 경우
        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) - 이 PIL 이미지는 이미 시각화하기 좋은 형태로 되어 있습니다.
        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)
        # depth_grayscale_pil = Image.fromarray(normalized_depth_np)

        return depth_image_pil, None # 뎁스 맵 PIL Image 반환
    except Exception as e:
        import traceback
        traceback.print_exc() # 에러 발생 시 상세 트레이스백 출력
        return None, f"Error processing image for depth: {e}"

# Gradio 인터페이스 정의
with gr.Blocks() as demo:
    gr.Markdown("# 🧑‍💻 얼굴 뎁스 맵 추출기")
    gr.Markdown("여러 장의 얼굴 사진을 업로드하면 각 사진에서 딥러닝을 통해 뎁스 맵(깊이 정보)을 추출합니다.")
    gr.Markdown("⚠️ **참고:** 무료 CPU 환경에서는 모델 로딩 및 이미지 처리 시간이 매우 오래 걸릴 수 있습니다.")
    
    with gr.Row():
        # file_count="multiple"로 여러 파일 업로드 가능하게 설정
        # type="filepath"로 설정하여 Gradio가 파일을 임시 경로에 저장하도록 함
        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_gallery = [] # Gradio Gallery에 표시할 (PIL Image, Label) 튜플 리스트
        
        for i, path in enumerate(image_paths):
            try:
                original_image = Image.open(path).convert("RGB")
                depth_map_pil, error = process_image_for_depth(original_image)
                
                # 원본 이미지를 먼저 추가
                results_for_gallery.append((original_image, f"원본 이미지 {i+1}"))
                
                if error:
                    print(f"Error processing image {i+1}: {error}")
                    # 오류가 발생한 경우 뎁스 맵 대신 오류 메시지 표시
                    # 이 부분은 Gradio Gallery가 (Image, Label)을 기대하므로,
                    # 오류 메시지를 텍스트 이미지로 만들거나, 오류 이미지를 사용할 수 있습니다.
                    # 여기서는 간단히 빈 이미지와 함께 오류 메시지를 표시하도록 합니다.
                    dummy_error_image = Image.new('RGB', original_image.size, color = 'red')
                    results_for_gallery.append((dummy_error_image, f"오류 발생: {error}"))
                else:
                    results_for_gallery.append((depth_map_pil, f"뎁스 맵 {i+1}"))
            except Exception as e:
                # 파일 자체를 여는 데 문제가 발생한 경우
                print(f"Failed to open or process file {path}: {e}")
                dummy_error_image = Image.new('RGB', (200, 200), color = 'red') # 작은 에러 이미지
                results_for_gallery.append((dummy_error_image, f"파일 처리 오류: {e}"))

        return results_for_gallery

    # 버튼 클릭 시 함수 연결
    process_button.click(
        fn=process_all_images,
        inputs=input_images,
        outputs=output_gallery
    )
    
# Gradio 앱을 FastAPI에 마운트
# 기본 경로 (path="/")에 Gradio 앱을 마운트하여 웹 페이지 접속 시 바로 UI가 보이도록 합니다.
app = gr.mount_gradio_app(app, demo, path="/")

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