import gradio as gr
import torch
from ultralyticsplus import YOLO
from PIL import Image, ImageDraw, ImageFont
import os
from transformers import pipeline
import numpy as np

# Tải model YOLO (detection)
yolo_model_path = "best.pt"  # Thay đường dẫn model YOLO của bạn
yolo_model = YOLO(yolo_model_path)

# Tải pipeline classification
class_pipe = pipeline("image-classification", model="Hemg/Acne-classification")

def detect_and_classify(image, image_size, conf_thresold=0.4, iou_thresold=0.5):
    # image là đường dẫn file, đọc thành PIL
    pil_image = Image.open(image).convert("RGB")

    # Detect với YOLO
    results = yolo_model.predict(pil_image, conf=conf_thresold, iou=iou_thresold, imgsz=image_size)
    boxes = results[0].boxes
    num_boxes = len(boxes)

    # Nếu không có vùng mụn, trả về kết quả
    if num_boxes == 0:
        severity = "Tốt"
        recommendation = "Làn da bạn khá ổn! Tiếp tục duy trì thói quen chăm sóc da."
        return image, f"Tình trạng mụn: {severity}", recommendation, "Không có loại mụn nào được phát hiện."

    # Chuyển boxes về array
    xyxy = boxes.xyxy.detach().cpu().numpy().astype(int)
    confidences = boxes.conf.detach().cpu().numpy()

    # Crop từng box và classify
    class_names = []
    for box in xyxy:
        x1, y1, x2, y2 = box
        crop = pil_image.crop((x1, y1, x2, y2))
        # Classification trên vùng crop
        results_class = class_pipe(crop)
        top_class = results_class[0]['label']
        class_names.append(top_class)

    # Đánh giá tình trạng dựa trên số lượng mụn
    if num_boxes > 10:
        severity = "Nặng"
        recommendation = "Bạn nên đến gặp bác sĩ da liễu và sử dụng liệu trình trị mụn chuyên sâu."
    elif 5 <= num_boxes <= 10:
        severity = "Trung bình"
        recommendation = "Hãy duy trì skincare đều đặn với sữa rửa mặt dịu nhẹ và dưỡng ẩm phù hợp."
    else:
        severity = "Tốt"
        recommendation = "Làn da bạn khá ổn! Tiếp tục duy trì thói quen chăm sóc da hiện tại."

    # Vẽ bounding box và class name lên ảnh
    draw = ImageDraw.Draw(pil_image)
    font = ImageFont.load_default()

    for i, (box, cname, conf) in enumerate(zip(xyxy, class_names, confidences), start=1):
        x1, y1, x2, y2 = box
        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
        text = f"#{i}: {cname} ({conf:.2f})"

        # Dùng textbbox để xác định kích thước text
        bbox = draw.textbbox((0,0), text, font=font)
        text_w = bbox[2]-bbox[0]
        text_h = bbox[3]-bbox[1]

        # Vẽ nền cho text
        draw.rectangle([x1, y1 - text_h, x1 + text_w, y1], fill="red")
        # Vẽ text
        draw.text((x1, y1 - text_h), text, fill="white", font=font)

    # Lưu ảnh kết quả
    predicted_image_save_path = "predicted_image.jpg"
    pil_image.save(predicted_image_save_path)

    # Liệt kê loại mụn theo số thứ tự
    acne_types_str = "Danh sách mụn phát hiện:\n"
    for i, cname in enumerate(class_names, start=1):
        acne_types_str += f"Mụn #{i}: {cname}\n"

    return predicted_image_save_path, f"Tình trạng mụn: {severity}", recommendation, acne_types_str

description_md = """
## Ứng dụng Nhận Diện & Phân Loại Mụn
- Sử dụng YOLO để phát hiện các vùng mụn trên khuôn mặt.
- Sử dụng mô hình phân loại (Hemg/Acne-classification) để xác định loại mụn.
- Hiển thị bounding box kèm số thứ tự và loại mụn tương ứng.
- Đánh giá tình trạng da và đưa ra khuyến nghị.
"""

custom_css = """
#component-0, #component-1, #component-2, #component-3, #component-4 {
    font-family: "Helvetica Neue", Helvetica, Arial, sans-serif;
    text-align: center;
}
#component-0 h1 {
    color: #F15B2A;
}
#component-0 h2, h3 {
    color: #333;
}
.gr-button {
    background-color: #F15B2A !important;
    color: #fff !important;
    border: none !important;
    font-weight: bold !important;
}
.gr-button:hover {
    background-color: #d94c1f !important;
}
"""

inputs = [
    gr.Image(type="filepath", label="Ảnh Khuôn Mặt"),
    gr.Slider(minimum=320, maximum=1280, step=32, value=640, label="Kích thước ảnh (Image Size)"),
    gr.Slider(minimum=0, maximum=1, step=0.05, value=0.4, label="Ngưỡng Confidence"),
    gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, label="Ngưỡng IOU")
]

outputs = [
    gr.Image(type="filepath", label="Ảnh Sau Khi Xử Lý"),
    gr.Textbox(label="Tình Trạng Mụn", interactive=False),
    gr.Textbox(label="Khuyến Nghị", interactive=False),
    gr.Textbox(label="Loại Mụn Phát Hiện", interactive=False)
]

yolo_app = gr.Interface(
    fn=detect_and_classify,
    inputs=inputs,
    outputs=outputs,
    title="YOLOv8 + Classification: Nhận Diện & Phân Loại Mụn",
    description=description_md,
    css=custom_css,
    theme="default"
)

yolo_app.launch(share=True)