yolov5 detect.py获得预测框坐标

### 获取YOLOv5中 `detect.py` 脚本预测框坐标的说明 在YOLOv5的 `detect.py` 中，获取预测框坐标的过程涉及多个步骤。主要通过调用模型并处理其输出来实现。 #### 主要流程概述 当执行检测操作时，程序会加载预训练好的YOLOv5模型并对输入图像进行推理。之后，会对得到的结果进一步处理以提取出边界框的位置信息以及其他属性（如类别标签和置信度）。具体来说： - **数据准备阶段**：确保输入图片尺寸符合网络的要求，并调整通道顺序等[^1]。 - **前向传播计算**：将处理后的图像送入神经网络完成特征抽取及分类回归任务；此过程由PyTorch框架自动管理[^2]。 - **后处理部分**：对原始输出应用非极大抑制算法(NMS)筛选掉冗余候选区域，最终保留下来的即是高质量的目标位置估计——即所需的矩形边界的左上角(xmin,ymin)与右下角(xmax,ymax)。 #### 关键代码片段展示 以下是简化版用于理解如何从模型输出中提取预测框坐标的Python代码示例: ```python import torch from pathlib import Path from models.experimental import attempt_load from utils.general import non_max_suppression, scale_coords from utils.datasets import letterbox def get_prediction_boxes(img_path, weights='yolov5s.pt', imgsz=640, conf_thres=0.25, iou_thres=0.45): device = 'cuda' if torch.cuda.is_available() else 'cpu' # 加载模型 model = attempt_load(weights, map_location=device) # 图像预处理 image = cv2.imread(str(Path(img_path))) im0_shape = image.shape[:2] processed_image = letterbox(image, new_shape=imgsz)[0] processed_image = processed_image[:, :, ::-1].transpose(2, 0, 1).copy() processed_image = torch.from_numpy(processed_image).to(device) processed_image = processed_image.float() / 255.0 if len(processed_image.shape) == 3: processed_image = processed_image[None] with torch.no_grad(): pred = model(processed_image, augment=False)[0] # 应用NMS过滤重复框 det = non_max_suppression(pred, conf_thres, iou_thres)[0] if det is not None and len(det): # 将检测到的对象缩放回原图大小 det[:, :4] = scale_coords(processed_image.shape[2:], det[:, :4], im0_shape).round() result = [] for *xyxy, conf, cls in reversed(det): xywh = (torch.tensor(xyxy).view(1, 4)).tolist()[0] class_id = int(cls.item()) confidence = float(conf.item()) box_info = { "bbox": [int(i) for i in xywh], "class": class_id, "confidence": confidence } result.append(box_info) return result ``` 上述代码展示了完整的预测管道，包括但不限于读取测试集中的单张图片、对其进行必要的转换以便适应模型需求、运行推断以及最后一步解析输出结果为易于理解和使用的格式。注意这里特别强调了对于三维张量添加批次维度的操作，这正是为了满足YOLOv5所期望接收的数据结构形式而采取的关键措施之一。