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


def process_image(input_img):
    # Create a temporary directory to store the input image
    with tempfile.TemporaryDirectory() as temp_input_dir:
        input_image_path = os.path.join(temp_input_dir, "input.jpg")
        input_img.save(input_image_path)

        # Create a temporary directory for the output image
        with tempfile.TemporaryDirectory() as temp_output_dir:
            # Command to run the YOLO model
            command = f"yolo task=detect mode=predict model=best.pt conf=0.25 source={temp_input_dir} save=True"
            subprocess.run(command, shell=True)

           
            # Get the most recent 'predict' folder in 'runs/detect'
            list_of_dirs = glob.glob('runs/detect/predict*')
            latest_dir = max(list_of_dirs, key=os.path.getctime)

            # Assuming YOLO saves the output with the same name in the latest 'predict' folder
            output_image_name = os.path.basename(input_image_path)
            output_image_path = os.path.join(latest_dir, output_image_name)

            if os.path.exists(output_image_path):
                output_img = Image.open(output_image_path)
                return output_img
            else:
                return "No output image found."

# Define the Gradio interface
demo = gr.Interface(
    fn=process_image, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Image(type="pil"),
    title="Object Detection with YOLO",
    description="Upload an image and the YOLO model will detect objects."
)

# Launch the app
demo.launch(share=True)