Files
saw_mill_knot_detection/fix_class_ids.py

51 lines
1.4 KiB
Python
Raw Permalink Normal View History

#!/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()