标签：xml join img voc coco path os dir

平常用coco格式的数据集比较多，所有这里整合一下数据集相关的常用的脚本。

pycocotools安装

这个非常重要，因为处理coco数据集时，用pycocotools包非常方便。

自行搜索一下怎么安装吧，windows安装比较麻烦。网上有很多方法，但是都有时效性，不定时就失效了。如果有好的安装pycocotools的文章，可以把链接评论在评论区。

xml转voc

xml2voc.py

# 命令行执行：  python xml2voc2007.py --input_dir data --output_dir VOCdevkit
# 输出文件夹必须为空文件夹

import argparse
import glob
import os
import random
import sys
import shutil


# 主程序执行
def xml2voc():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--input_dir", default="data", help="input annotated directory")
    parser.add_argument("--output_dir", default="VOCdevkit", help="output dataset directory")
    args = parser.parse_args()

    if os.path.exists(args.output_dir):
        print("Output directory already exists:", args.output_dir)
        sys.exit(1)
    os.makedirs(args.output_dir)
    print("| Creating dataset dir:", os.path.join(args.output_dir, "VOC2007"))

    # 创建保存的文件夹
    if not os.path.exists(os.path.join(args.output_dir, "VOC2007", "Annotations")):
        os.makedirs(os.path.join(args.output_dir, "VOC2007", "Annotations"))
    if not os.path.exists(os.path.join(args.output_dir, "VOC2007", "ImageSets")):
        os.makedirs(os.path.join(args.output_dir, "VOC2007", "ImageSets"))

    if not os.path.exists(os.path.join(args.output_dir, "VOC2007", "ImageSets", "Main")):
        os.makedirs(os.path.join(args.output_dir, "VOC2007", "ImageSets", "Main"))
    if not os.path.exists(os.path.join(args.output_dir, "VOC2007", "JPEGImages")):
        os.makedirs(os.path.join(args.output_dir, "VOC2007", "JPEGImages"))

    # 获取目录下所有的.jpg文件列表
    total_img = glob.glob(os.path.join(args.input_dir, "*.jpg"))
    print('| Image number: ', len(total_img))

    # 获取目录下所有的joson文件列表
    total_xml = glob.glob(os.path.join(args.input_dir, "*.xml"))
    print('| Xml number: ', len(total_xml))

    percent_trainval = 0.8
    percent_train = 0.75
    num_total = len(total_xml)
    data_list = range(num_total)

    num_tv = int(num_total * percent_trainval)
    num_tr = int(num_tv * percent_train)
    num_trainval = random.sample(data_list, num_tv)
    num_train = random.sample(num_trainval, num_tr)

    print('| Train number: ', num_tr)
    print('| Val number: ', num_tv - num_tr)
    print('| Test number: ', num_total - num_tv)

    file_trainval = open(
        os.path.join(args.output_dir, "VOC2007", "ImageSets", "Main", "trainval.txt"), 'w')
    file_test = open(
        os.path.join(args.output_dir, "VOC2007", "ImageSets", "Main", "test.txt"), 'w')
    file_train = open(
        os.path.join(args.output_dir, "VOC2007", "ImageSets", "Main", "train.txt"), 'w')
    file_val = open(
        os.path.join(args.output_dir, "VOC2007", "ImageSets", "Main", "val.txt"), 'w')

    for i in data_list:
        name = os.path.basename(total_xml[i])[:-4] + '\n'  # 去掉.xml后缀以及父级目录，只保留文件名
        if i in num_trainval:
            file_trainval.write(name)
            if i in num_train:
                file_train.write(name)
            else:
                file_val.write(name)
        else:
            file_test.write(name)

    file_trainval.close()
    file_train.close()
    file_val.close()
    file_test.close()

    if os.path.exists(args.input_dir):
        # root 所指的是当前正在遍历的这个文件夹的本身的地址  
        # dirs 是一个 list，内容是该文件夹中所有的目录的名字(不包括子目录)  
        # files 同样是 list, 内容是该文件夹中所有的文件(不包括子目录)  
        for root, dirs, files in os.walk(args.input_dir):
            for file in files:
                src_file = os.path.join(root, file)
                if src_file.endswith(".jpg"):
                    shutil.copy(src_file, os.path.join(args.output_dir, "VOC2007", "JPEGImages"))
                else:
                    shutil.copy(src_file, os.path.join(args.output_dir, "VOC2007", "Annotations"))

    print('| Done!')


