Spaces:
Sleeping
Sleeping
| """Create a deterministic, exact-duplicate-safe train/test image split.""" | |
| import argparse | |
| import hashlib | |
| import json | |
| import random | |
| import shutil | |
| from collections import defaultdict | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| AI_ROOT = Path(__file__).resolve().parents[2] | |
| DEFAULT_SOURCE = AI_ROOT.parent / 'dataset' / 'mission_dataset' | |
| SUPPORTED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.webp'} | |
| def sha256_file(path): | |
| digest = hashlib.sha256() | |
| with path.open('rb') as source: | |
| for chunk in iter(lambda: source.read(1024 * 1024), b''): | |
| digest.update(chunk) | |
| return digest.hexdigest() | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument('--source', type=Path, default=DEFAULT_SOURCE) | |
| parser.add_argument('--output', type=Path, required=True) | |
| parser.add_argument('--dataset-version', required=True) | |
| parser.add_argument('--train-ratio', type=float, default=0.8) | |
| parser.add_argument('--seed', type=int, default=17017) | |
| return parser.parse_args() | |
| def scan_unique_images(source): | |
| by_class = defaultdict(list) | |
| hash_labels = defaultdict(set) | |
| duplicate_count = 0 | |
| seen_hashes = set() | |
| for class_path in sorted(path for path in source.iterdir() if path.is_dir()): | |
| for image_path in sorted(class_path.rglob('*')): | |
| if not image_path.is_file() or image_path.suffix.lower() not in SUPPORTED_EXTENSIONS: | |
| continue | |
| content_hash = sha256_file(image_path) | |
| hash_labels[content_hash].add(class_path.name) | |
| if content_hash in seen_hashes: | |
| duplicate_count += 1 | |
| continue | |
| seen_hashes.add(content_hash) | |
| by_class[class_path.name].append((image_path, content_hash)) | |
| conflicts = { | |
| content_hash: sorted(labels) | |
| for content_hash, labels in hash_labels.items() | |
| if len(labels) > 1 | |
| } | |
| if conflicts: | |
| raise ValueError( | |
| f'{len(conflicts)} byte-identical images have conflicting class labels; fix them before splitting.' | |
| ) | |
| return by_class, duplicate_count | |
| def main(): | |
| args = parse_args() | |
| source = args.source.resolve() | |
| output = args.output.resolve() | |
| if not source.is_dir(): | |
| raise FileNotFoundError(f'Source dataset not found: {source}') | |
| if output.exists(): | |
| raise FileExistsError(f'Output already exists; choose a new versioned directory: {output}') | |
| if not 0.5 <= args.train_ratio < 1: | |
| raise ValueError('Train ratio must be at least 0.5 and less than 1.0.') | |
| by_class, duplicate_count = scan_unique_images(source) | |
| if not by_class: | |
| raise ValueError('No supported source images were found.') | |
| rng = random.Random(args.seed) | |
| manifest = { | |
| 'datasetVersion': args.dataset_version, | |
| 'createdAtUtc': datetime.now(timezone.utc).isoformat(), | |
| 'source': str(source), | |
| 'output': str(output), | |
| 'seed': args.seed, | |
| 'trainRatio': args.train_ratio, | |
| 'exactDuplicatesRemoved': duplicate_count, | |
| 'classes': {}, | |
| } | |
| copy_plan = [] | |
| for class_name in sorted(by_class): | |
| records = list(by_class[class_name]) | |
| rng.shuffle(records) | |
| split_index = int(len(records) * args.train_ratio) | |
| train_records = records[:split_index] | |
| test_records = records[split_index:] | |
| if not train_records or not test_records: | |
| raise ValueError(f'Class {class_name} does not have enough unique images for both splits.') | |
| manifest['classes'][class_name] = { | |
| 'uniqueSourceImages': len(records), | |
| 'trainImages': len(train_records), | |
| 'testImages': len(test_records), | |
| } | |
| copy_plan.extend(('train', class_name, *record) for record in train_records) | |
| copy_plan.extend(('test', class_name, *record) for record in test_records) | |
| for split_name, class_name, source_path, content_hash in copy_plan: | |
| destination_directory = output / split_name / class_name | |
| destination_directory.mkdir(parents=True, exist_ok=True) | |
| destination_name = f'{content_hash[:12]}-{source_path.name}' | |
| shutil.copy2(source_path, destination_directory / destination_name) | |
| train_hashes = {record[3] for record in copy_plan if record[0] == 'train'} | |
| test_hashes = {record[3] for record in copy_plan if record[0] == 'test'} | |
| overlap = train_hashes.intersection(test_hashes) | |
| if overlap: | |
| raise RuntimeError('Internal error: exact hashes overlap between generated splits.') | |
| manifest['totals'] = { | |
| 'trainImages': len(train_hashes), | |
| 'testImages': len(test_hashes), | |
| 'exactCrossSplitDuplicates': 0, | |
| } | |
| (output / 'dataset_manifest.json').write_text( | |
| json.dumps(manifest, indent=2, sort_keys=True), | |
| encoding='utf-8', | |
| ) | |
| print(json.dumps(manifest, indent=2, sort_keys=True)) | |
| return 0 | |
| if __name__ == '__main__': | |
| raise SystemExit(main()) | |