51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Fix COCO annotation class IDs to be 0-based instead of 1-based.
|
|
RF-DETR expects classes to start from 0 (background), but our annotations use 1-10.
|
|
"""
|
|
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
def fix_class_ids(coco_data):
|
|
"""Convert category IDs and annotation category_ids from 1-based to 0-based."""
|
|
# Fix categories: change id from 1-10 to 0-9
|
|
for cat in coco_data['categories']:
|
|
cat['id'] -= 1
|
|
|
|
# Fix annotations: change category_id from 1-10 to 0-9
|
|
for ann in coco_data['annotations']:
|
|
ann['category_id'] -= 1
|
|
|
|
return coco_data
|
|
|
|
def main():
|
|
dataset_dir = Path('dataset_coco')
|
|
|
|
for split in ['train', 'valid', 'test']:
|
|
ann_file = dataset_dir / split / '_annotations.coco.json'
|
|
|
|
if not ann_file.exists():
|
|
print(f"Warning: {ann_file} not found, skipping")
|
|
continue
|
|
|
|
print(f"Processing {ann_file}...")
|
|
|
|
# Load data
|
|
with open(ann_file, 'r') as f:
|
|
data = json.load(f)
|
|
|
|
# Fix class IDs
|
|
fixed_data = fix_class_ids(data)
|
|
|
|
# Save back
|
|
with open(ann_file, 'w') as f:
|
|
json.dump(fixed_data, f, indent=2)
|
|
|
|
print(f"✅ Fixed {len(data['annotations'])} annotations and {len(data['categories'])} categories")
|
|
|
|
print("🎉 All COCO annotations fixed for 0-based class indexing!")
|
|
|
|
if __name__ == '__main__':
|
|
main() |