#!/usr/bin/env python3
"""
Fix COCO annotations for RF-DETR compatibility by adding missing 'supercategory' field.
"""

import json
from pathlib import Path


def fix_coco_annotations(coco_file: Path):
    """Add supercategory field to categories in COCO annotation file."""
    print(f"Fixing {coco_file}...")

    # Load the COCO data
    with coco_file.open('r') as f:
        data = json.load(f)

    # Add supercategory field to all categories
    for category in data['categories']:
        if 'supercategory' not in category:
            category['supercategory'] = 'wood_defect'  # Default supercategory

    # Save the fixed data
    with coco_file.open('w') as f:
        json.dump(data, f, indent=2)

    print(f"✅ Fixed {len(data['categories'])} categories in {coco_file}")


def main():
    """Fix all COCO annotation files."""
    print("Fixing COCO annotations for RF-DETR compatibility...")

    coco_files = [
        Path("dataset_coco/train/_annotations.coco.json"),
        Path("dataset_coco/valid/_annotations.coco.json"),
        Path("dataset_coco/test/_annotations.coco.json")
    ]

    for coco_file in coco_files:
        if coco_file.exists():
            fix_coco_annotations(coco_file)
        else:
            print(f"⚠️  {coco_file} not found")

    print("\n🎉 All COCO annotations fixed for RF-DETR compatibility!")


if __name__ == "__main__":
    main()