File size: 3,884 Bytes
f7c0974 | 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 | """Create reproducible real-portrait splits, excluding benchmark duplicates."""
from __future__ import annotations
import argparse
import hashlib
import json
from pathlib import Path
import cv2
def pixel_digest(path):
image = cv2.imread(str(path))
if image is None:
raise ValueError(f'Unreadable dataset image: {path}')
return hashlib.sha256(str(image.shape).encode() + image.tobytes()).hexdigest()
def prepare_split(root: Path, holdout: Path, output: Path, folders=('faces', 'pinterest'), seed=42, exclude_unreadable=False):
root = root.resolve()
excluded = {s.strip() for s in holdout.read_text(encoding='utf-8').splitlines() if s.strip()}
heldout_hashes = set()
for name in sorted(excluded):
path = (root / name).resolve()
if not path.is_relative_to(root):
raise ValueError('Holdout path escapes dataset root')
heldout_hashes.add(pixel_digest(path))
seen, rows, rejected = set(), [], []
duplicates = 0
for folder in folders:
for path in sorted((root / folder).rglob('*')):
if not path.is_file() or path.suffix.lower() not in {'.png', '.jpg', '.jpeg', '.webp'}:
continue
try:
digest = pixel_digest(path)
except ValueError:
if not exclude_unreadable:
raise
rejected.append({'path': path.relative_to(root).as_posix(), 'reason': 'unreadable'})
continue
if digest in heldout_hashes or digest in seen:
duplicates += 1
continue
seen.add(digest)
rows.append({'path': path.relative_to(root).as_posix(), 'pixels_sha256': digest})
if len(rows) < 100:
raise ValueError(f'Only {len(rows)} unique non-holdout portraits; refusing a toy split')
rows.sort(key=lambda row: hashlib.sha256(f"{seed}:{row['pixels_sha256']}".encode()).hexdigest())
validation_count = max(1, round(len(rows) * .05))
validation, train = rows[:validation_count], rows[validation_count:]
if output.exists() and any(output.iterdir()):
raise FileExistsError('Choose a new split directory to preserve existing splits')
output.mkdir(parents=True, exist_ok=True)
for name, items in [('train', train), ('validation', validation)]:
(output / f'{name}.txt').write_text(''.join(row['path'] + '\n' for row in items), encoding='utf-8')
report = {'schema_version': 1, 'seed': seed, 'domains': list(folders),
'holdout_count': len(excluded), 'excluded_or_duplicate_count': duplicates,
'train': train, 'validation': validation, 'rejected': rejected,
'review_status': 'pending_identity_and_near_duplicate_review'}
benchmark_manifest = holdout.parent / 'manifest.csv'
if benchmark_manifest.is_file():
report['benchmark_manifest_sha256'] = hashlib.sha256(benchmark_manifest.read_bytes()).hexdigest()
(output / 'holdout_paths.txt').write_text(holdout.read_text(encoding='utf-8'), encoding='utf-8')
(output / 'split.json').write_text(json.dumps(report, indent=2), encoding='utf-8')
print(f'Split: {len(train)} train, {len(validation)} validation; {len(excluded)} holdout sources excluded')
return report
if __name__ == '__main__':
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--root', type=Path, default=Path('models/CodeFormer/datasets/ffhq/ffhq_512'))
parser.add_argument('--holdout', type=Path, default=Path('benchmarks/holdout_paths.txt'))
parser.add_argument('--output', type=Path, required=True)
parser.add_argument('--exclude-unreadable', action='store_true', help='Record unreadable images as excluded; never delete them')
args = parser.parse_args()
prepare_split(args.root, args.holdout, args.output, exclude_unreadable=args.exclude_unreadable)
|