深度学习常见数据集处理方法

1、数据集格式转换（json转txt）

import json
import os'''
任务：实例分割，labelme的json文件, 转txt文件
Ultralytics YOLO format
<class-index> <x1> <y1> <x2> <y2> ... <xn> <yn>
'''# 类别映射表，定义每个类别对应的ID
label_to_class_id = {"tree": 0# 根据需要添加更多类别
}# json转txt
def convert_labelme_json_to_yolo(json_file, output_dir, img_width, img_height):with open(json_file, 'r') as f:labelme_data = json.load(f)# 获取文件名（不含扩展名）file_name = os.path.splitext(os.path.basename(json_file))[0]# 输出的txt文件路径txt_file_path = os.path.join(output_dir, f"{file_name}.txt")with open(txt_file_path, 'w') as txt_file:for shape in labelme_data['shapes']:label = shape['label']points = shape['points']# 根据类别映射表获取类别ID，如果类别不在映射表中，跳过该标签class_id = label_to_class_id.get(label)if class_id is None:print(f"Warning: Label '{label}' not found in class mapping. Skipping.")continue# 将点的坐标归一化到0-1范围normalized_points = [(x / img_width, y / img_height) for x, y in points]# 写入类别IDtxt_file.write(f"{class_id}")# 写入多边形掩膜的所有归一化顶点坐标for point in normalized_points:txt_file.write(f" {point[0]:.6f} {point[1]:.6f}")txt_file.write("\n")if __name__ == "__main__":json_dir = "json"  # 替换为LabelMe标注的JSON文件目录output_dir = "labels"  # 输出的YOLO格式txt文件目录img_width = 500   # 图像宽度，根据实际图片尺寸设置img_height = 500  # 图像高度，根据实际图片尺寸设置# 创建输出文件夹if not os.path.exists(output_dir):os.makedirs(output_dir)# 批量处理所有json文件for json_file in os.listdir(json_dir):if json_file.endswith(".json"):json_path = os.path.join(json_dir, json_file)convert_labelme_json_to_yolo(json_path, output_dir, img_width, img_height)

2、数据集扩充（带json标签）

