| |
| |
| """ |
| 从 NYU 训练集 795 张中按场景均匀抽取量化校准图片。 |
| |
| 抽样规则： |
| - 按场景类型（scene 去掉 _0001 后缀，如 kitchen、bedroom）分组； |
| - 每个场景至少 1 张，其余名额按场景图片数比例分配（最大余数法）； |
| - 场景内按等距索引选取，结果确定可复现。 |
| |
| 输出： |
| <out-dir>/images/*.png 校准图片 |
| <out-dir>/selection.json 每张图的来源 mat_index / scene / 源文件 |
| <out-dir>/<name>.tar 给 pulsar2 的校准包，内部为 <name>/xxx.png |
| |
| 用法： |
| ../../venv/bin/python src/tool/make_nyu_calibration.py \ |
| --source-dir datasets/nyu_train795 --count 48 \ |
| --out-dir datasets/nyu_calibration_48 |
| """ |
|
|
| import argparse |
| import json |
| import shutil |
| import subprocess |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
|
|
|
|
| def scene_type(scene): |
| return scene.rsplit('_', 1)[0] |
|
|
|
|
| def allocate(groups, count): |
| """最大余数法分配名额，每个场景至少 1 张。""" |
| total = sum(len(v) for v in groups.values()) |
| if count < len(groups): |
| raise ValueError(f'抽取数量 {count} 少于场景数 {len(groups)}，无法保证每场景 1 张') |
| quota = {k: 1 for k in groups} |
| left = count - len(groups) |
| exact = {k: left * len(v) / total for k, v in groups.items()} |
| for k, v in exact.items(): |
| quota[k] += int(np.floor(v)) |
| remain = count - sum(quota.values()) |
| order = sorted(exact, key=lambda k: exact[k] - np.floor(exact[k]), reverse=True) |
| for k in order[:remain]: |
| quota[k] += 1 |
| return quota |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='按场景均匀抽取 NYU 校准集') |
| parser.add_argument('--source-dir', required=True, help='export_nyu_train795.py 的输出目录') |
| parser.add_argument('--count', type=int, default=48, help='抽取张数，默认 48') |
| parser.add_argument('--out-dir', required=True, help='校准集输出目录') |
| parser.add_argument('--name', help='tar 包内目录名，默认取 out-dir 名') |
| parser.add_argument('--ext', choices=('png', 'jpg'), default='png', |
| help='校准图片格式，默认 png；pulsar2 读取异常时可改 jpg') |
| args = parser.parse_args() |
|
|
| source = Path(args.source_dir) |
| split = json.load(open(source / 'split.json')) |
| groups = defaultdict(list) |
| for record in split: |
| groups[scene_type(record['scene'])].append(record) |
|
|
| quota = allocate(groups, args.count) |
| selected = [] |
| for scene, records in sorted(groups.items()): |
| records = sorted(records, key=lambda r: r['file']) |
| idx = np.linspace(0, len(records) - 1, quota[scene]).round().astype(int) |
| selected.extend(records[i] for i in sorted(set(idx.tolist()))) |
|
|
| if len(selected) != args.count: |
| raise RuntimeError(f'实际选出 {len(selected)} 张，与目标 {args.count} 不符') |
|
|
| out_dir = Path(args.out_dir) |
| image_dir = out_dir / 'images' |
| image_dir.mkdir(parents=True, exist_ok=True) |
| for record in selected: |
| src = source / 'images' / record['file'] |
| if args.ext == 'png': |
| dst = image_dir / record['file'] |
| shutil.copy2(src, dst) |
| else: |
| dst = image_dir / (Path(record['file']).stem + '.jpg') |
| image = cv2.imread(str(src), cv2.IMREAD_COLOR) |
| if image is None: |
| raise RuntimeError(f'读取失败: {src}') |
| cv2.imwrite(str(dst), image, [cv2.IMWRITE_JPEG_QUALITY, 95]) |
|
|
| with open(out_dir / 'selection.json', 'w') as fp: |
| json.dump(selected, fp, indent=2) |
|
|
| name = args.name or out_dir.name |
| tar_path = out_dir / f'{name}.tar' |
| subprocess.run( |
| ['tar', '-cf', str(tar_path), '--transform', f's,^,{name}/,', '-C', str(image_dir)] |
| + sorted(p.name for p in image_dir.iterdir()), |
| check=True) |
|
|
| print(f'已抽取 {len(selected)} 张，覆盖 {len(groups)} 个场景类型') |
| for scene in sorted(groups): |
| print(f' {scene:<16} {quota[scene]:>3} 张 / 共 {len(groups[scene])}') |
| print(f'校准包: {tar_path}') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
