| |
| |
| """ |
| NYU Depth V2 (labeled) 数据集查看与导出工具。 |
| |
| nyu_depth_v2_labeled.mat 是 MATLAB v7.3 (HDF5) 格式，scipy.io.loadmat 无法读取， |
| 这里统一用 h5py 按需读取，避免一次性载入 2.8GB。 |
| |
| 数据内容（N = 1449，图像为竖版 640x480，H=640, W=480，RGB 与深度已对齐）： |
| images [N, 3, 640, 480] uint8 RGB |
| depths [N, 640, 480] float32 已填补深度，单位米 |
| rawDepths [N, 640, 480] float32 原始深度，0 表示缺失 |
| labels [N, 640, 480] uint16 语义标签，0=未标注，1..894 对应 names |
| instances [N, 640, 480] uint8 实例标签，0=无实例 |
| names [1, 894] cell 类别名（1-based 下标即 label id） |
| namesToIds containers.Map MATLAB Map，h5py 无法直接解析 |
| scenes [1, N] cell 场景名，如 kitchen_0004 |
| sceneTypes [1, N] cell 场景类型，如 kitchen |
| rawRgbFilenames [1, N] cell 原始 RGB 文件名 |
| rawDepthFilenames [1, N] cell 原始深度文件名 |
| accelData [4, N] float32 加速度计数据 |
| |
| 用法（在 hs/DepthART 目录下执行）： |
| # 打印结构、场景/类别统计和深度数值范围（默认抽样 step=50） |
| ../../venv/bin/python src/tool/inspect_nyu_depth_v2.py --mat ../../datasets/nyu_depth_v2_labeled.mat |
| |
| # 全量扫描统计（较慢） |
| ... --mat ../../datasets/nyu_depth_v2_labeled.mat --full-stats |
| |
| # 列出全部 894 个类别名 |
| ... --mat ... --list-classes |
| |
| # 单样本导出：rgb.png / depth.npy / depth_vis.png / raw_depth.npy |
| ... --mat ... --index 0 --out-dir results/nyu_demo --with-labels |
| |
| # 批量导出，例如某场景全部帧 |
| ... --mat ... --export-dir results/nyu_export --scene kitchen_0004 --with-labels |
| |
| # 批量导出前 50 张 |
| ... --mat ... --export-dir results/nyu_export --count 50 |
| """ |
|
|
| import argparse |
| import collections |
| from pathlib import Path |
|
|
| import cv2 |
| import h5py |
| import numpy as np |
| from tqdm import tqdm |
|
|
| |
| DEPTH_VIS_RANGE = (0.5, 10.0) |
|
|
|
|
| def decode_cell(f, ref): |
| """把 MATLAB cell 中的 char 数组解码为 str。""" |
| return ''.join(chr(c) for c in np.asarray(f[ref]).ravel()) |
|
|
|
|
| def read_cell_list(dset): |
| return [decode_cell(dset.file, dset[0, i]) for i in range(dset.shape[1])] |
|
|
|
|
| def print_info(f, stride, full_stats): |
| print('== 变量 ==') |
| for k in f.keys(): |
| if k.startswith('#'): |
| continue |
| o = f[k] |
| line = f'{k:20s} shape={str(o.shape):24s} dtype={o.dtype}' |
| if o.compression: |
| line += f' compression={o.compression}' |
| print(line) |
|
|
| n = f['depths'].shape[0] |
| h, w = f['depths'].shape[1:] |
| names = read_cell_list(f['names']) |
| scenes = read_cell_list(f['scenes']) |
| scene_types = read_cell_list(f['sceneTypes']) |
|
|
| print('\n== 规模 ==') |
| print(f'样本数 {n}') |
| print(f'分辨率 W={w}, H={h}（data layout 为 [N, C, H, W]）') |
| print(f'场景数 {len(set(scenes))}') |
| print(f'场景类型数 {len(set(scene_types))}') |
| print(f'语义类别数 {len(names)}（label id 1..{len(names)}，0 为未标注）') |
| print('场景类型分布 ' + ', '.join(f'{k}:{v}' for k, v in |
| collections.Counter(scene_types).most_common())) |
|
|
| idx = range(0, n) if full_stats else range(0, n, stride) |
| idx = list(idx) |
| print(f'\n== 数值统计（{"全量" if full_stats else f"抽样 {len(idx)} 张, step={stride}"}） ==') |
| d_min, d_max = np.inf, -np.inf |
| raw_missing = 0 |
| raw_total = 0 |
| cls_count = collections.Counter() |
| inst_max = 0 |
| for i in tqdm(idx, desc='scan', ncols=80): |
| d = f['depths'][i] |
| r = f['rawDepths'][i] |
| d_min = min(d_min, float(d.min())) |
| d_max = max(d_max, float(d.max())) |
| raw_missing += int((r == 0).sum()) |
| raw_total += r.size |
| for c in np.unique(f['labels'][i]): |
| if c > 0: |
| cls_count[int(c)] += 1 |
| inst_max = max(inst_max, int(f['instances'][i].max())) |
|
|
| print(f'depths 范围 [{d_min:.4f}, {d_max:.4f}] m') |
| print(f'rawDepths 缺失 {raw_missing / raw_total * 100:.2f}% (值为 0)') |
| print(f'出现类别数 {len(cls_count)} / {len(names)}') |
| print(f'instances 最大 {inst_max}') |
| print('\n最常见类别 ' + ', '.join( |
| f'{names[c - 1]}({cnt})' for c, cnt in cls_count.most_common(15))) |
|
|
|
|
| def depth_to_vis(depth, raw=False): |
| lo, hi = DEPTH_VIS_RANGE |
| norm = np.clip((depth - lo) / (hi - lo), 0.0, 1.0) |
| vis = cv2.applyColorMap((norm * 255).astype(np.uint8), cv2.COLORMAP_TURBO) |
| if raw: |
| vis[depth <= 0] = 0 |
| return vis |
|
|
|
|
| def export_sample(f, index, out_dir, with_labels): |
| entry = out_dir / f'{index:05d}_{decode_cell(f, f["scenes"][0, index])}' |
| entry.mkdir(parents=True, exist_ok=True) |
|
|
| rgb = np.asarray(f['images'][index]).transpose(1, 2, 0) |
| depth = np.asarray(f['depths'][index], dtype=np.float32) |
| raw = np.asarray(f['rawDepths'][index], dtype=np.float32) |
|
|
| cv2.imwrite(str(entry / 'rgb.png'), cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)) |
| cv2.imwrite(str(entry / 'depth_vis.png'), depth_to_vis(depth)) |
| cv2.imwrite(str(entry / 'raw_depth_vis.png'), depth_to_vis(raw, raw=True)) |
| np.save(entry / 'depth.npy', depth) |
| np.save(entry / 'raw_depth.npy', raw) |
|
|
| if with_labels: |
| labels = np.asarray(f['labels'][index]) |
| np.save(entry / 'labels.npy', labels) |
| |
| rng = np.random.default_rng(0) |
| lut = rng.integers(40, 255, size=(65536, 3), dtype=np.uint8) |
| lut[0] = 0 |
| cv2.imwrite(str(entry / 'labels_vis.png'), lut[labels]) |
|
|
| return entry |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='查看/导出 NYU Depth V2 labeled 数据集') |
| parser.add_argument('--mat', required=True, help='nyu_depth_v2_labeled.mat 路径') |
| parser.add_argument('--list-classes', action='store_true', help='打印全部类别名及 id') |
| parser.add_argument('--index', type=int, help='导出单个样本的下标') |
| parser.add_argument('--out-dir', help='--index 模式的输出目录') |
| parser.add_argument('--export-dir', help='批量导出目录') |
| parser.add_argument('--count', type=int, help='批量导出数量，默认全部') |
| parser.add_argument('--step', type=int, default=1, help='批量导出间隔，默认 1') |
| parser.add_argument('--scene', help='只导出指定场景，如 kitchen_0004') |
| parser.add_argument('--with-labels', action='store_true', help='同时导出语义标签') |
| parser.add_argument('--full-stats', action='store_true', help='统计时扫描全部样本') |
| parser.add_argument('--stats-stride', type=int, default=50, help='抽样统计间隔，默认 50') |
| args = parser.parse_args() |
|
|
| with h5py.File(args.mat, 'r') as f: |
| if args.list_classes: |
| names = read_cell_list(f['names']) |
| for i, name in enumerate(names, 1): |
| print(f'{i:4d} {name}') |
| return |
|
|
| if args.index is not None: |
| out_dir = Path(args.out_dir or '.') |
| entry = export_sample(f, args.index, out_dir, args.with_labels) |
| print(f'已导出 {entry}') |
| return |
|
|
| if args.export_dir: |
| n = f['depths'].shape[0] |
| indices = list(range(0, n, args.step)) |
| if args.scene: |
| scenes = read_cell_list(f['scenes']) |
| indices = [i for i in indices if scenes[i] == args.scene] |
| print(f'场景 {args.scene} 命中 {len(indices)} 张') |
| if args.count: |
| indices = indices[:args.count] |
| out_dir = Path(args.export_dir) |
| for i in tqdm(indices, desc='export', ncols=80): |
| export_sample(f, i, out_dir, args.with_labels) |
| print(f'已导出 {len(indices)} 张到 {out_dir}') |
| return |
|
|
| print_info(f, args.stats_stride, args.full_stats) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
