Technical_Draw_Detection / Preprocess_datasets /convert_annotations_tech.py
TVQuyet05
init
248b460
import json
import os
def process_annotations_tech(file_path):
print(f"Processing {file_path}...")
with open(file_path, 'r', encoding='utf-8') as f:
data = json.load(f)
# Define new categories
new_categories = [
{'id': 0, 'name': 'PartDrawing', 'supercategory': 'none'},
{'id': 1, 'name': 'Note', 'supercategory': 'none'},
{'id': 2, 'name': 'Table', 'supercategory': 'none'}
]
# Create mapping from old category id to new category id
id_mapping = {}
for cat in data['categories']:
old_name = cat['name']
old_id = cat['id']
if old_name == 'figure':
id_mapping[old_id] = 0
elif old_name == 'note':
id_mapping[old_id] = 1
elif old_name == 'table':
id_mapping[old_id] = 2
# Update annotations
new_annotations = []
for ann in data['annotations']:
old_cat_id = ann['category_id']
if old_cat_id in id_mapping:
ann['category_id'] = id_mapping[old_cat_id]
new_annotations.append(ann)
data['annotations'] = new_annotations
data['categories'] = new_categories
with open(file_path, 'w', encoding='utf-8') as f:
json.dump(data, f)
print(f"Finished processing {file_path}")
if __name__ == "__main__":
base_dir = r"d:\Technical_Draw_Detection\Datasets\Data_tech"
# Data_tech only has train and valid according to the structure
files = [
os.path.join(base_dir, "train", "_annotations.coco.json"),
os.path.join(base_dir, "valid", "_annotations.coco.json")
]
for file_path in files:
if os.path.exists(file_path):
process_annotations_tech(file_path)
else:
print(f"File not found: {file_path}")