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

# Load the fine-tuned YOLOv10 model
model = YOLO('best.pt')

# Inference function for Gradio
def detect_objects(image):
    # Define the results directory
    results_dir = "runs/detect"
    
    # Clear previous detection results to avoid caching issues
    if os.path.exists(results_dir):
        shutil.rmtree(results_dir)  # Remove the entire 'runs/detect' folder
    
    # Run YOLO prediction and save results
    results = model.predict(source=image, save=True, save_txt=True, conf=0.5)

    # Get the latest result folder
    latest_result_folder = sorted(glob.glob(f"{results_dir}/*"), key=os.path.getmtime, reverse=True)[0]
    
    # Find the latest detected image inside the folder
    predicted_images = sorted(glob.glob(f"{latest_result_folder}/*.jpg") + glob.glob(f"{latest_result_folder}/*.png"), key=os.path.getmtime, reverse=True)

    if predicted_images:
        return predicted_images[0]  # Return the latest output image
    else:
        return "Error: No output image found!"

# Create the Gradio interface
app = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="filepath", label="📤 Upload Image"),
    outputs=gr.Image(type="filepath", label="✅ Detected Image"),
    title="🔬 YOLOv10 Blood Cell Detection App",
    description="Upload an image to detect blood cells (RBC, WBC, Platelets) using the fine-tuned YOLOv10 model."
)

# Launch the app
app.launch()