训练分割要准备好数据集和分割预训练权重文件
下面这张图是数据集的格式
下面这张图配置数据集，下面names 要和labelme转txt里配置的一样
下面这张图进行训练，配置一些全局参数 ，初始的yolov8s-seg.pt文件需要到github上yolov8开源项目里下
labelme转txt
# -*- coding: utf-8 -*-
import os
import numpy as np
import json
from glob import glob
import cv2
import shutil
import yaml
from sklearn.model_selection import train_test_split
from tqdm import tqdm
from PIL import Image
'''
辅助函数
'''
'''
转换，根据图像大小，返回box框的中点和高宽信息
'''
def convert(size, points):
# 转换后的返回列表
converted_points_list = []
# 宽
dw = 1. / (size[0])
# 高
dh = 1. / (size[1])
for point in points:
x = point[0] * dw
converted_points_list.append(x)
y = point[1] * dh
converted_points_list.append(y)
return converted_points_list
'''
# step1:统一图像格式
'''
def change_image_format(label_path, suffix='.jpg'):
"""
统一当前文件夹下所有图像的格式，如'.jpg'
:param suffix: 图像文件后缀
:param label_path:当前文件路径
:return:
"""
print("step1:统一图像格式")
# 修改尾缀列表
externs = ['png', 'jpg', 'JPEG', 'bmp']
files = list()
# 获取尾缀在externs中的所有图像
for extern in externs:
files.extend(glob(label_path + "\\*." + extern))
# 遍历所有图像，转换图像格式
for index,file in enumerate(tqdm(files)):
name = ''.join(file.split('.')[:-1])
file_suffix = file.split('.')[-1]
if file_suffix != suffix.split('.')[-1]:
# 重命名为jpg
new_name = name + suffix
# 读取图像
image = Image.open(file)
image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
# 重新存图为jpg格式
cv2.imwrite(new_name, image)
# 删除旧图像
os.remove(file)
'''
# step2:根据json文件划分训练集、验证集、测试集
'''
def split_dataset(ROOT_DIR, test_size=0.3, isUseTest=False, useNumpyShuffle=False):
"""
将文件分为训练集，测试集和验证集
:param useNumpyShuffle: 使用numpy方法分割数据集
:param test_size: 分割测试集或验证集的比例
:param isUseTest: 是否使用测试集，默认为False
:param label_path:当前文件路径
:return:
"""
# 获取所有json
print('step2:根据json文件划分训练集、验证集、测试集')
# 获取所有json文件路径
files = glob(ROOT_DIR + "\\*.json")
# 转换为json名称
files = [i.replace("\\", "/").split("/")[-1].split(".json")[0] for i in files]
# 是否打乱
if useNumpyShuffle:
file_length = len(files)
index = np.arange(file_length)
np.random.seed(32)
np.random.shuffle(index) # 随机划分
test_files = None
# 是否有测试集
if isUseTest:
trainval_files, test_files = np.array(files)[index[:int(file_length * (1 - test_size))]], np.array(files)[
index[int(file_length * (1 - test_size)):]]
else:
trainval_files = files
# 划分训练集和测试集
train_files, val_files = np.array(trainval_files)[index[:int(len(trainval_files) * (1 - test_size))]], \
np.array(trainval_files)[index[int(len(trainval_files) * (1 - test_size)):]]
else:
test_files = None
# 判断是否启用测试集
if isUseTest:
trainval_files, test_files = train_test_split(files, test_size=test_size, random_state=55)
else:
trainval_files = files
if len(trainval_files) != 0:
# 划分训练集和验证集
train_files, val_files = train_test_split(trainval_files, test_size=test_size, random_state=55)
else:
print("数据文件为空")
return train_files, val_files, test_files, files
'''
# step3:根据json文件，获取所有类别
'''
def get_all_class(file_list, ROOT_DIR):
"""
从json文件中获取当前数据的所有类别
:param file_list:当前路径下的所有文件名
:param label_path:当前文件路径
:return:
"""
print('step3:根据json文件，获取所有类别')
# 初始化类别列表
classes = list()
# 遍历所有json,读取shape中的label值内容，添加到classes
for filename in tqdm(file_list):
json_path = os.path.join(ROOT_DIR, filename + '.json')
json_file = json.load(open(json_path, "r", encoding="utf-8"))
for item in json_file["shapes"]:
label_class = item['label']
# 如果没有添加新的类型，则
if label_class not in classes:
classes.append(label_class)
print('read file done')
return classes
'''
# step4:将json文件转化为txt文件，并将json文件存放到指定文件夹
'''
def json2txt(classes, txt_Name='allfiles', ROOT_DIR="", suffix='.jpg'):
"""
将json文件转化为txt文件，并将json文件存放到指定文件夹
:param classes: 类别名
:param txt_Name:txt文件，用来存放所有文件的路径
:param label_path:当前文件路径
:param suffix:图像文件后缀
:return:
"""
print('step4:将json文件转化为txt文件，并将json文件存放到指定文件夹')
store_json = os.path.join(ROOT_DIR, 'json')
if not os.path.exists(store_json):
os.makedirs(store_json)
_, _, _, files = split_dataset(ROOT_DIR)
if not os.path.exists(os.path.join(ROOT_DIR, 'tmp')):
os.makedirs(os.path.join(ROOT_DIR, 'tmp'))
list_file = open(os.path.join(ROOT_DIR,'tmp/%s.txt'% txt_Name) , 'w')
for json_file_ in tqdm(files):
# json路径
json_filename = os.path.join(ROOT_DIR, json_file_ + ".json")
# 图像路径
imagePath = os.path.join(ROOT_DIR, json_file_ + suffix)
# 写入图像文件夹路径
list_file.write('%s\n' % imagePath)
# 转换后txt标签文件夹路径
out_file = open('%s/%s.txt' % (ROOT_DIR, json_file_), 'w')
# 加载标签json文件
json_file = json.load(open(json_filename, "r", encoding="utf-8"))
'''
核心：标签转换（json转txt）
'''
if os.path.exists(imagePath):
# 读取图像
image = Image.open(imagePath)
image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
# 获取图像高、宽、通道
height, width, channels = image.shape
# 获取shapes的Value值
for multi in json_file["shapes"]:
# 如果点位为空
if len(multi["points"][0]) == 0:
out_file.write('')
continue
# 获取单个缺陷的点位（x,y的list）
points = np.array(multi["points"])
# 标签
label = multi["label"]
# 类别id
cls_id = classes.index(label)
# 根据图像大小，返回box框的中点和高宽信息
xy_list = convert((width, height), points)
# 写txt标签文件
out_file.write(str(cls_id) + " " + " ".join([str(xy) for xy in xy_list]) + '\n')
if not os.path.exists(os.path.join(store_json, json_file_ + '.json')):
try:
shutil.move(json_filename, store_json)
except OSError:
pass
'''
# step5:创建yolov5训练所需的yaml文件
'''
def create_yaml(classes, ROOT_DIR, isUseTest=False,dataYamlName=""):
print('step5:创建yolov5训练所需的yaml文件')
classes_dict = {}
for index, item in enumerate(classes):
classes_dict[index] = item
if not isUseTest:
desired_caps = {
'path': ROOT_DIR,
'train': 'images/train',
'val': 'images/val',
'names': classes_dict
}
else:
desired_caps = {
'path': ROOT_DIR,
'train': 'images/train',
'val': 'images/val',
'test': 'images/test',
'names': classes_dict
}
yamlpath = os.path.join(ROOT_DIR, dataYamlName + ".yaml")
# 写入到yaml文件
with open(yamlpath, "w+", encoding="utf-8") as f:
for key, val in desired_caps.items():
yaml.dump({key: val}, f, default_flow_style=False)
'''
# step6:生成yolov5的训练、验证、测试集的文件夹
'''
def create_save_file(ROOT_DIR):
"""
按照训练时的图像和标注路径创建文件夹
:param label_path:当前文件路径
:return:
"""
print('step6:生成yolov5的训练、验证、测试集的文件夹')
# 生成训练集
train_image = os.path.join(ROOT_DIR, 'images','train')
if not os.path.exists(train_image):
os.makedirs(train_image)
train_label = os.path.join(ROOT_DIR, 'labels','train')
if not os.path.exists(train_label):
os.makedirs(train_label)
# 生成验证集
val_image = os.path.join(ROOT_DIR, 'images', 'val')
if not os.path.exists(val_image):
os.makedirs(val_image)
val_label = os.path.join(ROOT_DIR, 'labels', 'val')
if not os.path.exists(val_label):
os.makedirs(val_label)
# 生成测试集
test_image = os.path.join(ROOT_DIR, 'images', 'test')
if not os.path.exists(test_image):
os.makedirs(test_image)
test_label = os.path.join(ROOT_DIR, 'labels', 'test')
if not os.path.exists(test_label):
os.makedirs(test_label)
return train_image, train_label, val_image, val_label, test_image, test_label
'''
# step7:将所有图像和标注文件，移动到对应的训练集、验证集、测试集
'''
def push_into_file(file, images, labels, ROOT_DIR, suffix='.jpg'):
"""
最终生成在当前文件夹下的所有文件按image和label分别存在到训练集/验证集/测试集路径的文件夹下
:param file: 文件名列表
:param images: 存放images的路径
:param labels: 存放labels的路径
:param label_path: 当前文件路径
:param suffix: 图像文件后缀
:return:
"""
print('step7:将所有图像和标注文件，移动到对应的训练集、验证集、测试集')
# 遍历所有文件
for filename in tqdm(file):
# 图像文件
image_file = os.path.join(ROOT_DIR, filename + suffix)
# 标注文件
label_file = os.path.join(ROOT_DIR, filename + '.txt')
# yolov5存放图像文件夹
if not os.path.exists(os.path.join(images, filename + suffix)):
try:
shutil.move(image_file, images)
except OSError:
pass
# yolov5存放标注文件夹
if not os.path.exists(os.path.join(labels, filename + suffix)):
try:
shutil.move(label_file, labels)
except OSError:
pass
'''
labelme的json标签转yolo的txt标签
'''
def ChangeToYolo(ROOT_DIR="", suffix='.bmp',classes="", test_size=0.1, isUseTest=False,useNumpyShuffle=False,auto_genClasses = False,dataYamlName=""):
"""
生成最终标准格式的文件
:param test_size: 分割测试集或验证集的比例
:param label_path:当前文件路径
:param suffix: 文件后缀名
:param isUseTest: 是否使用测试集
:return:
"""
# step1:统一图像格式
change_image_format(ROOT_DIR, suffix=suffix)
# step2:根据json文件划分训练集、验证集、测试集
train_files, val_files, test_file, files = split_dataset(ROOT_DIR, test_size=test_size, isUseTest=isUseTest, useNumpyShuffle=useNumpyShuffle)
# step3:根据json文件，获取所有类别
classes = classes
# 是否自动从数据集中获取类别数
if auto_genClasses:
classes = get_all_class(files, ROOT_DIR)
'''
step4:（***核心***）将json文件转化为txt文件，并将json文件存放到指定文件夹
'''
json2txt(classes, txt_Name='allfiles', ROOT_DIR=ROOT_DIR, suffix=suffix)
# step5:创建yolov5训练所需的yaml文件
create_yaml(classes, ROOT_DIR, isUseTest=isUseTest,dataYamlName=dataYamlName)
# step6:生成yolov5的训练、验证、测试集的文件夹
train_image_dir, train_label_dir, val_image_dir, val_label_dir, test_image_dir, test_label_dir = create_save_file(ROOT_DIR)
# step7:将所有图像和标注文件，移动到对应的训练集、验证集、测试集file, images, labels, ROOT_DIR, suffix='.jpg'
# 将文件移动到训练集文件中
push_into_file(train_files, train_image_dir, train_label_dir,ROOT_DIR=ROOT_DIR, suffix=suffix)
# 将文件移动到验证集文件夹中
push_into_file(val_files, val_image_dir, val_label_dir,ROOT_DIR=ROOT_DIR, suffix=suffix)
# 如果测试集存在，则将文件移动到测试集文件中
if test_file is not None:
push_into_file(test_file, test_image_dir, test_label_dir, ROOT_DIR=ROOT_DIR, suffix=suffix)
print('create dataset done')
if __name__ == "__main__":
'''
1.ROOT_DIR:图像和json标签的路径
2.suffix：统一图像尾缀
3.classes=['dog', 'cat'], # 输入你标注的列表里的名称（注意区分大小写）,用于自定义类别名称和id对应
4.test_size：测试集和验证集所占比例
5.isUseTest：是否启用测试集
6.useNumpyShuffle：是否随机打乱
7.auto_genClasses：是否自动根据json标签生成类别列表
8.dataYamlName：数据集yaml文件名称
'''
ChangeToYolo(
ROOT_DIR = r'G:\down\labels', # 注意数据集路径不要含中文
suffix='.jpg', # 确定图像尾缀，用于统一图像尾缀
classes=['Arcella'], # 输入你标注的列表里的名称（注意区分大小写）
test_size=0.1, # 测试集大小占比
isUseTest=False, # 是否启用测试集
useNumpyShuffle=True, # 是否乱序
auto_genClasses = False, # 是否根据数据集自动生成类别id
dataYamlName= "catAndDog_data" # 数据集yaml文件名称
)
下面这个也可以将labelme转txt，不过转之后，自己在拆分
备注 需要在labelme文件夹里放一个classes.txt
import os
import json
import base64
import cv2
import numpy as np
def draw_polygon(image, points, color=(0, 255, 0), thickness=2):
"""
在给定图像上绘制一个多边形。
:param image: 待绘制多边形的图像（numpy数组）
:param points: 多边形顶点坐标列表
:param color: 多边形的颜色 (B, G, R)
:param thickness: 边缘线条厚度
"""
# 绘制多边形轮廓
cv2.polylines(image, [np.int32(points)], isClosed=True, color=color, thickness=thickness)
# 绘制多边形内部
cv2.fillPoly(image, [np.int32(points)], color=color)
def concat_images(image1, image2):
"""
拼接两个图像为一个垂直堆叠的图像。
:param image1: 第一张图像（numpy数组）
:param image2: 第二张图像（numpy数组）
:return: 垂直拼接后的图像（numpy数组）
"""
# 确保两个图像的宽度相同
max_height = max(image1.shape[0], image2.shape[0])
max_width = max(image1.shape[1], image2.shape[1])
# 调整图像大小以匹配最大宽度
if image1.shape[1] < max_width:
image1 = cv2.copyMakeBorder(image1, 0, 0, 0, max_width - image1.shape[1], cv2.BORDER_CONSTANT, value=[255, 255, 255])
if image2.shape[1] < max_width:
image2 = cv2.copyMakeBorder(image2, 0, 0, 0, max_width - image2.shape[1], cv2.BORDER_CONSTANT, value=[255, 255, 255])
# 如果需要，调整图像高度以匹配最大高度
if image1.shape[0] < max_height:
image1 = cv2.copyMakeBorder(image1, 0, max_height - image1.shape[0], 0, 0, cv2.BORDER_CONSTANT, value=[255, 255, 255])
if image2.shape[0] < max_height:
image2 = cv2.copyMakeBorder(image2, 0, max_height - image2.shape[0], 0, 0, cv2.BORDER_CONSTANT, value=[255, 255, 255])
# 拼接图像
return np.vstack((image1, image2))
def convert_labelme_to_yolov8(json_dir):
"""
将 LabelMe 格式的标注转换为 YOLOv8 格式，并绘制多边形到图像上。
:param json_dir: 包含 LabelMe JSON 文件的目录路径
"""
# 生成颜色列表
color_list = [
(0, 0, 255), # Red
(0, 255, 0), # Green
(255, 0, 0), # Blue
(0, 255, 255), # Yellow
(255, 255, 0), # Cyan
(255, 0, 255), # Magenta
(0, 165, 255), # Orange
(203, 192, 255), # Pink
(42, 42, 165), # Brown
(0, 128, 128), # Olive
(128, 128, 0), # Teal
(238, 130, 238), # Violet
(128, 128, 128), # Gray
(192, 192, 192), # Silver
(0, 0, 128), # Maroon
(128, 0, 128), # Purple
(0, 0, 128), # Navy
(0, 255, 0), # Lime
(0, 255, 255), # Aqua
(255, 0, 255), # Fuchsia
(255, 255, 255), # White
(0, 0, 0), # Black
(235, 206, 135), # Light Blue
(144, 238, 144), # Light Green
(193, 182, 255), # Light Pink
(224, 255, 255), # Light Yellow
(216, 191, 216), # Light Purple
(0, 128, 128), # Light Olive
(30, 105, 210), # Light Brown
(211, 211, 211) # Light Gray
]
# 加载类别文件
classes_file = os.path.join(json_dir, 'classes.txt')
if not os.path.exists(classes_file):
print("Error: Could not find 'classes.txt' in the specified directory.")
exit(1)
with open(classes_file, 'r') as f:
class_names = [line.strip() for line in f.readlines()]
# 获取 JSON 文件列表
json_files = [f for f in os.listdir(json_dir) if f.endswith('.json')]
for json_file in json_files:
json_file_path = os.path.join(json_dir, json_file)
# 输出文件名
output_file_name = json_file.replace('.json', '.txt')
output_file_path = os.path.join(json_dir, output_file_name)
# 读取 JSON 文件
with open(json_file_path, 'r') as f:
data = json.load(f)
image_width = data['imageWidth']
image_height = data['imageHeight']
# 解码图像数据
imageData = data.get('imageData')
if imageData is not None:
image_data = base64.b64decode(imageData)
image_np = np.frombuffer(image_data, dtype=np.uint8)
image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
# 创建一个副本用于绘制标注
annotated_image = image.copy()
# 绘制标注
for i, shape in enumerate(data['shapes']):
if shape['shape_type'] == 'polygon':
points = np.array(shape['points'], dtype=np.int32)
label = shape['label']
class_index = class_names.index(label)
color = color_list[class_index % len(color_list)]
draw_polygon(annotated_image, points, color=color)
# 保存原始图像
image_file_name = json_file.replace('.json', '.jpg')
image_file_path = os.path.join(json_dir, image_file_name)
cv2.imwrite(image_file_path, image)
# 将原始图像和带有标注的图像上下拼接
concatenated_image = concat_images(image, annotated_image)
# 保存拼接后的图像
concatenated_image_file_name = json_file.replace('.json', '_check.jpg')
concatenated_image_file_path = os.path.join(json_dir, concatenated_image_file_name)
cv2.imwrite(concatenated_image_file_path, concatenated_image)
# 显示带有标注的图像
cv2.imshow('Annotated Image', concatenated_image)
cv2.waitKey(0)
cv2.destroyAllWindows()
# 开始写入 YOLOv8 格式的文本文件
with open(output_file_path, 'w') as f:
for shape in data['shapes']:
if shape['shape_type'] == 'polygon':
label = shape['label']
points = shape['points']
# 获取类别索引
try:
class_index = class_names.index(label)
except ValueError:
print(f"Warning: Label '{label}' not found in 'classes.txt'. Skipping this label.")
continue
# 归一化坐标
normalized_points = []
for point in points:
x = point[0] / image_width
y = point[1] / image_height
normalized_points.extend([x, y])
# 写入 YOLOv8 格式的行
line = f"{class_index} {' '.join(map(str, normalized_points))}\n"
f.write(line)
if __name__ == '__main__':
import sys
# 从命令行参数获取 JSON 目录
if len(sys.argv) != 2:
print("Usage: python script.py /path/to/json/files")
exit(1)
json_dir = sys.argv[1]
convert_labelme_to_yolov8(json_dir)
FR：徐海涛(hunkxu)