from ultralytics import YOLO
import cv2
import numpy as np

# Define constants for size cut-off
MIN_W, MIN_H = 60, 60
MAX_W, MAX_H = 400, 400

# Load the YOLO model
model = YOLO('best (2).pt')
print("YOLO model loaded successfully.")
def predict_image(image_path):
    # Run inference
    results = model(image_path,
                    conf=0.5,
                    iou=0.15,
                    max_det=300,
                    augment=True)

    # Apply size-based post-processing
    if results[0].boxes:
        b = results[0].boxes.xywh # Get boxes in xywh format
        w = b[:, 2] # width
        h = b[:, 3] # height
        
        # Create a boolean mask for valid sizes
        keep = (w >= MIN_W) & (h >= MIN_H) & (w <= MAX_W) & (h <= MAX_H)
        
        # Filter the boxes
        results[0].boxes = results[0].boxes[keep]
        results[0].update() # Update the results object

    # Return the annotated image (as a numpy array in BGR format)
    annotated_image = results[0].plot(labels=True, conf=True)
    return annotated_image

import gradio as gr

# Create the Gradio interface
iface = gr.Interface(fn=predict_image,
                     inputs=gr.Image(type="filepath", label="Upload Image or Use Camera"),
                     outputs=gr.Image(type="numpy", label="Detected Objects"),
                     title="FAWDetect")

# Launch the interface
iface.launch(debug=True)