if __name__ == "__main__":
    print("—" * 50)
    xml2voc()
    print("—" * 50)

voc转coco

voc2coco.py

import json
import os
import shutil
import datetime
from PIL import Image
from tqdm import trange

root_dir = os.getcwd()


def voc2coco():
    # 处理coco数据集中category字段。
    # 创建一个 {类名 : id} 的字典，并保存到 总标签data 字典中。
    class_name_to_id = {'class1': 1, 'class2': 2, 'class3': 3, 'class4': 4, 'class5': 5, 'class6': 6, 'class7': 7, 'class8': 8}        # 改为自己的类别名称，以及对应的类别id

    # 创建coco的文件夹
    if not os.path.exists(os.path.join(root_dir, "coco")):
        os.makedirs(os.path.join(root_dir, "coco"))
        os.makedirs(os.path.join(root_dir, "coco", "annotations"))
        os.makedirs(os.path.join(root_dir, "coco", "train"))
        os.makedirs(os.path.join(root_dir, "coco", "val"))

    # 创建 总标签data
    now = datetime.datetime.now()
    data = dict(
        info=dict(
            description=None,
            url=None,
            version=None,
            year=now.year,
            contributor=None,
            date_created=now.strftime("%Y-%m-%d %H:%M:%S.%f"),
        ),
        licenses=[dict(url=None, id=0, name=None, )],
        images=[
            # license, file_name,url, height, width, date_captured, id
        ],
        type="instances",
        annotations=[
            # segmentation, area, iscrowd, image_id, bbox, category_id, id
        ],
        categories=[
            # supercategory, id, name
        ],
    )

    for name, id in class_name_to_id.items():
        data["categories"].append(
            dict(supercategory=None, id=id, name=name, )
        )

    # 处理coco数据集train中images字段。
    images_dir = os.path.join(root_dir, 'VOCdevkit', 'VOC2012', 'JPEGImages')
    images = os.listdir(images_dir)

    # 生成每个图片对应的image_id
    images_id = {}
    for idx, image_name in enumerate(images):
        images_id.update({image_name[:-4]: idx})

    # 获取训练图片
    train_img = []
    fp = open(os.path.join(root_dir, 'VOCdevkit', 'VOC2012', 'ImageSets', 'Main', 'train.txt'))
    for i in fp.readlines():
        train_img.append(i[:-1] + ".jpg")

    # 获取训练图片的数据
    for image in train_img:
        img = Image.open(os.path.join(images_dir, image))
        data["images"].append(
            dict(
                license=0,
                url=None,
                file_name=image,  # 图片的文件名带后缀
                height=img.height,
                width=img.width,
                date_captured=None,
                # id=image[:-4],
                id=images_id[image[:-4]],
            )
        )

    # 获取coco数据集train中annotations字段。
    train_xml = [i[:-4] + '.xml' for i in train_img]

    bbox_id = 0
    for xml in train_xml:
        category = []
        xmin = []
        ymin = []
        xmax = []
        ymax = []
        import xml.etree.ElementTree as ET
        tree = ET.parse(os.path.join(root_dir, 'VOCdevkit', 'VOC2012', 'Annotations', xml))
        root = tree.getroot()
        object = root.findall('object')
        for i in object:
            category.append(class_name_to_id[i.findall('name')[0].text])
            bndbox = i.findall('bndbox')
            for j in bndbox:
                xmin.append(float(j.findall('xmin')[0].text))
                ymin.append(float(j.findall('ymin')[0].text))
                xmax.append(float(j.findall('xmax')[0].text))
                ymax.append(float(j.findall('ymax')[0].text))
        for i in range(len(category)):
            data["annotations"].append(
                dict(
                    id=bbox_id,
                    image_id=images_id[xml[:-4]],
                    category_id=category[i],
                    area=(xmax[i] - xmin[i]) * (ymax[i] - ymin[i]),
                    bbox=[xmin[i], ymin[i], xmax[i] - xmin[i], ymax[i] - ymin[i]],
                    iscrowd=0,
                )
            )
            bbox_id += 1
    # 生成训练集的json
    json.dump(data, open(os.path.join(root_dir, 'coco', 'annotations', 'train.json'), 'w'))

    # 获取验证图片
    val_img = []
    fp = open(os.path.join(root_dir, 'VOCdevkit', 'VOC2012', 'ImageSets', 'Main', 'val.txt'))
    for i in fp.readlines():
        val_img.append(i[:-1] + ".jpg")

    # 将训练的images和annotations清空，
    del data['images']
    data['images'] = []
    del data['annotations']
    data['annotations'] = []

    # 获取验证集图片的数据
    for image in val_img:
        img = Image.open(os.path.join(images_dir, image))
        data["images"].append(
            dict(
                license=0,
                url=None,
                file_name=image,  # 图片的文件名带后缀
                height=img.height,
                width=img.width,
                date_captured=None,
                id=images_id[image[:-4]],   # 图片名作为id
            )
        )

    # 处理coco数据集验证集中annotations字段。
    val_xml = [i[:-4] + '.xml' for i in val_img]

    for xml in val_xml:
        category = []
        xmin = []
        ymin = []
        xmax = []
        ymax = []
        import xml.etree.ElementTree as ET
        tree = ET.parse(os.path.join(root_dir, 'VOCdevkit', 'VOC2012', 'Annotations', xml))
        root = tree.getroot()
        object = root.findall('object')
        for i in object:
            category.append(class_name_to_id[i.findall('name')[0].text])
            bndbox = i.findall('bndbox')
            for j in bndbox:
                xmin.append(float(j.findall('xmin')[0].text))
                ymin.append(float(j.findall('ymin')[0].text))
                xmax.append(float(j.findall('xmax')[0].text))
                ymax.append(float(j.findall('ymax')[0].text))
        for i in range(len(category)):
            data["annotations"].append(
                dict(
                    id=bbox_id,
                    image_id=images_id[xml[:-4]],
                    category_id=category[i],
                    area=(xmax[i] - xmin[i]) * (ymax[i] - ymin[i]),
                    bbox=[xmin[i], ymin[i], xmax[i] - xmin[i], ymax[i] - ymin[i]],
                    iscrowd=0,
                )
            )
            bbox_id += 1
    # 生成验证集的json
    json.dump(data, open(os.path.join(root_dir, 'coco', 'annotations', 'val.json'), 'w'))
    print('| VOC -> COCO annotations transform finish.')
    print('Start copy images...')

    # 复制图片
    m = len(train_img)
    for i in trange(m):
        shutil.copy(os.path.join(images_dir, train_img[i]), os.path.join(root_dir, 'coco', 'train', train_img[i]))
    print('| Train images copy finish.')

    m = len(val_img)
    for i in trange(m):
        shutil.copy(os.path.join(images_dir, val_img[i]), os.path.join(root_dir, 'coco', 'val', val_img[i]))
    print('| Val images copy finish.')


