ultralytics-yolov5库

### Ultralytics YOLOv5 库安装与使用指南 #### 安装依赖项 为了顺利运行 Ultralytics YOLOv5，需先确保环境已配置好 Python 和 pip 工具。接着通过命令行工具执行如下指令来完成必要的软件包安装： ```bash pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 ``` 上述命令会下载并安装 PyTorch 及其相关组件到本地环境中[^2]。 #### 下载并设置 YOLOv5 仓库 获取官方 GitHub 上托管的最新版本代码库，并切换至指定分支以获得稳定版特性支持： ```bash git clone https://github.com/ultralytics/yolov5.git # 克隆仓库 cd yolov5 # 进入目录 pip install -r requirements.txt # 安装需求文件中的Python模块 ``` 此过程将建立一个完整的开发平台用于后续训练、评估以及部署目标检测模型[^1]。 #### 训练自定义数据集上的新模型实例 利用预训练权重初始化网络参数可以加速收敛速度同时提高泛化能力。下面给出一段简单的脚本用来启动一次新的训练任务： ```python from pathlib import Path import yaml # 设置路径变量 data_path = 'path/to/custom_dataset.yaml' weights = 'yolov5s.pt' # 或者其他可用选项如'yolov5m', 'yolov5l', 'yolov5x' # 修改yaml文件指向正确的图片位置 with open(data_path, 'r') as f: data_dict = yaml.safe_load(f) for k in ['train', 'val']: path_key = Path(data_dict[k]) if not path_key.is_absolute(): abs_path = (Path.cwd() / path_key).resolve() data_dict[k] = str(abs_path) with open(data_path, 'w') as file: documents = yaml.dump(data_dict, file) # 启动训练进程 !python train.py --img 640 --batch 16 --epochs 50 --data {data_path} --weights {weights} ``` 这段代码展示了如何调整 `custom_dataset.yaml` 文件内的相对路径为绝对路径，并调用内置命令接口开始一轮迭代优化周期。 #### 推理预测图像或视频流 当拥有经过良好训练后的模型之后，就可以着手处理实际场景下的视觉识别挑战了。这里提供了一个简单例子演示怎样加载保存下来的 checkpoint 并对其输入样本做出响应： ```python import torch from utils.general import non_max_suppression, scale_coords from models.experimental import attempt_load from PIL import ImageDraw, ImageFont import numpy as np def detect(image_paths=['image.jpg']): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = attempt_load(weights='best.pt').to(device) # 加载最佳性能模型 results = [] for img_file in image_paths: im = cv2.imread(img_file)[..., ::-1] h, w = im.shape[:2] ratio = min(640/h, 640/w) nh, nw = int(h*ratio), int(w*ratio) pad_h, pad_w = ((nh+31)//32)*32-nh, ((nw+31)//32)*32-nw resized_im = cv2.resize(im, dsize=(nw,nh)) padded_im = np.zeros((nh+pad_h, nw+pad_w, 3)).astype(np.uint8)+114 padded_im[:nh,:nw,:] = resized_im input_tensor = torch.from_numpy(padded_im.transpose([2,0,1])[None]).float().div_(255.).sub_( .5).mul_(2.).contiguous().to(device) pred = model(input_tensor)[0].detach().cpu().numpy() boxes = non_max_suppression(pred=pred, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False)[0] draw = ImageDraw.Draw(Image.open(img_file)) font = ImageFont.load_default() colors = [(np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255)) for _ in range(len(boxes))] scaled_boxes = scale_coords(shape1=[input_tensor.size(-2), input_tensor.size(-1)], shape2=[im.shape[0], im.shape[1]], coords=torch.tensor([[box['xmin'], box['ymin'], box['xmax'], box['ymax']]]) ).round().tolist()[0] label = '%s %.2f' % ('class_name', float(box[-2])) t_size = draw.textbbox(xy=(scaled_boxes[0], max(scaled_boxes[1]-font.getsize(label)[-1][1], 0)), text=label, fill="black", outline="white") draw.rectangle(t_size, fill=colors[i]) draw.text(((t_size[0]+t_size[2])//2-(len(label)-1)*(font.getsize(label)[-1][-1]), max(scaled_boxes[1]-font.getsize(label)[-1][1], 0)),label,fill="white",font=font) draw.rectangle([(int(x.item()) for x in scaled_boxes)],outline=color,width=2) if __name__ == '__main__': detect(['test_image.png']) ``` 该函数接收一组或多张待测图片列表作为参数，在每一张上标记出所有被发现的目标物体边界框及其类别名称。