将SSDD数据集按比例划分训练集、验证集和测试集，并从VOC标注格式转换到COCO格式

SSDD数据集共1160张图片，该数据集使用voc格式标注，原始数据并没有划分好，需要自行划分。

在网上查找到voc格式数据集的相关划分代码，自己对其进行一定修改后可以成功运行，现在将代码记录如下：

'''
本脚本用来将所有的xml标注文件，随机划分成事先设定比例的训练集、验证集、训练验证集、测试集；结果是输出储存这些文件名（只是文件名而没有'.xml'这个后缀）的txt文件；从而得到到voc格式的数据集
本脚本在文件将数据集按7：2：1的比例分成训练集train、验证集val、测试集test
'''
import os
import random
random.seed(0)
xmlfilepath = r'/home/dwt/DataSets/SSDD_train_val_test/newAnnotations'
saveBasePath=r"/home/dwt/DataSets/SSDD_train_val_test/ImageSets/Main/"

#----------------------------------------------------------------------#
#   想要增加测试集修改trainval_percent
#   train_percent不需要修改
#----------------------------------------------------------------------#
train_percent = 0.7 #训练集在数据集中的比例
val_percent = 0.2   #验证集在数据集中的比例
trainval_percent = 0.9 #训练集和验证集加在一起在数据集中的比例
test_percent = 0.1   #测试集在数据集中的比例

temp_xml = os.listdir(xmlfilepath) #os.listdir() 方法用于返回指定的文件夹包含的文件或文件夹的名字的列表。
total_xml = []
for xml in temp_xml:
    if xml.endswith(".xml"):
        total_xml.append(xml)

num=len(total_xml) #num表示所有标注文件的总数量
total_list=range(num) #total_list是从0到1的总元素数量为num的列表。与total_xml列表一一对应
train_num = int(num*train_percent)#train_num表示所有标注文件中用于训练集的标注文件的数量
val_num = int(num*val_percent) #val_num表示所有标注文件中用于验证集的标注文件的数量
trainval_num = int(num*trainval_percent) #trainval_num表示所有标注文件中用于训练集和验证集的标注文件的数量
test_num = int(num*test_percent) #test_num表示所有标注文件中用于测试集的标注文件的数量

trainval_list = random.sample(total_list,trainval_num)#trainval_list时存储用于训练集和验证集的标注文件的序号的列表 #trainval = random.sample(list,tv)函数从列表list不分大小顺序的抽取tv个数，并抽取结果以列表的形式传给trainval
train_list = random.sample(trainval_list,train_num) #train_list是储存作为训练集中用于训练的标注文件在total_list列表中的序号

print("train and val size",trainval_num)
print("train size: {} , val size : {} , test size : {}".format(train_num,val_num,test_num))

ftrainval = open(os.path.join(saveBasePath,'trainval.txt'), 'w')
ftest = open(os.path.join(saveBasePath,'test.txt'), 'w')
ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')
fval = open(os.path.join(saveBasePath,'val.txt'), 'w')

for i  in total_list:
    name=total_xml[i][:-4]+'\n' #因为voc格式数据集的储存训练和测试集信息的txt文件是只有文件名的，所以诸如“aa.xml”的文件只将“aa”写入到储存训练和测试集信息的txt文件中
    if i in trainval_list:
        ftrainval.write(name)
        if i in train_list:
            ftrain.write(name)
        else:
            fval.write(name)
    else:
        ftest.write(name)

ftrainval.close()
ftrain.close()
fval.close()
ftest .close()

下面是将voc格式数据集注释（xml文件）转换到coco格式数据集（json）文件的代码（其中parseXMLFiles函数是将给定文件夹下的xml标注文件统一转换成json文件，parseXMLFiles_by_txt函数具体功能与parseXMLFiles函数一样，只是会根据voc格式数据集中区分训练集和测试集数据的train.txt和test.txt文件来将对应的xml文件转换成json文件）：

注：本段内容参考了博客VOC数据集格式转化成COCO数据集格式中的代码，对其原代码添加了注释方便理解，并对细节出做出适当修改。

import xml.etree.ElementTree as ET
import os
import json

coco = dict()
coco['images'] = []
coco['type'] = 'instances'
coco['annotations'] = []
coco['categories'] = []

category_set = dict() #该字典负责统计该数据集中出现的目标种类名称及其对应的id值
image_set = set()

category_item_id = -1
image_id = 0
annotation_id = 0