import time
import random
import cv2
import os
import numpy as np
from skimage.util import random_noise
import base64
import json
import re
from copy import deepcopy
import argparseclass DataAugmentForObjectDetection():#代码中包含五中数据增强的手段(噪声，光线，改变像素点，平移，镜像，打开后的数据增强为True，取消为False)def __init__(self, change_light_rate=0.5,add_noise_rate=0.2, random_point=0.5, flip_rate=0.5, shift_rate=0.5, rand_point_percent=0.03,is_addNoise=True, is_changeLight=False, is_random_point=True, is_shift_pic_bboxes=True,is_filp_pic_bboxes=True):self.change_light_rate = change_light_rateself.add_noise_rate = add_noise_rateself.random_point = random_pointself.flip_rate = flip_rateself.shift_rate = shift_rateself.rand_point_percent = rand_point_percent# 是否使用某种增强方式self.is_addNoise = is_addNoiseself.is_changeLight = is_changeLightself.is_random_point = is_random_pointself.is_filp_pic_bboxes = is_filp_pic_bboxesself.is_shift_pic_bboxes = is_shift_pic_bboxes# 加噪声(随机噪声)def _addNoise(self, img):return random_noise(img, seed=int(time.time())) * 255# 调整亮度def _changeLight(self, img):alpha = random.uniform(0.35, 1)blank = np.zeros(img.shape, img.dtype)return cv2.addWeighted(img, alpha, blank, 1 - alpha, 0)# 随机的改变点的值def _addRandPoint(self, img):percent = self.rand_point_percentnum = int(percent * img.shape[0] * img.shape[1])for i in range(num):rand_x = random.randint(0, img.shape[0] - 1)rand_y = random.randint(0, img.shape[1] - 1)if random.randint(0, 1) == 0:img[rand_x, rand_y] = 0else:img[rand_x, rand_y] = 255return img# 平移图像(注：需要到labelme工具上调整图像，部分平移的标注框可能会超出图像边界，对训练造成影响)def _shift_pic_bboxes(self, img, json_info):h, w, _ = img.shapex_min = wx_max = 0y_min = hy_max = 0shapes = json_info['shapes']for shape in shapes:points = np.array(shape['points'])x_min = min(x_min, points[:, 0].min())y_min = min(y_min, points[:, 1].min())x_max = max(x_max, points[:, 0].max())y_max = max(y_max, points[:, 0].max())d_to_left = x_mind_to_right = w - x_maxd_to_top = y_mind_to_bottom = h - y_maxx = random.uniform(-(d_to_left - 1) / 3, (d_to_right - 1) / 3)y = random.uniform(-(d_to_top - 1) / 3, (d_to_bottom - 1) / 3)M = np.float32([[1, 0, x], [0, 1, y]])shift_img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))for shape in shapes:for p in shape['points']:p[0] += xp[1] += yreturn shift_img, json_info# 图像镜像翻转def _filp_pic_bboxes(self, img, json_info):h, w, _ = img.shapesed = random.random()if 0 < sed < 0.33:flip_img = cv2.flip(img, 0)  # _flip_xinver = 0elif 0.33 < sed < 0.66:flip_img = cv2.flip(img, 1)  # _flip_yinver = 1else:flip_img = cv2.flip(img, -1)  # flip_x_yinver = -1shapes = json_info['shapes']for shape in shapes:for p in shape['points']:if inver == 0:p[1] = h - p[1]elif inver == 1:p[0] = w - p[0]elif inver == -1:p[0] = w - p[0]p[1] = h - p[1]return flip_img, json_infodef dataAugment(self, img, dic_info):change_num = 0while change_num < 1:if self.is_changeLight:if random.random() > self.change_light_rate:change_num += 1img = self._changeLight(img)if self.is_addNoise:if random.random() < self.add_noise_rate:change_num += 1img = self._addNoise(img)if self.is_random_point:if random.random() < self.random_point:change_num += 1img = self._addRandPoint(img)if self.is_shift_pic_bboxes:if random.random() < self.shift_rate:change_num += 1img, dic_info = self._shift_pic_bboxes(img, dic_info)if self.is_filp_pic_bboxes or 1:if random.random() < self.flip_rate:change_num += 1img, bboxes = self._filp_pic_bboxes(img, dic_info)return img, dic_infoclass ToolHelper():# 从json文件中提取原始标定的信息def parse_json(self, path):with open(path)as f:json_data = json.load(f)return json_data# 对图片进行字符编码def img2str(self, img_name):with open(img_name, "rb")as f:base64_data = str(base64.b64encode(f.read()))match_pattern = re.compile(r'b\'(.*)\'')base64_data = match_pattern.match(base64_data).group(1)return base64_data# 保存图片结果def save_img(self, save_path, img):cv2.imwrite(save_path, img)# 保持json结果def save_json(self, file_name, save_folder, dic_info):with open(os.path.join(save_folder, file_name), 'w') as f:json.dump(dic_info, f, indent=2)if __name__ == '__main__':need_aug_num = 5  #每张图片需要增强的次数toolhelper = ToolHelper()is_endwidth_dot = True  #文件是否以.jpg或者png结尾dataAug = DataAugmentForObjectDetection()parser = argparse.ArgumentParser()parser.add_argument('--source_img_json_path', type=str, default=r'/home/leeqianxi/YOLO/datasets/data/data')#需要更改的json地址parser.add_argument('--save_img_json_path', type=str, default=r'/home/leeqianxi/YOLO/datasets/data/new_data')#改变后的json保存地址args = parser.parse_args()source_img_json_path = args.source_img_json_path  # 图片和json文件原始位置save_img_json_path = args.save_img_json_path  # 图片增强结果保存文件# 如果保存文件夹不存在就创建if not os.path.exists(save_img_json_path):os.mkdir(save_img_json_path)for parent, _, files in os.walk(source_img_json_path):files.sort()  # 排序一下for file in files:if file.endswith('jpg') or file.endswith('png'):cnt = 0pic_path = os.path.join(parent, file)json_path = os.path.join(parent, file[:-4] + '.json')json_dic = toolhelper.parse_json(json_path)# 如果图片是有后缀的if is_endwidth_dot:# 找到文件的最后名字dot_index = file.rfind('.')_file_prefix = file[:dot_index]  # 文件名的前缀_file_suffix = file[dot_index:]  # 文件名的后缀img = cv2.imread(pic_path)while cnt < need_aug_num:  # 继续增强auged_img, json_info = dataAug.dataAugment(deepcopy(img), deepcopy(json_dic))img_name = '{}_{}{}'.format(_file_prefix, cnt + 1, _file_suffix)  # 图片保存的信息img_save_path = os.path.join(save_img_json_path, img_name)toolhelper.save_img(img_save_path, auged_img)  # 保存增强图片json_info['imagePath'] = img_namebase64_data = toolhelper.img2str(img_save_path)json_info['imageData'] = base64_datatoolhelper.save_json('{}_{}.json'.format(_file_prefix, cnt + 1),save_img_json_path, json_info)  # 保存xml文件print(img_name)cnt += 1  # 继续增强下一张

 3、数据集划分（训练集、测试集、验证集）

