import gradio as gr
import subprocess
import os
from PIL import Image

# YOLOv8 detection function
def detect_objects(image):
    # Save the uploaded image temporarily
    input_image_path = "input.jpg"
    image.save(input_image_path)

    # Define the YOLOv8 command
    output_dir = "./runs/detect"
    os.makedirs(output_dir, exist_ok=True)

    command = [
        "yolo",  # YOLOv8 CLI command
        "task=detect",  # Specify the task as detection
        "mode=predict",  # Set mode to predict
        f"model=best.pt",  # Path to your YOLOv8 model weights
        f"source={input_image_path}",  # Input image path (uploaded by user)
        f"project={output_dir}",  # Output directory
        f"name=result",  # Subfolder name for results
        "exist_ok=True"  # Allow overwriting existing results
    ]

    # Run YOLOv8 inference
    try:
        subprocess.run(command, check=True)
    except Exception as e:
        return f"Error during YOLO inference: {e}"

    # Get the path to the output image
    output_image_path = os.path.join(output_dir, "result", os.path.basename(input_image_path))

    # Return the output image for display as a PIL image
    if os.path.exists(output_image_path):
        return Image.open(output_image_path)
    else:
        return "Error: Output image not found."

# Create the Gradio interface
interface = gr.Interface(
    fn=detect_objects,  # The function to call when an image is uploaded
    inputs=gr.Image(type="pil"),  # Accept images as input
    outputs=gr.Image(type="pil"),  # Return a PIL image for display
    title="YOLOv8 Object Detection",
    description="Upload an image of floating waste in water, and this app will detect it using YOLOv8."
)

# Launch the Gradio app
interface.launch(share=True)  # Use share=True to get a public URL