yolo跑transmission_detection_512
时间: 2025-03-19 16:19:17 浏览: 33
### 运行 YOLO 框架进行传输检测任务
要使用 YOLO V5 框架运行 `transmission_detection_512` 任务,可以按照以下方法配置环境并执行相关命令。YOLO V5 是一种高效的实时目标检测框架,适用于多种场景下的物体识别任务[^1]。
#### 配置开发环境
首先需要安装 Python 和必要的依赖库。推荐使用虚拟环境来管理项目所需的包版本。创建虚拟环境并通过 pip 安装所需依赖项:
```bash
python -m venv yolov5_env
source yolov5_env/bin/activate
pip install torch torchvision numpy opencv-python matplotlib pyyaml tqdm
```
下载官方发布的 YOLOv5 源码仓库,并克隆到本地目录中:
```bash
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt
```
#### 数据集准备
为了完成 `transmission_detection_512` 的特定任务,需准备好数据集文件夹结构以及对应的标注信息。通常情况下,数据应遵循 COCO 或 Pascal VOC 格式。假设已有的训练图像尺寸为 512×512 像素,则可以直接利用预定义参数设置。
编辑 YAML 文件描述自定义的数据路径和类别名称:
```yaml
train: ../datasets/transmission/train/images/
val: ../datasets/transmission/valid/images/
nc: 80
names: ['bird', 'plane', ... ] # 自定义标签列表
```
保存上述内容至名为 `data/custom.yaml` 的位置以便后续调用。
#### 训练模型
启动训练过程前确认 GPU 是否可用加速计算速度。如果硬件支持 CUDA 加速则会显著提升效率。通过如下指令开始训练新的权重或者微调现有模型:
```bash
python train.py --img 512 --batch 16 --epochs 50 --data data/custom.yaml --cfg models/yolov5s.yaml --weights '' --name transmission_detection_results
```
此脚本中的选项解释如下:
- `--img`: 输入图片大小设为 512x512;
- `--batch`: 批处理数量设定为 16 张图每批次;
- `--epochs`: 总共迭代轮次定为 50 轮;
- `--data`: 使用刚才建立好的 dataset configuration file;
- `--cfg`: 模型架构选用 small 版本作为基础模板;
- `--weights`: 不加载任何预先训练过的权重(即从头学习);
- `--name`: 输出日志存储子目录命名规则。
#### 推理预测
当获得满意的模型之后即可用于实际应用当中做推理操作。下面展示一段简单的代码片段演示如何读取一张测试照片并标记出其中的目标对象边界框及其分类概率值。
```python
import torch
from PIL import Image
from utils.general import non_max_suppression, scale_coords
from utils.datasets import letterbox
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.load('runs/exp/weights/best.pt')['model'].float().eval()
half = device.type != 'cpu'
if half:
model.half()
def detect(image_path):
img_size = 512
conf_thres = 0.4
iou_thres = 0.5
original_image = Image.open(image_path).convert("RGB")
resized_img, _, _ = letterbox(original_image, new_shape=img_size)
input_tensor = (torch.from_numpy(resized_img.transpose(2, 0, 1)).to(device))
input_tensor = input_tensor.float() / 255.
if half:
input_tensor = input_tensor.half()
pred = model(input_tensor.unsqueeze(0))[0]
det = non_max_suppression(pred, conf_thres=conf_thres, iou_thres=iou_thres)[0]
if det is not None and len(det):
det[:, :4] = scale_coords(img_size, det[:, :4], original_image.size[::-1]).round()
for *xyxy, conf, cls in reversed(det):
label = f'{cls} {conf:.2f}'
plot_one_box(xyxy, np.array(original_image), label=label, color=(0, 255, 0))
detect('./test_images/sample.jpg')
original_image.show()
```
阅读全文
相关推荐
















