# -*- coding: utf-8 -*-
"""
แอปจำแนกอวัยวะ / ส่วนของผิวหนัง
โดยใช้โมเดลที่เทรนมาแล้วจาก Hugging Face
"""

from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch
from PIL import Image
import requests

# -----------------------------
# กำหนดชื่อโมเดล
# -----------------------------
model_name = "histai/SPIDER-skin-model"

# -----------------------------
# โหลดโมเดลและตัวประมวลผล
# -----------------------------
print("กำลังโหลดโมเดลและตัวประมวลผล...")
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)
print("โหลดโมเดลเรียบร้อย!")

# -----------------------------
# คำอธิบายของอวัยวะ / สภาพผิวหนัง (ตัวอย่าง)
# -----------------------------
organ_descriptions = {
    0: "ผิวหนังปกติ",
    1: "แผลผิวหนัง",
    2: "ตุ่มน้ำ / ฟองอากาศ",
    3: "จุดเลือดออกใต้ผิวหนัง",
    4: "รอยไหม้",
    5: "การอักเสบ",
    # สามารถเพิ่มอวัยวะ/สภาพผิวเพิ่มเติมได้
}

# -----------------------------
# โหลดภาพ (ใส่ URL หรือ path ของไฟล์)
# -----------------------------
# ตัวอย่าง: โหลดจาก URL
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers_logo.png"
image = Image.open(requests.get(url, stream=True).raw)

# หากต้องการใช้ไฟล์ในเครื่อง uncomment ด้านล่าง
# image_path = "path/to/your/image.jpg"
# image = Image.open(image_path)

# -----------------------------
# เตรียมข้อมูลสำหรับโมเดล
# -----------------------------
inputs = processor(images=image, return_tensors="pt")

# -----------------------------
# ทำนายผล
# -----------------------------
with torch.no_grad():
    logits = model(**inputs).logits
    predicted_class_idx = logits.argmax(-1).item()

# -----------------------------
# แสดงผลลัพธ์
# -----------------------------
description = organ_descriptions.get(predicted_class_idx, "ไม่ทราบอวัยวะ / สภาพผิว")
print("------ ผลลัพธ์ ------")
print(f"รหัสคลาสที่ทำนาย: {predicted_class_idx}")
print(f"คำอธิบาย: {description}")