| |
| |
| """ |
| 从 nyu_depth_v2_labeled.mat 导出 795 张训练集（trainNdxs）图片。 |
| |
| 划分来源与 export_nyu654_test.py 相同：BTS utils/splits.mat。 |
| trainNdxs 与 testNdxs（654）无重叠、并集覆盖 1..1449，可以与测试集严格隔离， |
| 用于量化校准或训练。 |
| |
| 只导出 RGB；GT 仍从 mat 按 split.json 的 mat_index 读取。 |
| |
| 用法： |
| ../../venv/bin/python src/tool/export_nyu_train795.py \ |
| --mat ../../datasets/nyu_depth_v2_labeled.mat \ |
| --splits datasets/nyu654_test/splits.mat \ |
| --out-dir datasets/nyu_train795 |
| """ |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import cv2 |
| import h5py |
| import numpy as np |
| import scipy.io as sio |
| from tqdm import tqdm |
|
|
|
|
| def decode_cell(f, ref): |
| return ''.join(chr(c) for c in np.asarray(f[ref]).ravel()) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='导出 NYU Depth V2 训练集 795 张') |
| parser.add_argument('--mat', required=True, help='nyu_depth_v2_labeled.mat') |
| parser.add_argument('--splits', required=True, help='含 trainNdxs 的 splits.mat') |
| parser.add_argument('--out-dir', required=True, help='输出目录') |
| args = parser.parse_args() |
|
|
| image_dir = Path(args.out_dir) / 'images' |
| image_dir.mkdir(parents=True, exist_ok=True) |
|
|
| ndxs = sio.loadmat(args.splits)['trainNdxs'].ravel().astype(int) |
| if len(ndxs) != 795: |
| raise ValueError(f'trainNdxs 应为 795 个，实际 {len(ndxs)}') |
|
|
| records = [] |
| with h5py.File(args.mat, 'r') as f: |
| for seq, mat_index in enumerate(tqdm(ndxs, desc='export', ncols=80)): |
| i = mat_index - 1 |
| scene = decode_cell(f, f['scenes'][0, i]) |
| rgb = np.asarray(f['images'][i]).transpose(1, 2, 0) |
| name = f'{seq:04d}_{scene}.png' |
| if not cv2.imwrite(str(image_dir / name), cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)): |
| raise RuntimeError(f'写入失败: {image_dir / name}') |
| records.append({'mat_index': int(i), 'scene': scene, 'file': name}) |
|
|
| with open(Path(args.out_dir) / 'split.json', 'w') as fp: |
| json.dump(records, fp, indent=2) |
| print(f'已导出 {len(records)} 张到 {image_dir}') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
