#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 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) # CHW -> HWC, RGB 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) # 语义标签着色：按 label id 的哈希色，0 保持黑色 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()