if __name__ == '__main__':
    voc2coco()

coco转yolo

coco数据集目录结构：

coco_small
├── train
│   ├── 001.jpg
│   ├── 002.jpg
│   ├── 003.jpg
│   └── ...
│
├── val
│   ├── 004.jpg
│   ├── 005.jpg
│   ├── 006.jpg
│   └── ...
├── test
│   ├── 007.jpg
│   ├── 008.jpg
│   ├── 009.jpg
│   └── ...
│
└── annotatoins
    ├── train.json
    ├── val.json
    └── test.json

生成的yolo数据集目录：

yolo_dataset
├── images
│   ├── 001.jpg
│   ├── 002.jpg
│   ├── 003.jpg
│   └── ...
│
├── labels
│   ├── 001.txt
│   ├── 002.txt
│   ├── 003.txt
│   └── ...
│
└── ImageSets
    ├── train.txt
    ├── val.txt
    └── test.txt

将coco的train，val，test分别转换为yolo的train，val，test
coco2yolo.py

# coco是x1,y1,w,h，yolo是x,y,w,h。 x1,y1是左上角坐标，x,y是中心坐标
import os
import shutil

from pycocotools.coco import COCO
from tqdm import trange

root_dir = os.getcwd()


# 将coco的bbox转换为yolo的bbox
def cocobbox2yolobbox(coco_box):
    x1, y1, w, h = coco_box
    x = x1 + w / 2
    y = y1 + h / 2
    return [x, y, w, h]


