yolov5 coco数据集预测

### 使用YOLOv5模型对COCO数据集执行预测操作 #### 准备工作 为了能够顺利运行YOLOv5模型并完成预测任务，需先确认软硬件环境已满足条件。具体来说，操作系统应为Ubuntu 18.04 64位；Python版本建议采用Anaconda下的3.7版；显卡方面推荐NVIDIA GTX 1070Ti及以上型号，并安装CUDA 10.1以及PyTorch 1.5来支持GPU加速计算[^2]。 #### 下载预训练权重与配置文件 获取官方提供的预训练权重文件对于提高检测精度至关重要。通常这些资源可以在GitHub仓库中找到对应的release页面下载链接。同时也要准备好`coco.yaml`这样的配置文件用于指定类别名称和其他参数设置[^1]。 #### 加载模型及设定推理选项 通过加载上述准备好的权重文件初始化YOLOv5网络结构，在此过程中可以调整一些超参比如置信度阈值(`conf`)、非极大抑制(NMS)的IOU阈值(`iou`)等以适应实际应用场景需求： ```python from pathlib import Path import torch from models.experimental import attempt_load from utils.general import check_img_size, non_max_suppression, scale_coords from utils.torch_utils import select_device from utils.datasets import LoadImages weights = 'yolov5s.pt' # 预训练权重路径 imgsz = 640 # 输入图像尺寸 device = '' # 设为空字符串表示自动选择CPU/GPU设备 half = False # 是否开启半精度浮点数运算 # 初始化设备和模型 device = select_device(device) model = attempt_load(weights, map_location=device) # 加载模型 stride = int(model.stride.max()) # 获取步幅大小 imgsz = check_img_size(imgsz, s=stride) if half: model.half() # 将模型转换成FP16模式 ``` #### 执行预测过程 定义好辅助函数之后就可以正式进入预测环节了。这里会遍历测试集中每一张待测图片，经过前处理送入神经网络得到边界框坐标及其对应类别的概率分布向量，最后利用NMS去除冗余候选区域只保留最优解作为最终输出结果之一： ```python source = './inference/images/' # 测试图片所在目录 dataset = LoadImages(source, img_size=imgsz, stride=stride) for path, img, im0s, vid_cap in dataset: img = torch.from_numpy(img).to(device) img = img.half() if half else img.float() img /= 255.0 # 归一化像素值范围至[0, 1] if img.ndimension() == 3: img = img.unsqueeze(0) pred = model(img)[0] conf_thres = 0.25 # 置信度阈值 iou_thres = 0.45 # NMS IOU阈值 det = non_max_suppression(pred, conf_thres=conf_thres, iou_thres=iou_thres)[0] if len(det): det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0s.shape).round() for *xyxy, conf, cls in reversed(det): label = f'{names[int(cls)]} {conf:.2f}' plot_one_box(xyxy, im0s, label=label, color=colors(int(cls)), line_thickness=3) cv2.imshow('result', im0s) cv2.waitKey(0) cv2.destroyAllWindows() ```