欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页  >  IT编程

DIOR数据集转化为COCO格式

程序员文章站 2022-06-22 16:29:09
DIOR数据集转化为COCO格式为了实验方便,需要将DIOR数据集转化为COCO格式,在此将代码共享,希望能帮助到同样研究方向的人。解压DIOR数据集的压缩文件之后,你的路径应该是这样的:第一个参数是你电脑里上图的路径,第二个参数是你想输出COCO格式文件的路径DIOR缺少很多COCO格式的数据,所以缺少的项都为空,代码如下import osimport cv2from tqdm import tqdmimport jsonimport xml.dom.minidomc...

DIOR数据集转化为COCO格式

为了实验方便,需要将DIOR数据集转化为COCO格式,在此将代码共享,希望能帮助到同样研究方向的人。

解压DIOR数据集的压缩文件之后,你的路径应该是这样的:

DIOR数据集转化为COCO格式

第一个参数是你电脑里上图的路径,第二个参数是你想输出COCO格式文件的路径

DIOR缺少很多COCO格式的数据,所以缺少的项都为空,代码如下

import os
import cv2
from tqdm import tqdm
import json
import xml.dom.minidom

category_list = ['airplane', 'airport', 'baseballfield', 'basketballcourt', 'bridge', 'chimney', 'dam',
              'Expressway-Service-area', 'Expressway-toll-station', 'golffield', 'groundtrackfield', 'harbor',
              'overpass', 'ship', 'stadium', 'storagetank', 'tenniscourt', 'trainstation', 'vehicle', 'windmill']


def convert_to_cocodetection(dir, output_dir):
    """
    input:
        dir:the path to DIOR dataset
        output_dir:the path write the coco form json file
    """
    annotations_path = dir
    namelist_path = os.path.join(dir, "Main")
    trainval_images_path = os.path.join(dir, "JPEGImages-trainval")
    test_images_path = os.path.join(dir, "JPEGImages-test")
    id_num = 0
    categories = [{"id": 0, "name": "Airplane"},
                  {"id": 1, "name": "Airport"},
                  {"id": 2, "name": "Baseball field"},
                  {"id": 3, "name": "Basketball court"},
                  {"id": 4, "name": "Bridge"},
                  {"id": 5, "name": "Chimney"},
                  {"id": 6, "name": "Dam"},
                  {"id": 7, "name": "Expressway service area"},
                  {"id": 8, "name": "Expressway toll station"},
                  {"id": 9, "name": "Golf course"},
                  {"id": 10, "name": "Ground track field"},
                  {"id": 11, "name": "Harbor"},
                  {"id": 12, "name": "Overpass"},
                  {"id": 13, "name": "Ship"},
                  {"id": 14, "name": "Stadium"},
                  {"id": 15, "name": "Storage tank"},
                  {"id": 16, "name": "Tennis court"},
                  {"id": 17, "name": "Train station"},
                  {"id": 18, "name": "Vehicle"},
                  {"id": 19, "name": "Wind mill"},
                  ]
    for mode in ["train", "val"]:
        images = []
        annotations = []
        print(f"start loading {mode} data...")
        if mode == "train":
            f = open(namelist_path + "/" + "train.txt", "r")
            images_path = trainval_images_path
        else:
            f = open(namelist_path + "/" + "val.txt", "r")
            images_path = trainval_images_path
        for name in tqdm(f.readlines()):
            # image part
            image = {}

            name = name.replace("\n", "")
            image_name = name + ".jpg"
            annotation_name = name + ".xml"
            height, width = cv2.imread(images_path + "/" + image_name).shape[:2]
            image["file_name"] = image_name
            image["height"] = height
            image["width"] = width
            image["id"] = name
            images.append(image)
            # anno part
            dom = xml.dom.minidom.parse(dir + "/" + annotation_name)
            root_data = dom.documentElement
            for i in range(len(root_data.getElementsByTagName('name'))):
                annotation = {}
                category = root_data.getElementsByTagName('name')[i].firstChild.data
                top_left_x = root_data.getElementsByTagName('xmin')[i].firstChild.data
                top_left_y = root_data.getElementsByTagName('ymin')[i].firstChild.data
                right_bottom_x = root_data.getElementsByTagName('xmax')[i].firstChild.data
                right_bottom_y = root_data.getElementsByTagName('ymax')[i].firstChild.data
                bbox = [top_left_x, top_left_y, right_bottom_x, right_bottom_y]
                bbox = [int(i) for i in bbox]
                bbox = xyxy_to_xywh(bbox)
                annotation["image_id"] = name
                annotation["bbox"] = bbox
                annotation["category_id"] = category_list.index(category)
                annotation["id"] = id_num
                annotation["iscrowd"] = 0
                annotation["segmentation"] = []
                annotation["area"] = bbox[2] * bbox[3]
                id_num += 1
                annotations.append(annotation)
        dataset_dict = {}
        dataset_dict["images"] = images
        dataset_dict["annotations"] = annotations
        dataset_dict["categories"] = categories
        json_str = json.dumps(dataset_dict)
        with open(f'{output_dir}/DIOR_{mode}_coco.json', 'w') as json_file:
            json_file.write(json_str)
    print("json file write done...")


def get_test_namelist(dir, out_dir):
    full_path = out_dir + "/" + "test.txt"
    file = open(full_path, 'w')
    for name in tqdm(os.listdir(dir)):
        name = name.replace(".txt", "")
        file.write(name + "\n")
    file.close()
    return None


def centerxywh_to_xyxy(boxes):
    """
    args:
        boxes:list of center_x,center_y,width,height,
    return:
        boxes:list of x,y,x,y,cooresponding to top left and bottom right
    """
    x_top_left = boxes[0] - boxes[2] / 2
    y_top_left = boxes[1] - boxes[3] / 2
    x_bottom_right = boxes[0] + boxes[2] / 2
    y_bottom_right = boxes[1] + boxes[3] / 2
    return [x_top_left, y_top_left, x_bottom_right, y_bottom_right]


def centerxywh_to_topleftxywh(boxes):
    """
    args:
        boxes:list of center_x,center_y,width,height,
    return:
        boxes:list of x,y,x,y,cooresponding to top left and bottom right
    """
    x_top_left = boxes[0] - boxes[2] / 2
    y_top_left = boxes[1] - boxes[3] / 2
    width = boxes[2]
    height = boxes[3]
    return [x_top_left, y_top_left, width, height]


def xyxy_to_xywh(boxes):
    width = boxes[2] - boxes[0]
    height = boxes[3] - boxes[1]
    return [boxes[0], boxes[1], width, height]


def clamp(coord, width, height):
    if coord[0] < 0:
        coord[0] = 0
    if coord[1] < 0:
        coord[1] = 0
    if coord[2] > width:
        coord[2] = width
    if coord[3] > height:
        coord[3] = height
    return coord


if __name__ == '__main__':
    convert_to_cocodetection(r"path to your DIOR dataset", r"the path you want to write the coco format json file")

 

本文地址:https://blog.csdn.net/qq_20793791/article/details/110881417