1.制作数据集
1.1 下载安装labelImg
下载好源码并 cd到源码路径下
安装需要的环境
# 安装lxml
pip install lxml
# 安装pyqt5
pip install pyqt5
# 将"resources.qrc" 文件编译成 Python 代码,并将输出到 "libs/resources.py" 文件中
pyrcc5 -o libs/resources.py resources.qrc
# 启动labelImg
python labelimg.py
1.2 开始制作数据集
打开data/predefined_classes.txt,修改默认类别,也可以直接清空,后面打标签的时候会自动生成,我这里将这个文件清空。
点击“创建区块” 开始打标签
准备训练
# 创建一个用于训练的虚拟环境
conda create --name YOLOv8_ultralytics python=3.8
#启动训练环境
conda activate YOLOv8_ultralytics
PyTorch 要求因操作系统和 CUDA 要求而异,因此建议先按照 https://pytorch.org/get-started/locally 中的说明安装 PyTorch。
#我这里cuda是11.6
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.6 -c pytorch -c nvidia
#安装ultralytics
pip install ultralytics
分配数据集
-
新建文件夹 datasets/data/images
然后将图片和标签全部放到这个文件夹下。 -
在datasets下新建一个分配数据集数据的脚本文件Process.py
import os
import random
import shutil
def split_dataset(data_dir,train_val_test_dir, train_ratio, val_ratio, test_ratio):
# 创建目标文件夹
train_dir = os.path.join(train_val_test_dir, 'train')
val_dir = os.path.join(train_val_test_dir, 'val')
test_dir = os.path.join(train_val_test_dir, 'test')
os.makedirs(train_dir, exist_ok=True)
os.makedirs(val_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)
# 获取数据集中的所有文件
files = os.listdir(data_dir)
# 过滤掉非图片文件
image_files = [f for f in files if f.endswith('.jpg') or f.endswith('.png')]
# 随机打乱文件列表
random.shuffle(image_files)
# 计算切分数据集的索引
num_files = len(image_files)
num_train = int(num_files * train_ratio)
num_val = int(num_files * val_ratio)
num_test = num_files - num_train - num_val
# 分离训练集
train_files = image_files[:num_train]
for file in train_files:
src_image_path = os.path.join(data_dir, file)
src_label_path = os.path.join(data_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
dst_image_path = os.path.join(train_dir, file)
dst_label_path = os.path.join(train_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
shutil.copy(src_image_path, dst_image_path)
shutil.copy(src_label_path, dst_label_path)
# 分离验证集
val_files = image_files[num_train:num_train+num_val]
for file in val_files:
src_image_path = os.path.join(data_dir, file)
src_label_path = os.path.join(data_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
dst_image_path = os.path.join(val_dir, file)
dst_label_path = os.path.join(val_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
shutil.copy(src_image_path, dst_image_path)
shutil.copy(src_label_path, dst_label_path)
# 分离测试集
test_files = image_files[num_train+num_val:]
for file in test_files:
src_image_path = os.path.join(data_dir, file)
src_label_path = os.path.join(data_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
dst_image_path = os.path.join(test_dir, file)
dst_label_path = os.path.join(test_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
shutil.copy(src_image_path, dst_image_path)
shutil.copy(src_label_path, dst_label_path)
print("数据集分离完成!")
print(f"训练集数量:{len(train_files)}")
print(f"验证集数量:{len(val_files)}")
print(f"测试集数量:{len(test_files)}")
def move_files(data_dir):
# 创建目标文件夹
images_dir = os.path.join(data_dir, 'images')
labels_dir = os.path.join(data_dir, 'labels')
os.makedirs(images_dir, exist_ok=True)
os.makedirs(labels_dir, exist_ok=True)
# 获取数据集中的所有文件
files = os.listdir(data_dir)
# 移动PNG文件到images文件夹
png_files = [f for f in files if f.endswith('.png')]
for file in png_files:
src_path = os.path.join(data_dir, file)
dst_path = os.path.join(images_dir, file)
shutil.move(src_path, dst_path)
# 移动TXT文件到labels文件夹
txt_files = [f for f in files if f.endswith('.txt')]
for file in txt_files:
src_path = os.path.join(data_dir, file)
dst_path = os.path.join(labels_dir, file)
shutil.move(src_path, dst_path)
print(f"{data_dir}文件移动完成!")
print(f"总共移动了 {len(png_files)} 个PNG文件到images文件夹")
print(f"总共移动了 {len(txt_files)} 个TXT文件到labels文件夹")
# 设置数据集路径和切分比例
data_dir = './data/images' # 图片和标签路径
train_val_test_dir= './data' # 目标文件夹
train_ratio = 0.7 # 训练集比例
val_ratio = 0.2 # 验证集比例
test_ratio = 0.1 # 测试集比例
# 调用函数分离数据集
split_dataset(data_dir, train_val_test_dir,train_ratio, val_ratio, test_ratio)
# 调用函数移动文件
move_files(os.path.join(train_val_test_dir, 'train'))
move_files(os.path.join(train_val_test_dir, 'val'))
move_files(os.path.join(train_val_test_dir, 'test'))
- 在datasets下运行这个脚本,然后生成这样的文件结构
- 在datasets同级目录下新建一个mydata.yaml文件
train: ./data/train/images
val: ./data/val/images
test: ./data/test/images
# 类别数
nc: 4
# 类别名称
names: ['CCTV2','fivestar','CCTV3',"CCTV1"]
训练
在datasets同级路径下新建train.py 训练脚本
from ultralytics import YOLO
# Load a model
#model = YOLO('yolov8n.yaml') # build a new model from YAML
model = YOLO('yolov8n.pt') # load a pretrained model (recommended for training)
#model = YOLO('yolov8n.yaml').load('yolov8n.pt') # build from YAML and transfer weights
# Train the model
model.train(data='./mydata.yaml', epochs=100, imgsz=640)
测试
在datasets同级路径下新建test.py测试脚本
from PIL import Image
from ultralytics import YOLO
# Load a pretrained YOLOv8n model
model = YOLO('./runs/detect/train/weights/last.pt')
# Run inference on 'bus.jpg'
results = model('./test1.png') # results list
# Show the results
for r in results:
im_array = r.plot() # plot a BGR numpy array of predictions
im = Image.fromarray(im_array[..., ::-1]) # RGB PIL image
im.save('results.jpg') # save image