def coco2yolo(dataset_type, json_fp, origin_imgs_dir, save_dir):
    imgs_dir = os.path.join(save_dir, 'images')
    labels_dir = os.path.join(save_dir, 'labels')
    ImageSets_dir = os.path.join(save_dir, 'ImageSets')
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
        os.mkdir(imgs_dir)
        os.mkdir(labels_dir)
        os.mkdir(ImageSets_dir)

    text_data_fp = os.path.join(ImageSets_dir, dataset_type + '.txt')
    text_data = []

    coco = COCO(json_fp)
    imgs = coco.imgs
    img_ids = coco.getImgIds()

    m = len(img_ids)
    for i in trange(m):
        img_id = img_ids[i]
        filename = imgs[img_id]['file_name']
        text_data.append(os.path.join(imgs_dir, filename))
        txt_name = filename.split('.')[0] + ".txt"  # 对应的txt名字，与jpg一致
        f_txt = open(os.path.join(labels_dir, txt_name), 'w')

        ann_ids = coco.getAnnIds(imgIds=img_id)
        anns = coco.loadAnns(ann_ids)
        for ann in anns:
            bbox = cocobbox2yolobbox(ann["bbox"])
            f_txt.write("%s %s %s %s %s\n" % (ann['category_id'], bbox[0], bbox[1], bbox[2], bbox[3]))
        f_txt.close()

    # 将数据集写入文件
    with open(text_data_fp, 'w') as f:
        for line in text_data:
            f.write(line + '\n')
    print('labels create done.')

    for i in trange(m):
        img_id = img_ids[i]
        filename = imgs[img_id]['file_name']
        shutil.copy(os.path.join(origin_imgs_dir, filename), os.path.join(imgs_dir, filename))
    print('images copy done.')


def coco2yolo_type(dataset_type):
    save_dir = os.path.join(root_dir, 'yolo_dataset')
    if dataset_type == 'train':
        json_fp = os.path.join(root_dir, 'coco_small', 'annotations', 'train.json')
        imgs_dir = os.path.join(root_dir, 'coco_small', 'train')
        coco2yolo(dataset_type, json_fp, imgs_dir, save_dir)
    elif dataset_type == 'val':
        json_fp = os.path.join(root_dir, 'coco_small', 'annotations', 'val.json')
        imgs_dir = os.path.join(root_dir, 'coco_small', 'val')
        coco2yolo(dataset_type, json_fp, imgs_dir, save_dir)
    elif dataset_type == 'test':
        json_fp = os.path.join(root_dir, 'coco_small', 'annotations', 'test.json')
        imgs_dir = os.path.join(root_dir, 'coco_small', 'test')
        coco2yolo(dataset_type, json_fp, imgs_dir, save_dir)


if __name__ == '__main__':
    coco2yolo_type('train')
    coco2yolo_type('val')
    coco2yolo_type('test')

coco数据集划分为train，val，test

划分前目录结构：

coco
   +annotations
   +val2017

划分后目录结构：

coco_small
   +annotations
      +train.json
      +val.json
      +test.json
   +train
   +val
   +test

split_coco.py

# 划分coco数据集
# 将coco数据集分为train和val两个子集
import json
import os
import random
import shutil

from pycocotools.coco import COCO
from tqdm import trange

train_percentage = 0.7
val_percentage = 0.2
root_dir = os.getcwd()


