YOLOV5小球定位

### 使用YOLOv5实现小球目标检测与定位

#### 准备工作
为了使用YOLOv5进行小球的目标检测与定位，需先安装必要的依赖库并下载预训练模型。具体操作如下：

对于Python环境而言，推荐创建虚拟环境来管理项目所需的包。

conda create -n yolov5 python=3.8
conda activate yolov5
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

上述命令会设置好运行YOLOv5所需的基础环境，并获取官方仓库中的源码以及相关依赖项[^1]。

#### 数据准备
针对特定的小球类别，应当收集足够的图片样本用于训练自定义模型。这些图像应该覆盖不同光照条件、背景变化等因素下的小球实例。如果可能的话，还可以利用数据增强技术增加多样性。标注工具可以选择LabelImg或其他开源软件完成边界框标记任务。

#### 模型微调
由于默认情况下YOLOv5并不包含专门针对小球这一类别的权重参数，因此建议基于现有网络结构对新加入的数据集执行迁移学习过程。通过调整最后一层输出节点数量适应新增加的对象种类即可满足需求。配置文件`yolov5s.yaml`可根据实际情况修改路径指向本地保存好的标签列表和图片位置。

train: ./data/custom/train/
val:   ./data/custom/test/

nc: 1 # number of classes
names: ['ball'] # class names

接着，在终端输入以下指令启动训练流程:

python train.py --img 416 --batch 16 --epochs 100 --data custom_data.yaml --weights yolov5s.pt --cache

这里指定了输入尺寸大小为416×416像素，批次规模设为16张图每轮迭代次数达到100次等超参设定；同时加载了轻量级版本的初始权值作为起点加快收敛速度。

#### 推理部署
当训练结束后，可选用最佳效果的一组参数导出简化后的ONNX格式供其他平台移植应用。而实时预测阶段则只需读取测试集中任意一张照片传递给已训练完毕的神经网络获得最终结果。

import cv2
from pathlib import Path
import torch
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.datasets import letterbox


def detect_ball(image_path='data/images/bus.jpg', weights='runs/train/exp/weights/best.pt'):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    model = attempt_load(weights, map_location=device)  # 加载模型
    
    img_size = 640
    conf_thres = 0.25
    iou_thres = 0.45
    
    im0 = cv2.imread(str(Path(image_path)))  # BGR
    img = letterbox(im0, new_shape=img_size)[0]

    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
    img = np.ascontiguousarray(img)

    img = torch.from_numpy(img).to(device)
    img = img.float()
    img /= 255.0  # 归一化到[0,1]
    if img.ndimension() == 3:
        img = img.unsqueeze(0)

    pred = model(img, augment=False)[0]
    det = non_max_suppression(pred, conf_thres, iou_thres)[0]

    if len(det):
        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

        for *xyxy, conf, cls in reversed(det):
            label = f'{conf:.2f}'
            plot_one_box(xyxy, im0, label=label, color=(0, 255, 0), line_thickness=3)


cv2.imshow('result', im0)
cv2.waitKey(0)
cv2.destroyAllWindows()

此段脚本实现了从小球所在区域裁剪出来的矩形框绘制功能，并显示带有置信度得分的结果画面。