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

# 모델과 feature extractor 로드
model_name = "shinyice/densenet121-dog-emotions"
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

def predict_emotion(image):
    inputs = feature_extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()
        return model.config.id2label[predicted_class_idx]

# Gradio 인터페이스 생성
interface = gr.Interface(
    fn=predict_emotion,
    inputs=gr.inputs.Image(type="pil"),
    outputs="text",
    title="Dog Emotion Recognition",
    description="Upload an image of your dog and get its predicted emotion."
)

interface.launch()