def addCatItem(name): #该函数对json文件中的categories这一项添加内容
    global category_item_id #category_item_id原本是该函数之外的变量，使用global关键字作用于category_item_id上后，该函数内对category_item_id的操作会影响到函数外category_item_id的值
    category_item = dict()
    category_item['supercategory'] = 'none'
    category_item_id += 1
    category_item['id'] = category_item_id
    category_item['name'] = name
    coco['categories'].append(category_item)
    category_set[name] = category_item_id #对category_set字典添加内容
    return category_item_id


def addImgItem(file_name, size):
    global image_id
    if file_name is None:
        raise Exception('Could not find filename tag in xml file.')
    if size['width'] is None:
        raise Exception('Could not find width tag in xml file.')
    if size['height'] is None:
        raise Exception('Could not find height tag in xml file.')
    image_id += 1
    image_item = dict()
    image_item['id'] = image_id
    image_item['file_name'] = file_name
    image_item['width'] = size['width']
    image_item['height'] = size['height']
    coco['images'].append(image_item)
    image_set.add(file_name)
    return image_id


def addAnnoItem(object_name, image_id, category_id, bbox):
    global annotation_id
    annotation_item = dict()
    annotation_item['segmentation'] = []
    seg = []
    # bbox[] is x,y,w,h
    # left_top
    seg.append(bbox[0])
    seg.append(bbox[1])
    # left_bottom
    seg.append(bbox[0])
    seg.append(bbox[1] + bbox[3])
    # right_bottom
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1] + bbox[3])
    # right_top
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1])

    annotation_item['segmentation'].append(seg)

    annotation_item['area'] = bbox[2] * bbox[3]
    annotation_item['iscrowd'] = 0
    annotation_item['ignore'] = 0
    annotation_item['image_id'] = image_id
    annotation_item['bbox'] = bbox
    annotation_item['category_id'] = category_id
    annotation_id += 1
    annotation_item['id'] = annotation_id
    coco['annotations'].append(annotation_item)


def _read_image_ids(image_sets_file):
    ids = []
    with open(image_sets_file) as f:
        for line in f:
            ids.append(line.rstrip())
    return ids


"""通过txt文件生成"""


# split ='train' 'va' 'trainval' 'test'
def parseXmlFiles_by_txt(data_dir, json_save_path, split='train'): #功能与parseXmlFiles相同，只是parseXmlFiles_by_txt根据txt文件中指示的文件进行对应xml的转化
    print("hello")
    labelfile = split + ".txt"
    image_sets_file = data_dir + "/ImageSets/Main/" + labelfile
    ids = _read_image_ids(image_sets_file)

    for _id in ids:
        xml_file = data_dir + f"/Annotations/{_id}.xml"

        bndbox = dict()
        size = dict()
        current_image_id = None
        current_category_id = None
        file_name = None
        size['width'] = None
        size['height'] = None
        size['depth'] = None

        tree = ET.parse(xml_file)
        root = tree.getroot()
        if root.tag != 'annotation':
            raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))

        # elem is <folder>, <filename>, <size>, <object>
        for elem in root:
            current_parent = elem.tag
            current_sub = None
            object_name = None

            if elem.tag == 'folder':
                continue

            if elem.tag == 'filename':
                file_name = elem.text
                if file_name in category_set:
                    raise Exception('file_name duplicated')

            # add img item only after parse <size> tag
            elif current_image_id is None and file_name is not None and size['width'] is not None:
                if file_name not in image_set:
                    current_image_id = addImgItem(file_name, size)
                    print('add image with {} and {}'.format(file_name, size))
                else:
                    raise Exception('duplicated image: {}'.format(file_name))
                    # subelem is <width>, <height>, <depth>, <name>, <bndbox>
            for subelem in elem:
                bndbox['xmin'] = None
                bndbox['xmax'] = None
                bndbox['ymin'] = None
                bndbox['ymax'] = None

                current_sub = subelem.tag
                if current_parent == 'object' and subelem.tag == 'name': #当循环来到xml文件的object中的name项（目标的种类名称信息）
                    object_name = subelem.text
                    if object_name not in category_set:#如果当前目标的种类名称先前没出现过则将该目标的种类名称进行统计整理，其中category_set字典负责统计该数据集中出现的目标种类名称及其对应的id值
                        current_category_id = addCatItem(object_name)#该函数对json文件中的categories这一项添加内容
                    else:
                        current_category_id = category_set[object_name] #如果当前目标的种类名称已出现过，则从category_set字典调出当前种类名称对应的编号

                elif current_parent == 'size': #当循环来到xml文件的size项（图片的尺寸信息）
                    if size[subelem.tag] is not None:
                        raise Exception('xml structure broken at size tag.')
                    size[subelem.tag] = int(subelem.text) #对图片size信息下面的width、height、depth三项信息进行记录

                # option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox>
                for option in subelem:
                    if current_sub == 'bndbox': #当循环来到xml文件的object中的bndbox项（目标框的坐标信息）
                        if bndbox[option.tag] is not None:
                            raise Exception('xml structure corrupted at bndbox tag.')
                        bndbox[option.tag] = int(option.text) #将当前object项中的目标框bndbox的xmin、ymin、xmax、ymax进行记录到bnbbox中临时储存

                # only after parse the <object> tag
                if bndbox['xmin'] is not None: #当获取到目标的目标框坐标信息后bndbox后，说明当前目标的信息已获取完毕，接下来将获取的该目标的信息进行添加记录操作
                    if object_name is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_image_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_category_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    bbox = []
                    # x
                    bbox.append(bndbox['xmin'])
                    # y
                    bbox.append(bndbox['ymin'])
                    # w
                    bbox.append(bndbox['xmax'] - bndbox['xmin'])
                    # h
                    bbox.append(bndbox['ymax'] - bndbox['ymin'])
                    print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id,
                                                                   bbox))
                    addAnnoItem(object_name, current_image_id, current_category_id, bbox)
    json.dump(coco, open(json_save_path, 'w'))


