from detect import *
from pathlib import Path
# Get the current script's directory
script_dir = Path(__file__).resolve().parent

# Define the relative path to the weights file
weights_path = script_dir / 'runs/train/exp17/weights/best.pt'


def yolo_inference(image):
    import tempfile
    temp_dir = tempfile.TemporaryDirectory()
    temp_path = Path(temp_dir.name) / "input_image.jpg"
    output_path = Path(temp_dir.name) / "output_image.jpg"
    image.save(temp_path)

    run(weights=weights_path, source=str(temp_path), project=str(temp_dir.name), name='result', exist_ok=True)
    
    result_image_path = list((Path(temp_dir.name) / 'result').glob('*.jpg'))[0]
    result_image = Image.open(result_image_path)
    
    return result_image

interface = gr.Interface(
    fn=yolo_inference,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="YOLO Object Detection",
    description="Upload an image to detect objects using YOLO."
)

interface.launch()