File size: 3,102 Bytes
6272613 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
"""
Dataset loader for training RF-DETR models.
Supports COCO format annotations.
"""
import json
import os
from pathlib import Path
from typing import Dict, List, Tuple
def load_coco_annotations(annotation_path: str) -> Dict:
"""Load COCO format annotation file."""
if not os.path.exists(annotation_path):
raise FileNotFoundError(f"Annotation file not found: {annotation_path}")
with open(annotation_path, 'r') as f:
return json.load(f)
def get_image_paths(images_dir: str) -> List[str]:
"""Get all image file paths from directory."""
image_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
image_paths = []
for ext in image_extensions:
image_paths.extend(Path(images_dir).glob(f'*{ext}'))
image_paths.extend(Path(images_dir).glob(f'*{ext.upper()}'))
return sorted([str(p) for p in image_paths])
def validate_dataset(dataset_path: str) -> Tuple[bool, List[str]]:
"""Validate dataset structure and return issues."""
issues = []
dataset_path = Path(dataset_path)
if not dataset_path.exists():
issues.append(f"Dataset directory does not exist: {dataset_path}")
return False, issues
images_dir = dataset_path / "images"
annotations_file = dataset_path / "annotations" / "annotations.json"
if not images_dir.exists():
issues.append(f"Images directory missing: {images_dir}")
if not annotations_file.exists():
issues.append(f"Annotations file missing: {annotations_file}")
else:
try:
coco_data = load_coco_annotations(str(annotations_file))
if 'images' not in coco_data:
issues.append("COCO format: missing 'images' key")
if 'annotations' not in coco_data:
issues.append("COCO format: missing 'annotations' key")
if 'categories' not in coco_data:
issues.append("COCO format: missing 'categories' key")
except json.JSONDecodeError as e:
issues.append(f"Invalid JSON in annotations: {e}")
return len(issues) == 0, issues
def get_dataset_info(dataset_path: str) -> Dict:
"""Get summary information about the dataset."""
dataset_path = Path(dataset_path)
info = {
'path': str(dataset_path),
'images_count': 0,
'annotations_count': 0,
'categories_count': 0,
'categories': []
}
images_dir = dataset_path / "images"
annotations_file = dataset_path / "annotations" / "annotations.json"
if images_dir.exists():
info['images_count'] = len(get_image_paths(str(images_dir)))
if annotations_file.exists():
try:
coco_data = load_coco_annotations(str(annotations_file))
info['annotations_count'] = len(coco_data.get('annotations', []))
info['categories_count'] = len(coco_data.get('categories', []))
info['categories'] = [cat['name'] for cat in coco_data.get('categories', [])]
except Exception as e:
info['error'] = str(e)
return info
|