def split_coco(json_fp, img_dir, save_dir):
    json_data = json.load(open(json_fp, 'r'))
    coco = COCO(json_fp)

    train_dir = os.path.join(save_dir, 'train')
    val_dir = os.path.join(save_dir, 'val')
    test_dir = os.path.join(save_dir, 'test')
    train_json = {'info': json_data['info'], 'licenses': json_data['licenses'], 'images': [], 'annotations': [],
                  'categories': json_data['categories']}
    val_json = {'info': json_data['info'], 'licenses': json_data['licenses'], 'images': [], 'annotations': [],
                'categories': json_data['categories']}
    test_json = {'info': json_data['info'], 'licenses': json_data['licenses'], 'images': [], 'annotations': [],
                 'categories': json_data['categories']}
    # 创建coco目录结构
    if os.path.exists(save_dir):
        shutil.rmtree(save_dir)
    os.makedirs(save_dir)
    os.mkdir(train_dir)
    os.mkdir(val_dir)
    os.mkdir(test_dir)
    os.mkdir(os.path.join(save_dir, 'annotations'))

    imgs = coco.imgs
    img_ids = coco.getImgIds()
    train_ids = random.sample(img_ids, int(len(img_ids) * train_percentage))
    val_ids = random.sample(list(set(img_ids) - set(train_ids)), int(len(img_ids) * val_percentage))

    m = len(img_ids)
    for i in trange(m):
        ann_ids = coco.getAnnIds(imgIds=img_ids[i])
        if img_ids[i] in train_ids:
            train_json['images'].append(imgs[img_ids[i]])
            shutil.copy(os.path.join(img_dir, imgs[img_ids[i]]['file_name']),
                        os.path.join(train_dir, imgs[img_ids[i]]['file_name']))
            for ann_id in ann_ids:
                train_json['annotations'].append(coco.anns[ann_id])
        elif img_ids[i] in val_ids:
            val_json['images'].append(imgs[img_ids[i]])
            shutil.copy(os.path.join(img_dir, imgs[img_ids[i]]['file_name']),
                        os.path.join(val_dir, imgs[img_ids[i]]['file_name']))
            for ann_id in ann_ids:
                val_json['annotations'].append(coco.anns[ann_id])
        else:
            test_json['images'].append(imgs[img_ids[i]])
            shutil.copy(os.path.join(img_dir, imgs[img_ids[i]]['file_name']),
                        os.path.join(test_dir, imgs[img_ids[i]]['file_name']))
            for ann_id in ann_ids:
                test_json['annotations'].append(coco.anns[ann_id])

    with open(os.path.join(save_dir, 'annotations', 'train.json'), 'x') as f:
        json.dump(train_json, f)
    with open(os.path.join(save_dir, 'annotations', 'val.json'), 'x') as f:
        json.dump(val_json, f)
    with open(os.path.join(save_dir, 'annotations', 'test.json'), 'x') as f:
        json.dump(test_json, f)


if __name__ == '__main__':
    json_fp = os.path.join(root_dir, 'coco', 'annotations', 'instances_val2017.json')
    img_dir = os.path.join(root_dir, 'coco', 'val2017')
    save_dir = os.path.join(root_dir, 'coco_small')

    split_coco(json_fp, img_dir, save_dir)
    print('done')

coco数据集图片检查

check_coco.py

# 检查coco数据集每张图片能否打开
import os
import cv2
from pycocotools.coco import COCO
from tqdm import trange
root_dir = os.getcwd()


def check_coco(json_fp, img_dir):
    coco = COCO(json_fp)
    imgs = coco.imgs
    img_ids = coco.getImgIds()

    for i in trange(len(img_ids)):
        img_fp = os.path.join(img_dir, imgs[img_ids[i]]['file_name'])
        img = cv2.imread(img_fp)
        try:
            img.shape
        except:
            print(img_fp)


if __name__ == '__main__':
    json_fp = os.path.join(root_dir, 'coco', 'annotations', 'instances_val2017.json')
    img_dir = os.path.join(root_dir, 'coco', 'val2017')
    check_coco(json_fp, img_dir)
    print('done')

coco数据集bbox可视化

visiual_coco.py

import os
import random
import cv2
from pycocotools.coco import COCO
from tqdm import trange
root_dir = os.getcwd()


# 检查转变后coco的json文件，坐标是否正确。
def visiual_coco(json_fp, images_dir):
    coco = COCO(json_fp)
    imgs = coco.imgs
    img_ids = coco.getImgIds()
    cats = coco.cats

    random.shuffle(img_ids)
    for i in trange(len(img_ids)):
        img_fp = os.path.join(images_dir, imgs[img_ids[i]]['file_name'])
        img = cv2.imread(img_fp)
        ann_ids = coco.getAnnIds(imgIds=img_ids[i])
        anns_ = coco.loadAnns(ann_ids)
        for ann in anns_:
            bbox = ann['bbox']
            left_top = (int(bbox[0]), int(bbox[1]))  # coco数据集中bbox的含义是x1,y1,w,h
            right_bottom = (int(bbox[0]) + int(bbox[2]), int(bbox[1]) + int(bbox[3]))
            cv2.rectangle(img, left_top, right_bottom, (0, 255, 0), 2)  # 图像，左上角，右下坐标，颜色，粗细
            cv2.putText(img, cats[ann['category_id']]['name'], left_top, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

        cv2.imshow('image', img)
        cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    json_fp = os.path.join(root_dir, 'coco', 'annotations', 'instances_val2017.json')
    images_dir = os.path.join(root_dir, 'coco', 'val2017')
    visiual_coco(json_fp, images_dir)

标签：xml,join,img,voc,coco,path,os,dir
来源： https://www.cnblogs.com/gy77/p/16473428.html

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。