1. 划分数据集
将图片数据集划随机划分为训练:验证:测试=6:2:2的比例,同时划分对应的json目录文件。废话不多说,直接上代码:
import os import random import shutil # 指定原始数据目录 images_dir = "images" json_dir = "Annotations" # 指定目标输出目录 output_dir_images = "../../datasets/fish/images" output_dir_json = "../../datasets/fish/annotations" # 创建目标输出目录 if not os.path.exists(output_dir_images): os.makedirs(output_dir_images) if not os.path.exists(output_dir_json): os.makedirs(output_dir_json) # 创建训练集、验证集和测试集子目录 subdirectories = ["train", "val", "test"] for subdir in subdirectories: subdir_path_images = os.path.join(output_dir_images, subdir) subdir_path_json = os.path.join(output_dir_json, subdir) if not os.path.exists(subdir_path_images): os.makedirs(subdir_path_images) if not os.path.exists(subdir_path_json): os.makedirs(subdir_path_json) # 获取图像文件和 JSON 文件列表 image_files = [f for f in os.listdir(images_dir) if f.endswith(".jpg")] json_files = [f for f in os.listdir(json_dir) if f.endswith(".json")] # 打乱数据的顺序 random.shuffle(image_files) # 计算划分的数据集大小 total_samples = len(image_files) train_size = int(total_samples * 0.6) val_size = int(total_samples * 0.2) test_size = total_samples - train_size - val_size # 划分数据集 train_images = image_files[:train_size] val_images = image_files[train_size:train_size + val_size] test_images = image_files[train_size + val_size:] # 移动图像文件到对应的子目录 for image_file in train_images: src_path_images = os.path.join(images_dir, image_file) dest_path_images = os.path.join(output_dir_images, "train", image_file) shutil.copy(src_path_images, dest_path_images) for image_file in val_images: src_path_images = os.path.join(images_dir, image_file) dest_path_images = os.path.join(output_dir_images, "val", image_file) shutil.copy(src_path_images, dest_path_images) for image_file in test_images: src_path_images = os.path.join(images_dir, image_file) dest_path_images = os.path.join(output_dir_images, "test", image_file) shutil.copy(src_path_images, dest_path_images) # 移动 JSON 文件到对应的子目录 for json_file in json_files: base_name = os.path.splitext(json_file)[0] if f"{base_name}.jpg" in train_images: src_path_json = os.path.join(json_dir, json_file) dest_path_json = os.path.join(output_dir_json, "train", json_file) shutil.copy(src_path_json, dest_path_json) elif f"{base_name}.jpg" in val_images: src_path_json = os.path.join(json_dir, json_file) dest_path_json = os.path.join(output_dir_json, "val", json_file) shutil.copy(src_path_json, dest_path_json) elif f"{base_name}.jpg" in test_images: src_path_json = os.path.join(json_dir, json_file) dest_path_json = os.path.join(output_dir_json, "test", json_file) shutil.copy(src_path_json, dest_path_json) print("数据集划分完成。")
复制
路径都用的是相对路径,如果不理解相对路径,可以这么理解:就是以此文件的路径开始,**举个例子:**我的images图片目录和此划分数据集的文件在同一个目录之下,则输入图片目录的时候,就直接写入图片的目录名称即可:
可以从图中看到,我的划分代码和images目录在在同一个目录VOCData下面。所以输入的图片目录就是:
…/ 这里的…(两个.)代表上级目录,因为我是用linux服务器进行划分的。
以下是我的目录格式,可以进行参考:
执行过后,两个数据集(images,json)就被划分为了train,val,test三个子目录。并且这两个数据集的子目录的文件都是一一对应的。如图所示:
此刻,划分数据集到此结束!
2. json转txt
因为上一步把json的文件夹划分为了train,val和test三个子目录,接下来要做的就是将这三个子目录的json文件都转换为txt文件。这个可以参考我的上一篇文章:python实现将一个文件夹中的多个 JSON 文件批量转换为 YOLO 可以识别的 TXT 文件