使用yolo-v8检测视频中的老鼠

2024-04-15 2024-04-15 约 643 字预计阅读 3 分钟

创建conda环境

conda create -n moused python=3.10

进入conda环境

conda activate moused

配置环境

mkdir moused
cd moused
pip install ultralytics

预测

yolo detect predict model=yolov8x.pt source=mouse.mp4

效果并不好, 需要使用自定义数据集微调.

使用自定义数据集微调

经过挑选, 下载安装X-AnyLabeling作为数据标注工具, 截取一些有老鼠的图像, 并且标注检测框, 标注完成后, 新建一个标签文件classes.txt, 内容如下:

mouse

选择导出为yolo标签, 会在图片目录同级出现labels目录.

拆分训练集

将image目录和labels目录放到一个data目录下, 然后运行如下代码:

import os
import random
import shutil

# 设置随机数种子
random.seed(123)

# 定义文件夹路径
root_dir = 'data'
image_dir = os.path.join(root_dir, 'images')
label_dir = os.path.join(root_dir, 'labels')
output_dir = 'dataset'

# 定义训练集、验证集和测试集比例
train_ratio = 0.7
valid_ratio = 0.15
test_ratio = 0.15

# 获取所有图像文件和标签文件的文件名（不包括文件扩展名）
image_filenames = [os.path.splitext(f)[0] for f in os.listdir(image_dir)]
label_filenames = [os.path.splitext(f)[0] for f in os.listdir(label_dir)]

# 随机打乱文件名列表
random.shuffle(image_filenames)

# 计算训练集、验证集和测试集的数量
total_count = len(image_filenames)
train_count = int(total_count * train_ratio)
valid_count = int(total_count * valid_ratio)
test_count = total_count - train_count - valid_count

# 定义输出文件夹路径
train_image_dir = os.path.join(output_dir, 'train', 'images')
train_label_dir = os.path.join(output_dir, 'train', 'labels')
valid_image_dir = os.path.join(output_dir, 'valid', 'images')
valid_label_dir = os.path.join(output_dir, 'valid', 'labels')
test_image_dir = os.path.join(output_dir, 'test', 'images')
test_label_dir = os.path.join(output_dir, 'test', 'labels')

# 创建输出文件夹
os.makedirs(train_image_dir, exist_ok=True)
os.makedirs(train_label_dir, exist_ok=True)
os.makedirs(valid_image_dir, exist_ok=True)
os.makedirs(valid_label_dir, exist_ok=True)
os.makedirs(test_image_dir, exist_ok=True)
os.makedirs(test_label_dir, exist_ok=True)

# 将图像和标签文件划分到不同的数据集中
for i, filename in enumerate(image_filenames):
    if i < train_count:
        output_image_dir = train_image_dir
        output_label_dir = train_label_dir
    elif i < train_count + valid_count:
        output_image_dir = valid_image_dir
        output_label_dir = valid_label_dir
    else:
        output_image_dir = test_image_dir
        output_label_dir = test_label_dir

    # 复制图像文件
    src_image_path = os.path.join(image_dir, filename + '.png')
    dst_image_path = os.path.join(output_image_dir, filename + '.png')
    shutil.copy(src_image_path, dst_image_path)

    # 复制标签文件
    src_label_path = os.path.join(label_dir, filename + '.txt')
    dst_label_path = os.path.join(output_label_dir, filename + '.txt')
    shutil.copy(src_label_path, dst_label_path)

将拆分出训练集/验证集/测试集.

训练

创建data.yml文件, 内容如下:

train: /home/user/moused/dataset/train  # train images (relative to 'path') 128 images
val: /home/user/moused/dataset/valid # val images (relative to 'path') 128 images
test: /home/user/moused/dataset/test # test images (optional)

# Classes
names:
  0: mouse

运行命令yolo detect train data=data.yml model=yolov8x.pt epochs=100 imgsz=640, 将开始训练.

预测

yolo detect predict model=runs/detect/train/weights/best.pt  source=mouse.mp4

效果还不错