File size: 4,527 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 96 97 98 99 100 101 102 103 104 105 106 | import os
import xml.etree.ElementTree as ET
import json
from collections import defaultdict
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/RUOD' # Replace with your actual path
anno_path = os.path.join(base_dir, 'RUOD_ANN', 'instances_train.json')
annotations = json.load(open(anno_path,"r"))
images_items = annotations["images"]
annos_items = annotations["annotations"]
cates_items = annotations["categories"]
category_dict = {}
for cate in cates_items:
category_dict[cate["id"]] = cate["name"]
category_dict_rev = {v: i for i, v in category_dict.items()}
novel_categories = ['corals', 'cuttlefish', 'turtle', 'jellyfish',]
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 = 10
caption_prefix = "An underwater image of "
thr = 15
os.makedirs(os.path.join(PROJECT_DIR, 'data', 'RUOD', 'metadatas', 'data_setting1'), exist_ok=True)
annos_dict = defaultdict(list)
for item in annos_items:
image_id = item["image_id"]
filename = images_items[image_id-1]['file_name']
annos_dict[filename].append(item["bbox"] + [item["category_id"]])
if __name__ == '__main__':
base_list, novel_list = [], [[] for i in range(len(novel_categories) + 1)]
for image_item in images_items:
dictin = {}
dictin['file_name'] = image_item['file_name']
width, height = image_item['width'], image_item['height']
categories_in_this_image = set()
categories, bndboxes, obndboxes= [], [], []
annos = annos_dict[image_item['file_name']]
for anno in annos:
xmin, ymin, w, h, category_id = anno
xmin, ymin, w, h = int(xmin), int(ymin), int(w), int(h)
xmin, ymin, xmax, ymax = xmin, ymin, xmin + w, ymin + h
xmin = xmin / width
ymin = ymin / height
xmax = xmax / width
ymax = ymax / height
categories.append(category_dict[category_id])
bndboxes.append([xmin, ymin, xmax, ymax])
obndboxes.append([xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax])
categories_in_this_image.add(category_id)
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
caption = [caption_prefix + ", ".join(categories)]
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/train/{image_item['file_name']}"
caplist = caption + categories
dictin["captions"] = caplist
dictin["bndboxes"] = bndboxes
dictin["obboxes"] = obndboxes
if is_novel:
if len(categories_in_this_image) > 1:
novel_list[-1].append(dictin.copy())
else:
novel_list[novel_dict_rev[category_dict[next(iter(categories_in_this_image))]]].append(dictin.copy())
else:
base_list.append(dictin.copy())
with open(os.path.join(PROJECT_DIR, "data/RUOD/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(len(novel_categories)):
with open(os.path.join(PROJECT_DIR, f"data/RUOD/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/RUOD/metadatas/data_setting1/train_novel_mixed.jsonl"), "w", encoding="utf-8") as f:
for item in novel_list[-1]:
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|