DIOR数据集转化为COCO格式

DIOR数据集转化为COCO格式

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

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

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

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

import os
import cv2
from tqdm import tqdm
import json
import xml.dom.minidomcategory_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 datasetoutput_dir:the path write the coco form json file"""annotations_path = dirnamelist_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 = 0categories = [{"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_pathelse:f = open(namelist_path + "/" + "val.txt", "r")images_path = trainval_images_pathfor name in tqdm(f.readlines()):# image partimage = {}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_nameimage["height"] = heightimage["width"] = widthimage["id"] = nameimages.append(image)# anno partdom = xml.dom.minidom.parse(dir + "/" + annotation_name)root_data = dom.documentElementfor i in range(len(root_data.getElementsByTagName('name'))):annotation = {}category = root_data.getElementsByTagName('name')[i].firstChild.datatop_left_x = root_data.getElementsByTagName('xmin')[i].firstChild.datatop_left_y = root_data.getElementsByTagName('ymin')[i].firstChild.dataright_bottom_x = root_data.getElementsByTagName('xmax')[i].firstChild.dataright_bottom_y = root_data.getElementsByTagName('ymax')[i].firstChild.databbox = [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"] = nameannotation["bbox"] = bboxannotation["category_id"] = category_list.index(category)annotation["id"] = id_numannotation["iscrowd"] = 0annotation["segmentation"] = []annotation["area"] = bbox[2] * bbox[3]id_num += 1annotations.append(annotation)dataset_dict = {}dataset_dict["images"] = imagesdataset_dict["annotations"] = annotationsdataset_dict["categories"] = categoriesjson_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 Nonedef 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] / 2y_top_left = boxes[1] - boxes[3] / 2x_bottom_right = boxes[0] + boxes[2] / 2y_bottom_right = boxes[1] + boxes[3] / 2return [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] / 2y_top_left = boxes[1] - boxes[3] / 2width = 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] = 0if coord[1] < 0:coord[1] = 0if coord[2] > width:coord[2] = widthif coord[3] > height:coord[3] = heightreturn coordif __name__ == '__main__':convert_to_cocodetection(r"path to your DIOR dataset", r"the path you want to write the coco format json file")