import gradio as gr
import torch
from ultralyticsplus import YOLO, render_result
from PIL import Image
import os

def yolov8_func(image, 
                image_size, 
                conf_thresold=0.4,
                iou_thresold=0.50):

    # Load the YOLOv8 model
    model_path = "best.pt"
    model = YOLO(model_path)  # Thay bằng đường dẫn model của bạn

    # Dự đoán
    result = model.predict(image, conf=conf_thresold, iou=iou_thresold, imgsz=image_size)
    boxes = result[0].boxes
    num_boxes = len(boxes)

    # In thông tin (tùy chọn)
    print("Object type: ", boxes.cls)
    print("Confidence: ", boxes.conf)
    print("Coordinates: ", boxes.xyxy)
    print(f"Số lượng vùng phát hiện: {num_boxes}")

    # Phân loại mức độ mụn và gợi ý
    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."

    print(f"Tình trạng mụn: {severity}")
    print(f"Khuyến nghị: {recommendation}")

    # Kết xuất ảnh kết quả
    render = render_result(model=model, image=image, result=result[0])
    predicted_image_save_path = "predicted_image.jpg"
    render.save(predicted_image_save_path)

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

# Tạo phần mô tả và hướng dẫn
description_md = """
## Ứng dụng Nhận Diện Mụn bằng YOLOv8
Chọn một ảnh chụp khuôn mặt của bạn. Ứng dụng sẽ phân tích và đánh giá tình trạng mụn dựa trên mô hình YOLOv8.
- **Bước 1:** Tải lên ảnh khuôn mặt.
- **Bước 2:** Điều chỉnh các tham số (nếu cần).
- **Bước 3:** Nhấn **Submit** để nhận kết quả.

Phần kết quả sẽ cho ra:
- Ảnh với bounding box quanh các nốt mụn được phát hiện.
- Đánh giá tình trạng mụn.
- Lời khuyên chăm sóc da.
"""

# CSS tùy chỉnh để giao diện đẹp hơn
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 (Input Image)"),
    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.15, label="Ngưỡng Confidence"),
    gr.Slider(minimum=0, maximum=1, step=0.05, value=0.15, 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)
]

yolo_app = gr.Interface(
    fn=yolov8_func,
    inputs=inputs,
    outputs=outputs,
    title="YOLOv8: Nhận Diện Mụn",
    description=description_md,
    css=custom_css,
    theme="default"  # Bạn có thể thử "default" hoặc "huggingface"
)

yolo_app.launch(debug=True)