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

# โหลดโมเดล
model_name = "microsoft/resnet-50"
model = ResNetForImageClassification.from_pretrained(model_name)
processor = AutoImageProcessor.from_pretrained(model_name)

# Mock-up: แมป class_id ไปยัง BMI category
def map_to_bmi(class_id):
    if class_id < 250:
        return "Underweight"
    elif class_id < 500:
        return "Normal"
    elif class_id < 750:
        return "Overweight"
    else:
        return "Obese"

# Mock-up: แมป class_id ไปยัง Body Type
def map_to_body_type(class_id):
    if class_id % 3 == 0:
        return "Ectomorph (ผอมเพรียว)"
    elif class_id % 3 == 1:
        return "Mesomorph (สมส่วน/ล่ำ)"
    else:
        return "Endomorph (ล่ำอวบ)"

# ฟังก์ชันประมวลผลภาพ
def analyze_image(image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits
    class_id = logits.argmax(-1).item()
    
    bmi = map_to_bmi(class_id)
    body_type = map_to_body_type(class_id)
    
    return f"🧍 Body Type: {body_type}\n📏 BMI Category: {bmi}"

# Gradio Interface
demo = gr.Interface(
    fn=analyze_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="BMI + Body Type Estimator (Demo)",
    description="วิเคราะห์ BMI และลักษณะรูปร่างจากภาพถ่ายด้วย ResNet-50 (จำลอง)"
)

demo.launch()