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)  
    
    # Make predictions
    result = model.predict(image, conf=conf_thresold, iou=iou_thresold, imgsz=image_size)

    # Access object detection results
    boxes = result[0].boxes  
    num_boxes = len(boxes)   

    # Categorize based on number of boxes (detections) and provide recommendations
    if num_boxes > 10:
        severity = "Worse"
        recommendation = "It is recommended to see a dermatologist and start stronger acne treatment."
    elif 5 <= num_boxes <= 10:
        severity = "Medium"
        recommendation = "You should follow a consistent skincare routine with proper cleansing and moisturizing."
    else:
        severity = "Good"
        recommendation = "Your skin looks good! Keep up with your current skincare routine."

    # Render the result (with bounding boxes/labels)
    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"Acne condition: {severity}", recommendation

# Create the Gradio
with gr.Blocks() as yolo_app:
    gr.Markdown("# YOLOv8: An Object Detection for Acne")

    with gr.Row():
        with gr.Column(scale=1):  # Left side with input
            input_image = gr.Image(type="filepath", label="Input Image")
            image_size = gr.Slider(minimum=320, maximum=1280, step=32, value=640, label="Image Size")
            conf_thresh = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.15, label="Confidence Threshold")
            iou_thresh = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.2, label="IOU Threshold")
            submit_btn = gr.Button("Submit")

        with gr.Column(scale=1):  # Right side with output
            output_image = gr.Image(type="filepath", label="Output Image")
            acne_condition = gr.Textbox(label="Acne Condition")
            recommendation = gr.Textbox(label="Recommendation")

    # Link the submit button to the function
    submit_btn.click(fn=yolov8_func, 
                     inputs=[input_image, image_size, conf_thresh, iou_thresh], 
                     outputs=[output_image, acne_condition, recommendation])

# Launch the app
yolo_app.launch(debug=True)