| | from pathlib import Path |
| |
|
| | class_name = 'strawberry' |
| |
|
| | def rewrite_single_class_data_yaml(dataset_dir, class_name='strawberry'): |
| | dataset_dir = Path(dataset_dir) |
| | data_yaml_path = dataset_dir / 'data.yaml' |
| | if not data_yaml_path.exists(): |
| | print('⚠️ data.yaml not found, skipping rewrite.') |
| | return |
| |
|
| | train_path = dataset_dir / 'train' / 'images' |
| | val_path = dataset_dir / 'valid' / 'images' |
| | test_path = dataset_dir / 'test' / 'images' |
| |
|
| | content_lines = [ |
| | '# Strawberry-only dataset', |
| | f'train: {train_path}', |
| | f'val: {val_path}', |
| | f"test: {test_path if test_path.exists() else ''}", |
| | '', |
| | 'nc: 1', |
| | f"names: ['{class_name}']", |
| | ] |
| |
|
| | data_yaml_path.write_text('\n'.join(content_lines) + '\n') |
| | print(f"✅ data.yaml updated for single-class training ({class_name}).") |
| |
|
| |
|
| | def enforce_single_class_dataset(dataset_dir, target_class=0, class_name='strawberry'): |
| | dataset_dir = Path(dataset_dir) |
| | stats = {'split_kept': {}, 'labels_removed': 0, 'images_removed': 0} |
| | allowed_ext = ['.jpg', '.jpeg', '.png'] |
| |
|
| | for split in ['train', 'valid', 'test']: |
| | labels_dir = dataset_dir / split / 'labels' |
| | images_dir = dataset_dir / split / 'images' |
| | if not labels_dir.exists(): |
| | continue |
| |
|
| | kept = 0 |
| | for label_path in labels_dir.glob('*.txt'): |
| | kept_lines = [] |
| | for raw_line in label_path.read_text().splitlines(): |
| | line = raw_line.strip() |
| | if not line: |
| | continue |
| | parts = line.split() |
| | if not parts: |
| | continue |
| | try: |
| | class_id = int(parts[0]) |
| | except ValueError: |
| | continue |
| | if class_id == target_class: |
| | kept_lines.append(line) |
| |
|
| | if kept_lines: |
| | label_path.write_text('\n'.join(kept_lines) + '\n') |
| | kept += len(kept_lines) |
| | else: |
| | label_path.unlink() |
| | stats['labels_removed'] += 1 |
| | for ext in allowed_ext: |
| | candidate = images_dir / f"{label_path.stem}{ext}" |
| | if candidate.exists(): |
| | candidate.unlink() |
| | stats['images_removed'] += 1 |
| | break |
| |
|
| | stats['split_kept'][split] = kept |
| |
|
| | rewrite_single_class_data_yaml(dataset_dir, class_name) |
| |
|
| | print('\n🍓 Strawberry-only filtering summary:') |
| | for split, count in stats['split_kept'].items(): |
| | print(f" {split}: {count} annotations kept") |
| | print(f" Label files removed: {stats['labels_removed']}") |
| | print(f" Images removed (non-strawberry or empty labels): {stats['images_removed']}") |
| |
|
| | return stats |
| |
|
| |
|
| | if __name__ == "__main__": |
| | |
| | dataset_path = "path/to/your/dataset" |
| |
|
| | if dataset_path and Path(dataset_path).exists(): |
| | strawberry_stats = enforce_single_class_dataset(dataset_path, target_class=0, class_name=class_name) |
| | else: |
| | print("Please set a valid dataset_path variable.") |
