import gradio as gr
from transformers import AutoProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# 모델 로딩
model_id = "dima806/ai_vs_human_generated_image_detection"
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForImageClassification.from_pretrained(model_id)

def detect_image(image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    probs = torch.nn.functional.softmax(logits, dim=1)
    confidences = probs[0].tolist()
    labels = model.config.id2label

    # 결과 정렬
    result = sorted(zip(labels.values(), confidences), key=lambda x: x[1], reverse=True)

    top_label, top_conf = result[0]
    explanation = "이 이미지는 '{0}'로 분류되었습니다.\n\n신뢰도: {1:.2f}%\n\n".format(top_label, top_conf * 100)

    # 부가 설명 추가
    if "AI" in top_label:
        explanation += "배경 흐림, 비현실적인 디테일, 반복되는 패턴 등 AI 이미지의 전형적 특징이 감지되었습니다."
    else:
        explanation += "조명, 질감, 왜곡 없는 디테일 등 실제 촬영 이미지의 특징이 확인되었습니다."

    return f"{top_label} ({top_conf*100:.2f}%)", explanation

# 블랙&화이트 UI + 왼/오 비율 레이아웃
with gr.Blocks(theme=gr.themes.Base(primary_hue="gray")) as demo:
    gr.Markdown(
        "# 🎨 AI vs Human 이미지 판별기",
        elem_id="title"
    )

    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="이미지 업로드", elem_id="image-box")

        with gr.Column(scale=2):
            label_output = gr.Textbox(label="판별 결과 (가장 높은 신뢰도)", elem_id="result-label", max_lines=1, interactive=False)
            text_output = gr.Textbox(label="판별 이유", lines=10, elem_id="result-text", interactive=False)

    image_input.change(fn=detect_image, inputs=image_input, outputs=[label_output, text_output])

# CSS로 스타일 커스터마이징
demo.css = """
body {
    background-color: #111;
    color: #eee;
    font-family: 'Helvetica Neue', sans-serif;
}
#title {
    text-align: center;
    font-size: 32px;
    color: white;
    margin-bottom: 20px;
}
#image-box {
    background-color: #222;
    border: 1px solid #444;
    padding: 10px;
}
#result-label textarea {
    font-size: 24px;
    color: #00ffcc;
    background-color: #1a1a1a;
}
#result-text textarea {
    font-size: 16px;
    background-color: #1a1a1a;
    color: #ccc;
}
"""

# 실행
if __name__ == "__main__":
    demo.launch()