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