File size: 11,416 Bytes
4a6e43e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 |
# -*- 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
# 获取当前路径
ROOT_DIR = os.getcwd()
'''
统一图像格式
'''
def change_image_format(label_path=ROOT_DIR, suffix='.png'):
"""
统一当前文件夹下所有图像的格式,如'.jpg'
:param suffix: 图像文件后缀
:param label_path:当前文件路径
:return:
"""
externs = ['png', 'jpg', 'JPEG', 'BMP', 'bmp']
files = list()
# 获取尾缀在ecterns中的所有图像
for extern in externs:
files.extend(glob(label_path + "\\*." + extern))
# 遍历所有图像,转换图像格式
for file in files:
name = ''.join(file.split('.')[:-1])
file_suffix = file.split('.')[-1]
if file_suffix != suffix.split('.')[-1]:
# 重命名为jpg
new_name = name + suffix
# 读取图像
image = cv2.imread(file)
# 重新存图为jpg格式
cv2.imwrite(new_name, image)
# 删除旧图像
os.remove(file)
'''
读取所有json文件,获取所有的类别
'''
def get_all_class(file_list, label_path=ROOT_DIR):
"""
从json文件中获取当前数据的所有类别
:param file_list:当前路径下的所有文件名
:param label_path:当前文件路径
:return:
"""
# 初始化类别列表
classes = list()
# 遍历所有json,读取shape中的label值内容,添加到classes
for filename in tqdm(file_list):
json_path = os.path.join(label_path, 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
'''
划分训练集、验证机、测试集
'''
def split_dataset(label_path, test_size=0.3, isUseTest=False, useNumpyShuffle=False):
"""
将文件分为训练集,测试集和验证集
:param useNumpyShuffle: 使用numpy方法分割数据集
:param test_size: 分割测试集或验证集的比例
:param isUseTest: 是否使用测试集,默认为False
:param label_path:当前文件路径
:return:
"""
# 获取所有json
files = glob(label_path + "\\*.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
train_files, val_files = train_test_split(trainval_files, test_size=test_size, random_state=55)
return train_files, val_files, test_files, files
'''
生成yolov5的训练、验证、测试集的文件夹
'''
def create_save_file(label_path=ROOT_DIR):
"""
按照训练时的图像和标注路径创建文件夹
:param label_path:当前文件路径
:return:
"""
# 生成训练集
train_image = os.path.join(label_path, 'train', 'images')
if not os.path.exists(train_image):
os.makedirs(train_image)
train_label = os.path.join(label_path, 'train', 'labels')
if not os.path.exists(train_label):
os.makedirs(train_label)
# 生成验证集
val_image = os.path.join(label_path, 'valid', 'images')
if not os.path.exists(val_image):
os.makedirs(val_image)
val_label = os.path.join(label_path, 'valid', 'labels')
if not os.path.exists(val_label):
os.makedirs(val_label)
# 生成测试集
test_image = os.path.join(label_path, 'test', 'images')
if not os.path.exists(test_image):
os.makedirs(test_image)
test_label = os.path.join(label_path, 'test', 'labels')
if not os.path.exists(test_label):
os.makedirs(test_label)
return train_image, train_label, val_image, val_label, test_image, test_label
'''
转换,根据图像大小,返回box框的中点和高宽信息
'''
def convert(size, box):
# 宽
dw = 1. / (size[0])
# 高
dh = 1. / (size[1])
x = (box[0] + box[1]) / 2.0 - 1
y = (box[2] + box[3]) / 2.0 - 1
# 宽
w = box[1] - box[0]
# 高
h = box[3] - box[2]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return x, y, w, h
'''
移动图像和标注文件到指定的训练集、验证集和测试集中
'''
def push_into_file(file, images, labels, label_path=ROOT_DIR, suffix='.jpg'):
"""
最终生成在当前文件夹下的所有文件按image和label分别存在到训练集/验证集/测试集路径的文件夹下
:param file: 文件名列表
:param images: 存放images的路径
:param labels: 存放labels的路径
:param label_path: 当前文件路径
:param suffix: 图像文件后缀
:return:
"""
# 遍历所有文件
for filename in file:
# 图像文件
image_file = os.path.join(label_path, filename + suffix)
# 标注文件
label_file = os.path.join(label_path, 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
'''
'''
def json2txt(classes, txt_Name='allfiles', label_path=ROOT_DIR, suffix='.png'):
"""
将json文件转化为txt文件,并将json文件存放到指定文件夹
:param classes: 类别名
:param txt_Name:txt文件,用来存放所有文件的路径
:param label_path:当前文件路径
:param suffix:图像文件后缀
:return:
"""
store_json = os.path.join(label_path, 'json')
if not os.path.exists(store_json):
os.makedirs(store_json)
_, _, _, files = split_dataset(label_path)
if not os.path.exists(os.path.join(label_path, 'tmp')):
os.makedirs(os.path.join(label_path, 'tmp'))
list_file = open('tmp/%s.txt' % txt_Name, 'w')
for json_file_ in tqdm(files):
json_filename = os.path.join(label_path, json_file_ + ".json")
imagePath = os.path.join(label_path, json_file_ + suffix)
list_file.write('%s\n' % imagePath)
out_file = open('%s/%s.txt' % (label_path, json_file_), 'w')
json_file = json.load(open(json_filename, "r", encoding="utf-8"))
if os.path.exists(imagePath):
height, width, channels = cv2.imread(imagePath).shape
for multi in json_file["shapes"]:
if len(multi["points"][0]) == 0:
out_file.write('')
continue
points = np.array(multi["points"])
xmin = min(points[:, 0]) if min(points[:, 0]) > 0 else 0
xmax = max(points[:, 0]) if max(points[:, 0]) > 0 else 0
ymin = min(points[:, 1]) if min(points[:, 1]) > 0 else 0
ymax = max(points[:, 1]) if max(points[:, 1]) > 0 else 0
label = multi["label"]
if xmax <= xmin:
pass
elif ymax <= ymin:
pass
else:
cls_id = classes.index(label)
b = (float(xmin), float(xmax), float(ymin), float(ymax))
bb = convert((width, height), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
# print(json_filename, xmin, ymin, xmax, ymax, cls_id)
if not os.path.exists(os.path.join(store_json, json_file_ + '.json')):
try:
shutil.move(json_filename, store_json)
except OSError:
pass
'''
创建yaml文件
'''
def create_yaml(classes, label_path, isUseTest=False):
nc = len(classes)
if not isUseTest:
desired_caps = {
'path': label_path,
'train': 'train/images',
'val': 'valid/images',
'nc': nc,
'names': classes
}
else:
desired_caps = {
'path': label_path,
'train': 'train/images',
'val': 'valid/images',
'test': 'test/images',
'nc': nc,
'names': classes
}
yamlpath = os.path.join(label_path, "data" + ".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)
# 首先确保当前文件夹下的所有图片统一后缀,如.jpg,如果为其他后缀,将suffix改为对应的后缀,如.png
def ChangeToYolo5(label_path=r"D:\storydata", suffix='.png', test_size=0.1, isUseTest=False):
"""
生成最终标准格式的文件
:param test_size: 分割测试集或验证集的比例
:param label_path:当前文件路径
:param suffix: 文件后缀名
:param isUseTest: 是否使用测试集
:return:
"""
# step1:统一图像格式
change_image_format(label_path)
# step2:根据json文件划分训练集、验证集、测试集
train_files, val_files, test_file, files = split_dataset(label_path, test_size=test_size, isUseTest=isUseTest)
# step3:根据json文件,获取所有类别
classes = get_all_class(files)
# step4:将json文件转化为txt文件,并将json文件存放到指定文件夹
json2txt(classes)
# step5:创建yolov5训练所需的yaml文件
create_yaml(classes, label_path, isUseTest=isUseTest)
# step6:生成yolov5的训练、验证、测试集的文件夹
train_image, train_label, val_image, val_label, test_image, test_label = create_save_file(label_path)
# step7:将所有图像和标注文件,移动到对应的训练集、验证集、测试集
push_into_file(train_files, train_image, train_label, suffix=suffix) # 将文件移动到训练集文件中
push_into_file(val_files, val_image, val_label, suffix=suffix) # 将文件移动到验证集文件夹中
if test_file is not None: # 如果测试集存在,则将文件移动到测试集文件中
push_into_file(test_file, test_image, test_label, suffix=suffix)
print('create dataset done')
if __name__ == "__main__":
ChangeToYolo5()
|