from transformers import AutoImageProcessor, AutoModelForImageSegmentation
from PIL import Image
import torch
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt

# --- ตั้งชื่อโมเดล (ใช้ของ Hugging Face) ---
model_name = "Arya6967/SAMMed-2d"  # ตัวอย่าง
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageSegmentation.from_pretrained(model_name)

# --- โหลดภาพ ---
img_path = "your_image.png"  # เปลี่ยนเป็นภาพร่างกาย/diagram ของคุณ
img = Image.open(img_path).convert("RGB")

# --- ทำ segmentation ---
inputs = processor(img, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)

# --- แปลง logits -> mask ---
mask = outputs.logits.argmax(dim=1)[0].cpu().numpy()  # mask ของแต่ละ pixel
# สมมติ organ_id เริ่มจาก 0
organ_names = [
    "Heart", "Liver", "Lung L", "Lung R", "Stomach", "Kidney L", "Kidney R",
    "Brain", "Pancreas", "Spleen", "Bladder", "Intestine", "Esophagus", "Gallbladder",
    "Appendix", "Thyroid", "Trachea", "Aorta", "Vena Cava", "Pelvis",
    "Femur L", "Femur R", "Humerus L", "Humerus R", "Radius L", "Radius R",
    "Ulna L", "Ulna R", "Clavicle L", "Clavicle R", "Scapula L", "Scapula R",
    "Skull", "Spine", "Ribcage", "Sternum", "Patella L", "Patella R",
    "Tibia L", "Tibia R", "Fibula L", "Fibula R", "Shoulder Joint L", "Shoulder Joint R",
    "Elbow Joint L", "Elbow Joint R", "Knee Joint L", "Knee Joint R", "Ankle Joint L", "Ankle Joint R"
]

# --- ฟังก์ชันคลิก ---
def show_organ(click_data):
    # click_data: (x, y) pixel
    x = int(click_data[0])
    y = int(click_data[1])
    organ_id = mask[y, x]
    if organ_id < len(organ_names):
        return f"อวัยวะ: {organ_names[organ_id]}"
    else:
        return "ไม่ทราบอวัยวะ"

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## คลิกบนภาพเพื่อดูอวัยวะ")
    img_component = gr.Image(value=np.array(img), interactive=True, tool="select")
    output_text = gr.Textbox(label="ผลลัพธ์")
    img_component.select(show_organ, inputs=img_component, outputs=output_text)

demo.launch()