File size: 937 Bytes
81ad36a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
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()
|