"""直接从xml文件夹中生成"""


def parseXmlFiles(xml_path, json_save_path):
    for f in os.listdir(xml_path):
        if not f.endswith('.xml'):
            continue

        bndbox = dict()
        size = dict()
        current_image_id = None
        current_category_id = None
        file_name = None
        size['width'] = None
        size['height'] = None
        size['depth'] = None

        xml_file = os.path.join(xml_path, f) #xml_file为当前xml文件的具体地址
        print(xml_file)

        tree = ET.parse(xml_file)
        root = tree.getroot()
        if root.tag != 'annotation':
            raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))

        # elem is <folder>, <filename>, <size>, <object>
        for elem in root:
            current_parent = elem.tag
            current_sub = None
            object_name = None

            if elem.tag == 'folder':
                continue

            if elem.tag == 'filename':
                file_name = elem.text
                if file_name in category_set:
                    raise Exception('file_name duplicated')

            # add img item only after parse <size> tag
            elif current_image_id is None and file_name is not None and size['width'] is not None:
                if file_name not in image_set:
                    current_image_id = addImgItem(file_name, size)
                    print('add image with {} and {}'.format(file_name, size))
                else:
                    raise Exception('duplicated image: {}'.format(file_name))
                    # subelem is <width>, <height>, <depth>, <name>, <bndbox>
            for subelem in elem:
                bndbox['xmin'] = None
                bndbox['xmax'] = None
                bndbox['ymin'] = None
                bndbox['ymax'] = None

                current_sub = subelem.tag
                if current_parent == 'object' and subelem.tag == 'name':#当elem项为object且subelem为name时才进行下面的操作
                    object_name = subelem.text #获取当前目标的种类名称
                    if object_name not in category_set:
                        current_category_id = addCatItem(object_name)
                    else:
                        current_category_id = category_set[object_name]

                elif current_parent == 'size':
                    if size[subelem.tag] is not None:
                        raise Exception('xml structure broken at size tag.')
                    size[subelem.tag] = int(subelem.text)

                # option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox>
                for option in subelem:
                    if current_sub == 'bndbox':
                        if bndbox[option.tag] is not None:
                            raise Exception('xml structure corrupted at bndbox tag.')
                        bndbox[option.tag] = int(option.text)

                # only after parse the <object> tag
                if bndbox['xmin'] is not None:
                    if object_name is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_image_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_category_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    bbox = []
                    # x
                    bbox.append(bndbox['xmin'])
                    # y
                    bbox.append(bndbox['ymin'])
                    # w
                    bbox.append(bndbox['xmax'] - bndbox['xmin'])
                    # h
                    bbox.append(bndbox['ymax'] - bndbox['ymin'])
                    print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id,
                                                                   bbox))
                    addAnnoItem(object_name, current_image_id, current_category_id, bbox)
    json.dump(coco, open(json_save_path, 'w'))


