Spaces:

sukeshs
/

FAWDetect

Sleeping

Update app.py

4eba9eb verified 10 months ago

1.4 kB

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)

	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)