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)  # Use your custom model path here

    # Make predictions
    result = model.predict(image, conf=conf_thresold, iou=iou_thresold, imgsz=image_size)

    # Access object detection results
    boxes = result[0].boxes  # Bounding boxes
    num_boxes = len(boxes)   # Count the number of bounding boxes (detections)

    # 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])
    
    # Save the rendered image (with predictions)
    predicted_image_save_path = "predicted_image.jpg"
    render.save(predicted_image_save_path)

    # Return the saved image, severity, and recommendation for Gradio output
    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.15, label="IOU Threshold")
]

# Define the output for the Gradio app (image + text for severity and recommendation)
outputs = [
    gr.Image(type="filepath", label="Output Image"),
    gr.Textbox(label="Acne Condition"),
    gr.Textbox(label="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)