动机

在做一个关于目标检测的项目的时候，用到了部分coco格式的自定义数据集，只有一个images图片文件夹和json文件，为了项目能够更方便运行，就尽量和coco数据集结构对齐，因此对其进行train_val_test划分。
前提是已经有该数据集整体的coco json文件，目的就是将数据集划分的同时也方便划分json文件。

• coco结构

${CenterNet_ROOT}
|-- data
`-- |-- coco
    `-- |-- annotations
        |   |-- train2017.json
        |   |-- val2017.json

        |-- train2017
        |-- val2017
        |-- test2017

代码

1. 将图片数据集划分成train、test和val三个文件夹数据集

# 工具类
import os
import random
import shutil
from shutil import copy2


def data_set_split(src_data_folder, target_data_folder, train_scale=0.8, val_scale=0.1, test_scale=0.1):
    '''
    读取源数据文件夹，生成划分好的文件夹，分为trian、val、test三个文件夹进行
    :param src_data_folder: 源文件夹 
    :param target_data_folder: 目标文件夹 
    :param train_scale: 训练集比例
    :param val_scale: 验证集比例
    :param test_scale: 测试集比例
    :return:
    '''
    print("开始数据集划分")
    class_names = os.listdir(src_data_folder)
    # 在目标目录下创建文件夹
    split_names = ['train', 'val', 'test']
    for split_name in split_names:
        split_path = os.path.join(target_data_folder, split_name)
        if os.path.isdir(split_path):
            pass
        else:
            os.mkdir(split_path)
        # 然后在split_path的目录下创建类别文件夹
        for class_name in class_names:
            class_split_path = os.path.join(split_path, class_name)
            if os.path.isdir(class_split_path):
                pass
            else:
                os.mkdir(class_split_path)

    # 按照比例划分数据集，并进行数据图片的复制
    # 首先进行分类遍历
    for class_name in class_names:
        current_class_data_path = os.path.join(src_data_folder, class_name)
        current_all_data = os.listdir(current_class_data_path)
        current_data_length = len(current_all_data)
        current_data_index_list = list(range(current_data_length))
        random.shuffle(current_data_index_list)

        train_folder = os.path.join(os.path.join(target_data_folder, 'train'), class_name)
        val_folder = os.path.join(os.path.join(target_data_folder, 'val'), class_name)
        test_folder = os.path.join(os.path.join(target_data_folder, 'test'), class_name)
        train_stop_flag = current_data_length * train_scale
        val_stop_flag = current_data_length * (train_scale + val_scale)
        current_idx = 0
        train_num = 0
        val_num = 0
        test_num = 0
        for i in current_data_index_list:
            src_img_path = os.path.join(current_class_data_path, current_all_data[i])
            if current_idx <= train_stop_flag:
                copy2(src_img_path, train_folder)
                # print("{}复制到了{}".format(src_img_path, train_folder))
                train_num = train_num + 1
            elif (current_idx > train_stop_flag) and (current_idx <= val_stop_flag):
                copy2(src_img_path, val_folder)
                # print("{}复制到了{}".format(src_img_path, val_folder))
                val_num = val_num + 1
            else:
                copy2(src_img_path, test_folder)
                # print("{}复制到了{}".format(src_img_path, test_folder))
                test_num = test_num + 1

            current_idx = current_idx + 1

        print("*********************************{}*************************************".format(class_name))
        print(
            "{}类按照{}：{}：{}的比例划分完成，一共{}张图片".format(class_name, train_scale, val_scale, test_scale, current_data_length))
        print("训练集{}：{}张".format(train_folder, train_num))
        print("验证集{}：{}张".format(val_folder, val_num))
        print("测试集{}：{}张".format(test_folder, test_num))


if __name__ == '__main__':
    src_data_folder = "./dataset"
    target_data_folder = "./output"
    data_set_split(src_data_folder, target_data_folder)

2、根据图片名和COCO总JSON文件生成对应的JSON

根据下面代码，分别生成train、test、val的json文件

# -*- coding: utf-8 -*-
from __future__ import print_function
"""
Created on 2021/12/9

@author: Totie
"""

'''
从coco的标注文件里提取一张图片对应的json信息，并保存成新的json文件（以instance为例，其他的类似）。
'''
import matplotlib.pyplot as plt
import os, sys, zipfile
import urllib.request
import shutil
import numpy as np
# import skimage.io as io
import pylab
import json
from pycocotools.coco import COCO

pylab.rcParams['figure.figsize'] = (8.0, 10.0)
dataset_path = './coco'
image_path = 'val2017'
json_file='/root/CenterNet-ht/Center/data/coco/annotations/train.json' # # json源文件

coco=COCO(json_file)
data=json.load(open(json_file,'r')) 

data_2={}   #新json文件

data_2['info']=data['info']
data_2['licenses']=data['licenses']
data_2['categories']=data['categories']

annotation=[]
images = []
# print(data_img['images'])

imagename = [f for f in os.listdir(os.path.join(dataset_path, image_path))] #读取文件夹下图片名字
print(len(data['images'])) 

#根据图片数量找到每张图片对应的annotation，即每个‘images’可能有多个annotation（一张图片有多个可识别的目标）
for name_index in range(0,len(imagename)):
    # 通过imgID 找到其所有instance
    imgID = 0
    for i in range(0,len(data['images'])):
        if data['images'][i]['file_name'] == imagename[name_index]:  #根据图片名找到对应的json中的'images'
            imgID = imgID=data['images'][i]['id']
            print(name_index, imgID)  

    for ann in data['annotations']:    #根据image_id找到对应的annotation
        if ann['image_id']==imgID:
            annotation.append(ann)

data_2['annotations']=annotation
data_2['images'] = images
print(len(data_2['annotations']))
# 保存到新的json

json.dump(data_2,open('./val2017.json','w'),indent=4)
# 从coco标注json中提取单张图片的标注信息

最后

以上就是个性花生为你收集整理的根据图片名生成COCO格式json文件，train_val_test数据集的划分动机代码的全部内容，希望文章能够帮你解决根据图片名生成COCO格式json文件，train_val_test数据集的划分动机代码所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。