小目标检测遇到瓶颈,总是有检错或者漏检,我想问一下 在YOLOv5怎么用滑动窗口检测,有相关博文推荐吗?
1条回答 默认 最新
关注给你写个实例代码,希望采纳:
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()解决 无用评论 打赏 举报