Spaces:
Runtime error
Runtime error
File size: 1,967 Bytes
88439ab 4685ada | 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 47 48 49 50 51 52 53 | import gradio as gr
import torch
from torchvision import transforms
from models.experimental import attempt_load
from utils.general import non_max_suppression, plot_one_box
import numpy as np
# Load YOLOv7 model
weights_path = "cattle.pt"
config_path = "yolov7.yaml"
# Initialize your YOLOv7 model here
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = attempt_load(weights_path, map_location=device)
# Define function to detect objects using YOLOv7
def detect_objects(image):
global net
# Perform any necessary preprocessing on the image
transform = transforms.Compose([
transforms.Resize((416, 416)), # Resize image to expected input size
transforms.ToTensor(), # Convert image to tensor
])
image_tensor = transform(image).unsqueeze(0)
# Perform object detection
with torch.no_grad():
# Forward pass through the network
outs = net(image_tensor)
# Apply non-maximum suppression
pred = non_max_suppression(outs, conf_thres=0.4, iou_thres=0.5, classes=None, agnostic=False)
# Process detection results and draw bounding boxes
object_count = 0
for i, det in enumerate(pred):
# Skip if no detections
if len(det):
# Increment object count
object_count += len(det)
# Loop over the detections
for *xyxy, conf, cls in reversed(det):
# Draw bounding box
plot_one_box(xyxy, image, label=f'{conf:.2f}', color=(255, 0, 0), line_thickness=3)
return image, object_count
# Create Gradio interface
inputs = gr.inputs.Image(label="Upload Image or Video")
outputs = [gr.outputs.Image(label="Output Image with Objects Detected"), gr.outputs.Text(label="Object Count")]
gr.Interface(detect_objects, inputs, outputs, title="YOLOv7 Object Detection").launch()
|