import gradio as gr
from transformers import AutoImageProcessor, SiglipForImageClassification
from PIL import Image
import torch

# 1. Tải mô hình AI và bộ xử lý ảnh từ Hugging Face
model_name = "prithivMLmods/Dog-Breed-120"
processor = AutoImageProcessor.from_pretrained(model_name)
model = SiglipForImageClassification.from_pretrained(model_name)

# 2. Định nghĩa hàm dự đoán
def predict(image):
    # Xử lý ảnh đầu vào để mô hình có thể hiểu được
    inputs = processor(images=image, return_tensors="pt")

    # Đưa ảnh đã xử lý vào mô hình để nhận kết quả
    with torch.no_grad():
        outputs = model(**inputs)
    
    logits = outputs.logits
    
    # Tìm ra giống chó có xác suất cao nhất
    predicted_class_idx = logits.argmax(-1).item()
    
    # Lấy tên của giống chó và độ tin cậy
    breed_name = model.config.id2label[predicted_class_idx]
    confidence_score = torch.nn.functional.softmax(logits, dim=1)[0][predicted_class_idx].item()
    
    # Định dạng kết quả trả về
    return f"{breed_name} ({confidence_score:.1%})"

# 3. Tạo giao diện người dùng
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Tải ảnh chú chó của bạn lên đây"),
    outputs=gr.Textbox(label="Kết quả dự đoán"),
    title="AI Nhận Diện Giống Chó",
    description="Ứng dụng này sử dụng mô hình AI để nhận diện 120 giống chó khác nhau. Hãy thử tải một bức ảnh lên!",
    
)

# 4
iface.launch()