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)   

    # Print object detection details (optional)
    print("Object type: ", boxes.cls)
    print("Confidence: ", boxes.conf)
    print("Coordinates: ", boxes.xyxy)
    print(f"Number of bounding boxes: {num_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."

    print(f"Acne condition: {severity}")
    print(f"Recommendation: {recommendation}")

    # 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

# Define inputs for the Gradio app
inputs = [
    gr.Image(type="filepath", label="Input Image"),
    gr.Slider(minimum=320, maximum=1280, step=32, value=640, label="Image Size"),
    gr.Slider(minimum=0, maximum=1, step=0.05, value=0.15, label="Confidence Threshold"),
    gr.Slider(minimum=0, maximum=1, step=0.05, value=0.2, label="IOU Threshold")
]

# Use a Row layout to align the textboxes for condition and recommendation
output_image = gr.Image(type="filepath", label="Output Image")
acne_condition = gr.Textbox(label="Acne Condition")
recommendation = gr.Textbox(label="Recommendation")

# Define the layout using Rows and Columns
outputs = [
    output_image,
    gr.Row([acne_condition, recommendation])
]

# Set the title of the Gradio app
title = "YOLOv8: An Object Detection for Acne"

# Create the Gradio interface
yolo_app = gr.Interface(fn=yolov8_func,
                    inputs=inputs,
                    outputs=outputs,
                    title=title)

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