# 将图片和标注数据按比例切分为 训练集和测试集
import shutil
import random
import os# 原始路径
image_original_path = "/home/leeqianxi/YOLO/ultralytics/pic/"
label_original_path = "/home/leeqianxi/YOLO/ultralytics/labels/"cur_path = os.getcwd()
# 训练集路径
train_image_path = os.path.join(cur_path, "datasets/images/train/")
train_label_path = os.path.join(cur_path, "datasets/labels/train/")# 验证集路径
val_image_path = os.path.join(cur_path, "datasets/images/val/")
val_label_path = os.path.join(cur_path, "datasets/labels/val/")# 测试集路径
test_image_path = os.path.join(cur_path, "datasets/images/test/")
test_label_path = os.path.join(cur_path, "datasets/labels/test/")# 训练集目录
list_train = os.path.join(cur_path, "datasets/train.txt")
list_val = os.path.join(cur_path, "datasets/val.txt")
list_test = os.path.join(cur_path, "datasets/test.txt")train_percent = 0.8
val_percent = 0.2
test_percent = 0def del_file(path):for i in os.listdir(path):file_data = path + "\\" + ios.remove(file_data)def mkdir():if not os.path.exists(train_image_path):os.makedirs(train_image_path)else:del_file(train_image_path)if not os.path.exists(train_label_path):os.makedirs(train_label_path)else:del_file(train_label_path)if not os.path.exists(val_image_path):os.makedirs(val_image_path)else:del_file(val_image_path)if not os.path.exists(val_label_path):os.makedirs(val_label_path)else:del_file(val_label_path)if not os.path.exists(test_image_path):os.makedirs(test_image_path)else:del_file(test_image_path)if not os.path.exists(test_label_path):os.makedirs(test_label_path)else:del_file(test_label_path)def clearfile():if os.path.exists(list_train):os.remove(list_train)if os.path.exists(list_val):os.remove(list_val)if os.path.exists(list_test):os.remove(list_test)def main():mkdir()clearfile()file_train = open(list_train, 'w')file_val = open(list_val, 'w')file_test = open(list_test, 'w')total_txt = os.listdir(label_original_path)num_txt = len(total_txt)list_all_txt = range(num_txt)num_train = int(num_txt * train_percent)num_val = int(num_txt * val_percent)num_test = num_txt - num_train - num_valtrain = random.sample(list_all_txt, num_train)# train从list_all_txt取出num_train个元素# 所以list_all_txt列表只剩下了这些元素val_test = [i for i in list_all_txt if not i in train]# 再从val_test取出num_val个元素，val_test剩下的元素就是testval = random.sample(val_test, num_val)print("训练集数目：{}, 验证集数目：{}, 测试集数目：{}".format(len(train), len(val), len(val_test) - len(val)))for i in list_all_txt:name = total_txt[i][:-4]srcImage = image_original_path + name + '.png'srcLabel = label_original_path + name + ".txt"if i in train:dst_train_Image = train_image_path + name + '.png'dst_train_Label = train_label_path + name + '.txt'shutil.copyfile(srcImage, dst_train_Image)shutil.copyfile(srcLabel, dst_train_Label)file_train.write(dst_train_Image + '\n')elif i in val:dst_val_Image = val_image_path + name + '.png'dst_val_Label = val_label_path + name + '.txt'shutil.copyfile(srcImage, dst_val_Image)shutil.copyfile(srcLabel, dst_val_Label)file_val.write(dst_val_Image + '\n')else:dst_test_Image = test_image_path + name + '.jpg'dst_test_Label = test_label_path + name + '.txt'shutil.copyfile(srcImage, dst_test_Image)shutil.copyfile(srcLabel, dst_test_Label)file_test.write(dst_test_Image + '\n')file_train.close()file_val.close()file_test.close()if __name__ == "__main__":main()

4、图像裁剪为固定大小

from PIL import Image
import osdef crop_image(image_path, output_dir, crop_c, crop_size=(500, 500)):# 打开原始图片img = Image.open(image_path)img_width, img_height = img.sizecrop_width, crop_height = crop_size# 确保输出目录存在if not os.path.exists(output_dir):os.makedirs(output_dir)# 计算可以裁剪的行数和列数horizontal_crops = img_width // crop_widthvertical_crops = img_height // crop_height# 裁剪并保存子图crop_count = crop_cfor i in range(vertical_crops):for j in range(horizontal_crops):left = j * crop_widthupper = i * crop_heightright = left + crop_widthlower = upper + crop_height# 裁剪图像cropped_img = img.crop((left, upper, right, lower))# 保存裁剪后的图像output_path = os.path.join(output_dir, f"crop_{crop_count+1}.png")cropped_img.save(output_path)crop_count += 1print(f"裁剪完成，共裁剪 {crop_count} 张图片。")if __name__ == "__main__":image_path = "img.png"  # 输入图片的路径output_dir = 'cropped_images'  # 输出文件夹路径crop_c = 0crop_image(image_path, output_dir, crop_c)