小目标检测遇到瓶颈,总是有检错或者漏检,我想问一下 在YOLOv5怎么用滑动窗口检测,有相关博文推荐吗?
给你写个实例代码,希望采纳:
import torch
import cv2
from PIL import Image
from numpy import random
from matplotlib import pyplot as plt
# Load YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model='path/to/weights.pt')
# Define input image resolution and stride
img_size = 640 # input image size
stride = img_size // 2 # stride for sliding window
# Load input image
img_path = 'path/to/input/image.jpg'
img = Image.open(img_path)
# Convert image to numpy array
img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
# Define sliding window parameters
window_size = (img_size, img_size) # window size
overlap = 0.5 # overlap ratio
# Compute number of windows and positions
height, width, _ = img.shape
x_steps = int((width - window_size[0]) / (stride * overlap)) + 1
y_steps = int((height - window_size[1]) / (stride * overlap)) + 1
# Initialize empty detection list
detections = []
# Loop over windows
for i in range(x_steps):
for j in range(y_steps):
# Compute window position
x = i * stride * overlap
y = j * stride * overlap
# Crop window from input image
window = img[y:y+window_size[1], x:x+window_size[0]]
# Convert window to PIL image
window = Image.fromarray(window)
# Run YOLOv5 model on window
results = model(window)
# Filter results by confidence threshold and class
results.filter('class', 0, '>', 'confidence', 0.5)
# Add window position to detection boxes
for result in results.xyxy[0]:
box = result.tolist()
box[0] += x
box[1] += y
box[2] += x
box[3] += y
detections.append(box)
# Draw detection boxes on input image
for box in detections:
x1, y1, x2, y2, conf, cls = box
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
# Show output image
plt.imshow(img[:, :, ::-1])
plt.show()