### YOLOv5 单目标检测 Python 实现 以下是基于 PyTorch 的 YOLOv5 单目标检测的代码示例。此代码假设已经安装并配置好了 YOLOv5 环境，并且提供了预训练权重文件 `yolov5x.pt`。 #### 修改后的推理脚本 通过调整默认的 `detect.py` 脚本，使其专注于单目标检测： ```python import torch from models.experimental import attempt_load from utils.general import non_max_suppression, scale_coords from utils.datasets import letterbox import cv2 import numpy as np def preprocess_image(image_path, img_size=640): """加载图像并进行预处理""" img = cv2.imread(image_path) img = letterbox(img, new_shape=img_size)[0] img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB and HWC to CHW img = np.ascontiguousarray(img).astype(np.float32) / 255.0 # 归一化到 [0, 1] return torch.from_numpy(img).unsqueeze(0) def postprocess_predictions(preds, names, conf_thres=0.4, iou_thres=0.5): """后处理预测结果""" det = non_max_suppression(preds, conf_thres, iou_thres)[0] # NMS 过滤 if det is not None and len(det): for *xyxy, conf, cls in reversed(det): label = f'{names[int(cls)]} {conf:.2f}' print(f'检测到的目标: {label}, 坐标: {tuple(map(int, xyxy))}') else: print('未检测到任何目标') if __name__ == "__main__": weights = './weights/yolov5x.pt' # 权重路径 source = './inference/images/zidane.jpg' # 测试图片路径 device = 'cuda' if torch.cuda.is_available() else 'cpu' model = attempt_load(weights, map_location=device) # 加载模型 stride = int(model.stride.max()) # 模型步幅 imgsz = 640 # 输入尺寸 names = model.module.names if hasattr(model, 'module') else model.names # 类别名称 img = preprocess_image(source, img_size=imgsz) # 预处理输入图像 pred = model(img.to(device))[0] # 推理 postprocess_predictions(pred.cpu(), names) # 后处理预测结果 ``` --- #### 关键点说明 1. **模型加载**: 使用 `attempt_load` 函数加载预训练权重文件[^1]。 2. **数据预处理**: 图像被缩放到统一大小 (通常为 640×640)，并通过 `letterbox` 方法填充以保持宽高比不变。 3. **推理阶段**: 将预处理后的张量传递给模型，获取原始预测结果。 4. **NMS 处理**: 应用非极大值抑制 (`non_max_suppression`) 对预测框去冗余[^1]。 5. **类别映射**: 输出检测到的目标及其置信度分数和边界框坐标[^1]。 --- #### 性能与精度权衡 YOLOv5 提供了多种模型变体（如 `yolov5n`, `yolov5s`, ..., `yolov5x`），它们在速度和准确性之间进行了不同的平衡[^2]。对于单目标检测场景，可以选择较大的模型（如 `yolov5x`）来获得更高的检测精度，尽管其运行时间较长。 --- #### 训练与部署注意事项 如果需要自定义训练模型，可以在训练结束后导出最佳权重文件 `best.pt` 并用于推理[^5]。此外，TensorFlow 用户可能更熟悉 `.pb` 文件格式，而 YOLOv5 默认使用 PyTorch 的 `.pt` 格式存储模型参数[^3]。 ---