| 知乎专栏 |
初始化环境
neo@Netkiller-Mac-mini-M4 netkiller % python3 -m venv .venv
我使用的是 fish shell
neo@Netkiller-Mac-mini-M4 netkiller % fish Welcome to fish, the friendly interactive shell Type help for instructions on how to use fish
切换到 venv 环境
neo@Netkiller-Mac-mini-M4 ~/P/netkiller (yolo)> source .venv/bin/activate.fish (.venv) neo@Neo-Mac-mini-M4 ~/P/netkiller (yolo)>
这里我们使用 captcha 随机生成验证码
(.venv) neo@Neo-Mac-mini-M4 ~/P/netkiller (yolo)> pip install captcha
为了降低学习难度，我们生成4为纯数字验证码，其中 50 张用于训练，10张用于测试
import os
import random
from captcha.image import ImageCaptcha
# captcha_character = list("0123456789abcdefghijklmnopqrstuvwxyz")
captcha_character = list("0123456789")
captcha_length = 4
if __name__ == "__main__":
image = ImageCaptcha()
for i in range(50):
code = "".join(random.sample(captcha_character, captcha_length))
path = "./captcha/images/{}_{}.png".format(i,code)
print(code)
image.write(code, path)
for i in range(10):
code = "".join(random.sample(captcha_character, captcha_length))
path = "./captcha/test/{}_{}.png".format(i, code)
print(code)
image.write(code, path)
安装图像标注工具
pip install labelme labelme2yolo
启动 labelme，然后进入漫长苦哈哈的标注工作
(.venv) neo@Neo-Mac-mini-M4 ~/P/netkiller (yolo)> labelme
标注完成之后，使用 labelme2yolo 工具将 json 转成 yolo 格式
(.venv) neo@Neo-Mac-mini netkiller % labelme2yolo --json_dir captcha/images [2024-11-19T07:06:52Z INFO labelme2yolo] Starting the conversion process... [2024-11-19T07:06:52Z INFO labelme2yolo] Read and parsed 20 JSON files. [Train] [00:00:00] [########################################] 16/16 (0s) [Val] [00:00:00] [########################################] 4/4 (0s) [2024-11-19T07:06:52Z INFO labelme2yolo] Creating dataset.yaml file... [2024-11-19T07:06:52Z INFO labelme2yolo] Conversion process completed successfully. (.venv) neo@Neo-Mac-mini netkiller % ls captcha/images/YOLODataset dataset.yaml images labels
将 YOLODataset 复制 datasets 目录
rm -rf datasets/YOLODataset mv captcha/images/YOLODataset datasets
修改 captcha/images/YOLODataset/dataset.yaml 文件，确认 path 路径正确
path: /Users/neo/PycharmProjects/netkiller/datasets/YOLODataset
train: images/train
val: images/val
test:
names:
0: 8
1: 2
2: 4
3: 9
4: 7
5: 3
6: 0
7: 6
8: 5
9: 1
新建一个模型，参考 https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/models/11/yolo11.yaml
# Ultralytics YOLO 🚀, AGPL-3.0 license # YOLO11 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect # Parameters nc: 10 # number of classes scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n' # [depth, width, max_channels] n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs s: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs m: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs l: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs x: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs # YOLO11n backbone backbone: # [from, repeats, module, args] - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 - [-1, 2, C3k2, [256, False, 0.25]] - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 - [-1, 2, C3k2, [512, False, 0.25]] - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 - [-1, 2, C3k2, [512, True]] - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 - [-1, 2, C3k2, [1024, True]] - [-1, 1, SPPF, [1024, 5]] # 9 - [-1, 2, C2PSA, [1024]] # 10 # YOLO11n head head: - [-1, 1, nn.Upsample, [None, 2, "nearest"]] - [[-1, 6], 1, Concat, [1]] # cat backbone P4 - [-1, 2, C3k2, [512, False]] # 13 - [-1, 1, nn.Upsample, [None, 2, "nearest"]] - [[-1, 4], 1, Concat, [1]] # cat backbone P3 - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small) - [-1, 1, Conv, [256, 3, 2]] - [[-1, 13], 1, Concat, [1]] # cat head P4 - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium) - [-1, 1, Conv, [512, 3, 2]] - [[-1, 10], 1, Concat, [1]] # cat head P5 - [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large) - [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)
修改 nc: 10 # number of classes 我们训练的纯数字验证码，只有 0～9，所以 nc 是 10 个分类
yolo task=detect mode=train model=captcha.yaml data=datasets/YOLODataset/dataset.yaml epochs=100 workers=0 batch=10 imgsz=320
yolo val model=runs/detect/train/weights/best.pt data=datasets/YOLODataset/dataset.yaml
neo@Netkiller-Mac-mini-M4 ~/P/netkiller (yolo)> yolo val model=runs/detect/train/weights/best.pt data=datasets/YOLODataset/dataset.yaml
Ultralytics 8.3.34 🚀 Python-3.12.7 torch-2.5.1 CPU (Apple M4)
captcha summary (fused): 238 layers, 2,584,102 parameters, 0 gradients, 6.3 GFLOPs
val: Scanning /Users/neo/PycharmProjects/netkiller/datasets/YOLODataset/labels/val.cache... 20 images, 0 backgrounds, 0 corrupt: 100%|██████████| 20/20 [00:00<?, ?it/s]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 2/2 [00:00<00:00, 4.17it/s]
all 20 80 0.406 0.538 0.532 0.366
8 9 9 0.408 0.889 0.668 0.495
2 7 7 0.614 0.571 0.627 0.429
4 9 9 0.415 0.222 0.342 0.199
9 6 6 0.337 0.667 0.538 0.348
7 9 9 0.102 0.111 0.224 0.142
3 8 8 0.452 0.315 0.541 0.406
0 9 9 0.391 0.778 0.736 0.593
6 10 10 0.454 0.6 0.549 0.384
1 13 13 0.484 0.692 0.562 0.293
Speed: 0.2ms preprocess, 21.8ms inference, 0.0ms loss, 0.6ms postprocess per image
Results saved to /Users/neo/PycharmProjects/netkiller/runs/detect/val
💡 Learn more at https://docs.ultralytics.com/modes/val
我们训练好的模型 runs/detect/train/weights/best.pt，预测的图片 captcha/test/0_5902.png
yolo predict model=runs/detect/train/weights/best.pt source=captcha/test/0_5902.png imgsz=320
from ultralytics import YOLO
# 加载模型
model = YOLO('runs/detect/train/weights/best.pt')
# 进行预测
results = model.predict('captcha/test/0_5902.png')
# 提取检测结果
for result in results:
boxes = result.boxes.xyxy # 边界框坐标
scores = result.boxes.conf # 置信度分数
classes = result.boxes.cls # 类别索引
# 如果有类别名称，可以通过类别索引获取
class_names = [model.names[int(cls)] for cls in classes]
# 打印检测结果
for box, score, class_name in zip(boxes, scores, class_names):
print(f"Class: {class_name}, Score: {score:.2f}, Box: {box}")
# 可视化检测结果
annotated_img = result.plot()
# 显示图像
result.show()
# -*- coding: utf-8 -*-
from tqdm import tqdm
import shutil
import random
import os
import argparse
from collections import Counter
import yaml
import json
def mkdir(path):
if not os.path.exists(path):
os.makedirs(path)
def convert_label_json(json_dir, save_dir, classes):
# json_paths = os.listdir(json_dir)
suffix = ".json"
json_files = [file for file in os.listdir(json_dir) if file.endswith(suffix)]
# print(json_files)
classes = classes.split(',')
mkdir(save_dir)
for json_path in tqdm(json_files):
# print(json_path)
# for json_path in json_paths:
path = os.path.join(json_dir, json_path)
with open(path, 'r') as load_f:
json_dict = json.load(load_f)
h, w = json_dict['imageHeight'], json_dict['imageWidth']
# save txt path
txt_path = os.path.join(save_dir, json_path.replace('json', 'txt'))
txt_file = open(txt_path, 'w')
for shape_dict in json_dict['shapes']:
label = shape_dict['label']
label_index = classes.index(label)
points = shape_dict['points']
points_nor_list = []
for point in points:
points_nor_list.append(point[0] / w)
points_nor_list.append(point[1] / h)
points_nor_list = list(map(lambda x: str(x), points_nor_list))
points_nor_str = ' '.join(points_nor_list)
label_str = str(label_index) + ' ' + points_nor_str + '\n'
txt_file.writelines(label_str)
def get_classes(json_dir):
'''
统计路径下 JSON 文件里的各类别标签数量
'''
names = []
json_files = [os.path.join(json_dir, f) for f in os.listdir(json_dir) if f.endswith('.json')]
for json_path in json_files:
with open(json_path, 'r') as f:
data = json.load(f)
for shape in data['shapes']:
name = shape['label']
names.append(name)
result = Counter(names)
return result
def main(image_dir, json_dir, txt_dir, save_dir):
# 创建文件夹
mkdir(save_dir)
images_dir = os.path.join(save_dir, 'images')
labels_dir = os.path.join(save_dir, 'labels')
img_train_path = os.path.join(images_dir, 'train')
img_val_path = os.path.join(images_dir, 'val')
label_train_path = os.path.join(labels_dir, 'train')
label_val_path = os.path.join(labels_dir, 'val')
mkdir(images_dir)
mkdir(labels_dir)
mkdir(img_train_path)
mkdir(img_val_path)
mkdir(label_train_path)
mkdir(label_val_path)
# 数据集划分比例，训练集75%，验证集15%，测试集15%，按需修改
train_percent = 0.90
val_percent = 0.10
total_txt = os.listdir(txt_dir)
num_txt = len(total_txt)
list_all_txt = range(num_txt) # 范围 range(0, num)
num_train = int(num_txt * train_percent)
num_val = int(num_txt * val_percent)
train = random.sample(list_all_txt, num_train)
# 在全部数据集中取出train
val = [i for i in list_all_txt if not i in train]
# 再从val_test取出num_val个元素，val_test剩下的元素就是test
# val = random.sample(list_all_txt, num_val)
print("训练集数目：{}, 验证集数目：{}".format(len(train), len(val)))
for i in list_all_txt:
name = total_txt[i][:-4]
srcImage = os.path.join(image_dir, name + '.png')
srcLabel = os.path.join(txt_dir, name + '.txt')
if i in train:
dst_train_Image = os.path.join(img_train_path, name + '.png')
dst_train_Label = os.path.join(label_train_path, name + '.txt')
shutil.copyfile(srcImage, dst_train_Image)
shutil.copyfile(srcLabel, dst_train_Label)
elif i in val:
dst_val_Image = os.path.join(img_val_path, name + '.png')
dst_val_Label = os.path.join(label_val_path, name + '.txt')
shutil.copyfile(srcImage, dst_val_Image)
shutil.copyfile(srcLabel, dst_val_Label)
obj_classes = get_classes(json_dir)
classes = list(obj_classes.keys())
# 编写yaml文件
classes_txt = {i: classes[i] for i in range(len(classes))} # 标签类别
data = {
'path': os.path.join(os.getcwd(), save_dir),
'train': "images/train",
'val': "images/val",
'names': classes_txt,
'nc': len(classes)
}
with open(save_dir + '/segment.yaml', 'w', encoding="utf-8") as file:
yaml.dump(data, file, allow_unicode=True)
print("标签：", dict(obj_classes))
if __name__ == "__main__":
"""
python json2txt_nomalize.py --json-dir my_datasets/color_rings/jsons --save-dir my_datasets/color_rings/txts --classes "cat,dogs"
"""
classes_list = '0,1,2,3,4,5,6,7,8,9' # 类名
parser = argparse.ArgumentParser(description='json convert to txt params')
parser.add_argument('--image-dir', type=str, default='', help='图片地址')
parser.add_argument('--json-dir', type=str, default='', help='json地址')
parser.add_argument('--txt-dir', type=str, default='', help='保存txt文件地址')
parser.add_argument('--save-dir', default='', type=str, help='保存最终分割好的数据集地址')
parser.add_argument('--classes', type=str, default=classes_list, help='classes')
args = parser.parse_args()
# print(args)
# parser.print_help()
if 'image_dir' not in args:
parser.print_help()
exit(128)
else:
image_dir = args.image_dir
if 'json_dir' not in args:
parser.print_help()
exit(128)
else:
json_dir = args.json_dir
if 'txt_dir' not in args:
parser.print_help()
exit(128)
else:
txt_dir = args.txt_dir
if 'save_dir' not in args:
parser.print_help()
exit(128)
else:
save_dir = args.save_dir
classes = args.classes
# json格式转txt格式
convert_label_json(json_dir, txt_dir, classes)
# 划分数据集，生成yaml训练文件
main(image_dir, json_dir, txt_dir, save_dir)