import torch
import cv2
import numpy as np
from PIL import Image
import gradio as gr
from ultralytics import YOLO
from transformers import AutoImageProcessor, ResNetForImageClassification

# โหลด YOLOv8 โมเดลสำหรับตรวจจับบุคคล
yolo_model = YOLO("yolov8n.pt")  # ใช้เวอร์ชันเล็ก (nano)

# โหลด ResNet-50 สำหรับจำแนกภาพ
resnet_model_name = "microsoft/resnet-50"
resnet_model = ResNetForImageClassification.from_pretrained(resnet_model_name)
processor = AutoImageProcessor.from_pretrained(resnet_model_name)

# ฟังก์ชันสำหรับ mapping class_id → BMI
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"

# ฟังก์ชันสำหรับ mapping 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 detect_and_crop_person(image_np):
    results = yolo_model(image_np)

    for result in results:
        for box in result.boxes:
            class_id = int(box.cls)
            if yolo_model.names[class_id] == 'person':
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                cropped = image_np[y1:y2, x1:x2]
                cropped_rgb = cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB)
                return Image.fromarray(cropped_rgb)
    return None

# ฟังก์ชันหลักสำหรับ Gradio
def process_image(image):
    # Convert PIL to numpy
    image_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    
    # ตรวจจับและครอบตัดบุคคล
    cropped = detect_and_crop_person(image_np)
    if cropped is None:
        return "⚠️ ไม่พบบุคคลในภาพ"

    # วิเคราะห์ด้วย ResNet
    inputs = processor(images=cropped, return_tensors="pt")
    with torch.no_grad():
        logits = resnet_model(**inputs).logits
    class_id = logits.argmax(-1).item()

    # แมปผลลัพธ์
    bmi = map_to_bmi(class_id)
    body_type = map_to_body_type(class_id)
    label = resnet_model.config.id2label[class_id]

    return f"📷 ResNet Label: {label}\n🧍 Body Type: {body_type}\n📏 BMI Category: {bmi}"

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

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