import gradio as gr
from ultralytics import YOLO
import os
import shutil

# 1. Load the fine-tuned YOLO model
model_path = "best.pt"
model = YOLO(model_path)

# Define the directory to save inference results
output_dir = "inference_results"
os.makedirs(output_dir, exist_ok=True)

# 2. Define an inference function for images
def predict_image(image_file, conf_threshold):
    if image_file is None:
        return "No input image provided."

    print(f"Processing: {image_file}")
    print(f"Confidence threshold: {conf_threshold}")

    # Clear previous results if any
    for item in os.listdir(output_dir):
        item_path = os.path.join(output_dir, item)
        if os.path.isfile(item_path):
            os.remove(item_path)
        elif os.path.isdir(item_path):
            shutil.rmtree(item_path)

    results = model.predict(source=image_file, conf=conf_threshold, save=True, project=output_dir, name="run", exist_ok=True)

    # The results are saved in a subdirectory within output_dir/run/
    # We need to find the actual path to the saved file.
    # Ultralytics saves to runs/detect/runX (where X is an incrementing number)
    # We need to find the latest run folder.

    # Get the latest run folder within the output_dir
    run_folders = [d for d in os.listdir(output_dir) if os.path.isdir(os.path.join(output_dir, d))]
    run_folders.sort(key=lambda x: os.path.getmtime(os.path.join(output_dir, x)), reverse=True)

    if not run_folders:
        return "No detection results saved."

    latest_run_path = os.path.join(output_dir, run_folders[0])

    # Now find the actual image file inside this run_folder
    detected_files = [f for f in os.listdir(latest_run_path) if f.endswith(('.jpg', '.jpeg', '.png'))]

    if not detected_files:
        return "No detected image found in results."

    # Assuming only one file is processed at a time
    result_path = os.path.join(latest_run_path, detected_files[0])

    print(f"Results saved to: {result_path}")
    return result_path

# 3. Create a Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# YOLOv8 Signature Detection")
    gr.Markdown("Upload an image to perform signature detection using a fine-tuned YOLOv8n model.")

    with gr.Tab("Image Detection"):
        image_input = gr.Image(type="filepath", label="Upload Image")
        image_conf_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold")
        image_output = gr.Image(label="Detection Results")
        image_button = gr.Button("Detect Signature in Image")
        image_button.click(predict_image, inputs=[image_input, image_conf_slider], outputs=image_output)

# 4. Launch the Gradio interface
demo.launch(share=True)