File size: 4,180 Bytes
c881b77 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import os
import xml.etree.ElementTree as ET
import json
PROJECT_DIR = os.getenv('DSP_PROJECT_DIR', '/path/to/DSP_PROJECT_DIR') # Set this manually if the environment variable is unavailable
base_dir = '/path/to/DIOR-VOC/Annotations/Horizontal_Bounding_Boxes' # Replace with your actual path
category_list = [
'vehicle', 'baseballfield', 'groundtrackfield', 'windmill', 'bridge',
'overpass', 'ship', 'airplane', 'tenniscourt', 'airport',
'expressway-service-area', 'basketballcourt', 'stadium', 'storagetank', 'chimney',
'dam', 'expressway-toll-station', 'golffield', 'trainstation', 'harbor'
]
category_dict_rev = {v: i for i, v in enumerate(category_list)}
novel_categories = [
'windmill', 'airport', 'chimney', 'dam', 'trainstation'
]
novel_ids = set([category_dict_rev[cate] for cate in novel_categories])
novel_dict_rev = {v: i for i, v in enumerate(novel_categories)}
num_classes = 20
width_height = 800
caption_prefix = "The aerial image features a city with "
thr = 15
os.makedirs(os.path.join(PROJECT_DIR, 'data', 'DIOR', 'metadatas', 'data_setting1'), exist_ok=True)
if __name__ == '__main__':
base_list, novel_list = [], [[], [], [], [], [], []]
filenames = sorted(os.listdir(base_dir))[:5862]
for filename in filenames:
dictin = {}
root = ET.parse(os.path.join(base_dir, filename)).getroot()
categories_in_this_image = set()
categories, bndboxes, obndboxes= [], [], []
for object in root.findall('object'):
category = object.find('name').text.lower()
category_id = category_dict_rev[category]
xmin, ymin, xmax, ymax = [int(child.text) / width_height for child in object.find('bndbox')]
categories.append(category)
bndboxes.append([xmin, ymin, xmax, ymax])
obndboxes.append([xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax])
categories_in_this_image.add(category_id)
# caption put to line 58
caption = [caption_prefix + ", ".join(categories)]
is_novel = False
if categories_in_this_image & set(novel_ids):
is_novel = True
tmp_categories, tmp_bndboxes, tmp_obndboxes= [], [], []
for i in range(len(categories)):
if categories[i] in novel_categories:
tmp_categories.append(categories[i])
tmp_bndboxes.append(bndboxes[i])
tmp_obndboxes.append(obndboxes[i])
categories, bndboxes, obndboxes = tmp_categories, tmp_bndboxes, tmp_obndboxes
if len(categories) > thr:
categories = categories[:thr]
bndboxes = bndboxes[:thr]
obndboxes = obndboxes[:thr]
while len(categories) < thr:
categories.append("")
bndboxes.append([0,0,0,0])
obndboxes.append([0,0,0,0,0,0,0,0])
dictin["file_name"] = f"../../images/{filename[:5]}.jpg"
caplist = caption + categories
dictin["captions"] = caplist
dictin["bndboxes"] = bndboxes
dictin["obboxes"] = obndboxes
if is_novel:
if len(categories_in_this_image) > 1:
novel_list[5].append(dictin.copy())
else:
novel_list[novel_dict_rev[category_list[next(iter(categories_in_this_image))]]].append(dictin.copy())
else:
base_list.append(dictin.copy())
with open(os.path.join(PROJECT_DIR, "data/DIOR/metadatas/data_setting1/train_base.jsonl"), "w", encoding="utf-8") as f:
for item in base_list:
f.write(json.dumps(item, ensure_ascii=False) + "\n")
for i in range(5):
with open(os.path.join(PROJECT_DIR, f"data/DIOR/metadatas/data_setting1/train_novel_{novel_categories[i]}.jsonl"), "w", encoding="utf-8") as f:
for item in novel_list[i]:
f.write(json.dumps(item, ensure_ascii=False) + "\n")
with open(os.path.join(PROJECT_DIR, "data/DIOR/metadatas/data_setting1/train_novel_mixed.jsonl"), "w", encoding="utf-8") as f:
for item in novel_list[5]:
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|