import gradio as gr
import cv2
import time
import numpy as np
from sahi import AutoDetectionModel
from sahi.predict import get_sliced_prediction
from pathlib import Path

# Load the detection model
detection_model = AutoDetectionModel.from_pretrained(
    model_type='ultralytics',
    model_path="./DDR.pt",  # Replace with your model path
    confidence_threshold=0.01,
    device="cpu"  # Change to 'cuda:0' if you have a GPU
)

OUTPUT_PATH = "./pred_image.jpg"
TEMP_PNG_PATH = "./pred_image.png"

def wait_for_file(file_path, timeout=10):
    """Poll for the file to exist until the timeout (in seconds) is reached."""
    start_time = time.time()
    while not Path(file_path).exists():
        if time.time() - start_time > timeout:
            return False
        time.sleep(0.5)
    return True

def run_inference(image):
    # Perform sliced prediction on the input image.
    result = get_sliced_prediction(
        image,
        detection_model,
        slice_height=256,
        slice_width=256,
        overlap_height_ratio=0.2,
        overlap_width_ratio=0.2
    )
    
    # Export visualization to a temporary PNG file.
    result.export_visuals(export_dir=Path(TEMP_PNG_PATH).parent, file_name=Path(TEMP_PNG_PATH).name)
    
    # Wait for the PNG file to be created.
    if not wait_for_file(TEMP_PNG_PATH, timeout=10):
        raise FileNotFoundError(f"SAHI did not save the PNG file at {TEMP_PNG_PATH}")
    
    # Read the PNG image, convert it to JPG, and remove the temporary file.
    processed_image = cv2.imread(TEMP_PNG_PATH)
    cv2.imwrite(OUTPUT_PATH, processed_image)
    Path(TEMP_PNG_PATH).unlink()  # Delete the temporary PNG
    
    return OUTPUT_PATH

demo = gr.Interface(
    fn=run_inference,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(type="filepath"),
    title="YOLO11 Object Detection",
    description="Upload an image to run inference using YOLO11"
)

demo.launch(share=True)