Blame for src/tool/inspect_nyu_depth_v2.py

#!/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()