import gradio as gr
from transformers import pipeline
import spaces  # 허깅페이스의 ZeroGPU 환경을 사용하기 위한 필수 모듈입니다.

# 1. 이미지 분류 모델 로드
image_classifier = pipeline("image-classification", model="google/vit-base-patch16-224")

@spaces.GPU  # [핵심] ZeroGPU 환경에서 이 함수가 실행될 때 GPU 자원을 할당받도록 지시합니다.
def analyze_image(image):
    """이미지를 분석하여 확률 딕셔너리와 설명을 반환하는 함수"""
    if image is None:
        return None, "이미지를 업로드해주세요."
    
    results = image_classifier(image)
    
    confidences = {res['label']: res['score'] for res in results}
    top_label = results[0]['label']
    description = f"이 이미지는 높은 확률로 **'{top_label}'**과(와) 관련이 있어 보입니다."
    
    return confidences, description

# 화면 구성 (Gradio Blocks)
with gr.Blocks() as demo:  # Gradio 6.0 규격에 맞춰 이곳에 있던 theme 속성을 제거했습니다.
    gr.Markdown("## 🖼️ AI 이미지 분류기")
    
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="이미지 업로드")
            img_btn = gr.Button("이미지 분석", variant="primary")
            
        with gr.Column():
            img_label_output = gr.Label(label="Top 5 클래스 확률", num_top_classes=5)
            img_desc_output = gr.Markdown(label="AI 이미지 설명")
            
    img_btn.click(fn=analyze_image, inputs=img_input, outputs=[img_label_output, img_desc_output])

if __name__ == "__main__":
    # Gradio 6.0 최신 문법에 따라 theme 옵션을 launch() 함수 내부로 이동시켰습니다.
    demo.launch(theme=gr.themes.Soft())