DIOR数据集转化为COCO格式

DIOR数据集转化为COCO格式

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

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

第一个参数是你电脑里上图的路径，第二个参数是你想输出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")