if __name__ == '__main__':
    # 通过txt文件生成
    voc_data_dir="/home/dwt/MyCode/pycharm_projects/YOLOX_sample/datasets/VOCdevkit-hrsid/VOC2011"
    json_save_path= "data/coco/hrsid_coco_annotations/train.json"
    parseXmlFiles_by_txt(voc_data_dir,json_save_path,"train")

    # # 通过文件夹生成
    # ann_path = "E:/VOCdevkit/VOC2007/Annotations" #voc标注文件地址
    # json_save_path = "E:/VOCdevkit/test.json" #json文件保存地址
    # parseXmlFiles(ann_path, json_save_path)

下面的代码是在将SSDD从voc格式转换到coco格式（即：将voc格式标注信息xml文件转换到coco格式标注信息json文件）之后，再根据train、val、test这三个json文件中对应的图片归类成对应的三个文件夹:image_train、image_val、image_test，所写的脚本的代码记录：

（1）初始代码

#move_rename_image.py
#在这个脚本中实现根据annotation中图片id与图片名称来对image图片改名为其对应id并另存到单独的文件夹中
import os
import json
import argparse
import  shutil

parser = argparse.ArgumentParser()
parser.add_argument('--source_path', default='D:/engineering/intrusion detection/dataset/preliminary_contest_crane_federal',type=str)
parser.add_argument('--save_path', default='D:/engineering/intrusion detection/dataset/diaoche_data', type=str, help="specify where to save the output dir of labels")
arg = parser.parse_args()

c = 0 #设置这个变量用来计数处理的标签数，在训练过程中用print打印c来指示进度

if __name__ == '__main__':
    json_file = os.path.join(arg.source_path, 'annotations', 'train2.json')
    #image_file_catalog = os.path.join(arg.save_path, 'images')
    data = json.load(open(json_file, encoding="utf-8"))
    images = data['images']

    #在下面的循环中将图片改成其对应id的名称并另存到指定的文件夹
    #在处理train2中的数据时，因为计划将train2中图片合在train1之后，故train2图片的重命名根据其对应annotation中id号加1159
    for img_info in images:
        img_name = img_info['file_name']
        img_id = img_info['id']
        src = os.path.join(arg.source_path,'images','train2',img_name)
        dst = os.path.join(arg.save_path,'images',f'{str(img_id + 1159)}.jpg')
        shutil.copy(src,dst)

        c += 1
        if c % 50 == 0:
            print("当前已经重命名并另存%d个吊机图片" % (c))

（2）根据上面的初始代码又改写的代码

#image_tidy.py
'''在voc转化成coco格式后，需要将原本混在一起的图片分成train、val、test这三个文件夹'''
import os
import json
import argparse
import shutil


# train_txt = '/home/dwt/DataSets/SSDD_train_val_test/VOCdevkit/VOC2011/ImageSets/Main/train.txt'
# val_txt = '/home/dwt/DataSets/SSDD_train_val_test/VOCdevkit/VOC2011/ImageSets/Main/val.txt'
# test_txt = '/home/dwt/DataSets/SSDD_train_val_test/VOCdevkit/VOC2011/ImageSets/Main/test.txt'

train_json = '/home/dwt/DataSets/SSDD_train_val_test/coco_format/voc2011_train.json'
val_json = '/home/dwt/DataSets/SSDD_train_val_test/coco_format/voc2011_val.json'
test_json = '/home/dwt/DataSets/SSDD_train_val_test/coco_format/voc2011_test.json'

src_front = '/home/dwt/DataSets/SSDD_train_val_test/VOCdevkit'
dst_front = '/home/dwt/DataSets/SSDD_train_val_test/coco_format'

files = ["train","val","test"]
for file in files:
    if file == "train":
        dir_json = train_json
        dst_file = os.path.join(dst_front,"image_train")
    if file == "val":
        dir_json = val_json
        dst_file = os.path.join(dst_front, "image_val")
    if file == "test":
        dir_json = test_json
        dst_file = os.path.join(dst_front, "image_test")
    data = json.load(open(dir_json,encoding="utf-8"))
    images = data['images']

    for img_info in images:
        img_dir = img_info["file_name"]# 例如："VOC2011/JPEGImages/000078.jpg"
        src = os.path.join(src_front,img_dir)
        dst = os.path.join(dst_file,img_dir[-10:])
        shutil.copy(src,dst)

    print("done")

print("Done!")