File size: 1,404 Bytes
68f71b7 4eba9eb 68f71b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | 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) |