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

# ✅ ใช้โมเดลเปิดสาธารณะจาก Google
model_name = "google/vit-base-patch16-224"

# โหลดโมเดลและ processor
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

# ตัวอย่าง label ที่เรากำหนดเอง (อวัยวะหลัก ๆ)
organ_labels = ["สมอง", "หัวใจ", "ปอด", "ตับ", "ไต", "กระเพาะอาหาร", "ลำไส้", "ตา", "มือ", "เท้า"]

# ฟังก์ชันวิเคราะห์ภาพ
def classify_organ(image):
    image = Image.fromarray(image)
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
        top_probs, top_indices = torch.topk(probs, 3)
    # ✅ แก้บรรทัดนี้ให้ถูก
    results = {organ_labels[i.item() % len(organ_labels)]: float(top_probs[0][j]) for j, i in enumerate(top_indices[0])}
    return results

# ส่วนติดต่อ Gradio
app = gr.Interface(
    fn=classify_organ,
    inputs=gr.Image(type="numpy", label="📸 อัปโหลดภาพอวัยวะ"),
    outputs=gr.Label(num_top_classes=3, label="ผลลัพธ์จาก AI"),
    title="🧠 AI จำแนกอวัยวะในร่างกาย",
    description="อัปโหลดภาพอวัยวะ แล้วให้ AI พยายามทำนายว่าเป็นส่วนไหนของร่างกาย (จำลองการทำงานจาก ViT)"
)

if __name__ == "__main__":
    app.launch()