File size: 2,601 Bytes
efaa3d7 | 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 | #!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从 nyu_depth_v2_labeled.mat 导出官方 Eigen test split 的 654 张 RGB 图片。
划分来源：BTS 仓库 utils/splits.mat 的 testNdxs（1-based，共 654 个），
与 train_test_inputs/nyudepthv2_test_files_with_gt.txt 的 654 行逐场景帧数一致，
已核对无重叠、并集覆盖 1..1449。
只导出 RGB，GT 不落盘：评测脚本按 split.json 里的 mat_index 直接从 mat 读取，
避免重复存储。
用法：
# 导出 654 张到 datasets/nyu654_test/
../../venv/bin/python src/tool/export_nyu654_test.py \
--mat ../../datasets/nyu_depth_v2_labeled.mat \
--splits datasets/nyu654_test/splits.mat \
--out-dir datasets/nyu654_test
输出：
<out-dir>/images/{序号:04d}_{scene}.png
<out-dir>/split.json [{mat_index, scene, file}, ...]
"""
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 Eigen test split 654 张')
parser.add_argument('--mat', required=True, help='nyu_depth_v2_labeled.mat')
parser.add_argument('--splits', required=True, help='含 testNdxs 的 splits.mat')
parser.add_argument('--out-dir', required=True, help='输出目录')
args = parser.parse_args()
out_dir = Path(args.out_dir)
image_dir = out_dir / 'images'
image_dir.mkdir(parents=True, exist_ok=True)
ndxs = sio.loadmat(args.splits)['testNdxs'].ravel().astype(int)
if len(ndxs) != 654:
raise ValueError(f'testNdxs 应为 654 个，实际 {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 # matlab 1-based
scene = decode_cell(f, f['scenes'][0, i])
rgb = np.asarray(f['images'][i]).transpose(1, 2, 0) # CHW -> HWC
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(out_dir / 'split.json', 'w') as fp:
json.dump(records, fp, indent=2)
print(f'已导出 {len(records)} 张到 {image_dir}，划分文件 {out_dir / "split.json"}')
if __name__ == '__main__':
main()
|