# --- Phần 1: Nhập khẩu các thư viện cần thiết ---
import gradio as gr
from transformers import AutoImageProcessor, SiglipForImageClassification
from PIL import Image
import torch

# --- Phần 2: Tải mô hình AI và bộ xử lý ảnh từ Hugging Face ---
# (Chỉ chạy một lần khi ứng dụng khởi động)
model_name = "prithivMLmods/Dog-Breed-120"
print("Bắt đầu tải model và processor... Quá trình này có thể mất vài phút.")
processor = AutoImageProcessor.from_pretrained(model_name)
model = SiglipForImageClassification.from_pretrained(model_name)
print("Tải model và processor thành công!")

# Chuyển mô hình sang "chế độ dự đoán" để tối ưu hóa hiệu suất
model.eval()

# --- Phần 3: Định nghĩa hàm dự đoán chính ---
def predict(image):
    # Dùng processor để xử lý ảnh đầu vào thành định dạng mà mô hình có thể hiểu
    inputs = processor(images=image, return_tensors="pt")

    # Đưa ảnh đã xử lý vào mô hình để nhận kết quả
    # torch.no_grad() giúp chạy nhanh hơn vì không cần tính toán cho việc training
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    probs = torch.nn.functional.softmax(logits,dim=1).squeeze()
    top_idx = probs.argmax().item()
    label = model.config.id2label[top_idx]
    confidence = probs[top_idx].item()
    return f"{label} ({round(confidence*100, 1)}%)"

# --- Phần 4: Tạo và khởi chạy giao diện người dùng với Gradio ---
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Tải ảnh chú chó của bạn lên đây"),
    outputs="text",
    title="AI Nhận Diện Giống Chó 🐶",
    description="Ứng dụng này sử dụng mô hình AI 'Dog-Breed-120' để nhận diện 120 giống chó khác nhau. Hãy thử tải một bức ảnh lên!",
)

iface.launch()