50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
|
#!/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()
|