如何用Python加载NYU Depth V2数据集？

Question

**问题描述：** 在使用Python进行深度估计或计算机视觉任务时，如何正确加载和解析NYU Depth V2数据集？该数据集通常以.mat文件格式存储，包含RGB图像、深度图及场景标签等信息。开发者常遇到的问题包括：如何读取.mat文件中的结构化数据、如何提取并可视化RGB与深度图像、如何处理不同版本的MATLAB文件（如v7.3与早期版本），以及如何划分训练集与测试集。此外，由于数据量较大，内存管理与数据增强策略也是常见挑战。本文将介绍如何使用Python（如`scipy`, `h5py`, `numpy`, `opencv`等库）高效加载、处理并预览NYU Depth V2数据集，为后续模型训练打下基础。

Answer 1

一、概述与数据集结构解析

NYU Depth V2 是一个广泛用于深度估计和场景理解任务的室内RGB-D图像数据集。该数据集包含来自Microsoft Kinect采集的约1449张对齐的RGB图像及其对应的深度图，并附带场景类别标签。数据以MATLAB格式（.mat）存储，开发者在使用Python进行加载时需特别注意其版本兼容性。

常见的.mat文件版本包括：

• v4: 旧版MATLAB矩阵文件，可由scipy.io.loadmat()直接读取。
• v6, v7: 支持更复杂的数据类型，同样适用于scipy.io。
• v7.3及以上：基于HDF5格式，必须使用h5py库进行读取。

开发者需要首先确认所下载的NYU Depth V2 .mat文件版本，以便选择合适的解析方法。

二、加载与解析NYU Depth V2数据

由于官方发布的NYU Depth V2数据集通常为v7.3格式，推荐使用h5py进行读取。以下是一个基础示例代码：

import h5py
import numpy as np

file_path = 'nyu_depth_v2_labeled.mat'
with h5py.File(file_path, 'r') as f:
    # 查看所有顶层变量名
    print("Top-level keys:", list(f.keys()))

    # 获取RGB图像和深度图数据
    images = f['images']  # shape: (N, 3, H, W)
    depths = f['depths']  # shape: (N, H, W)

    # 取出第一张图像并转置为(H, W, C)格式
    rgb_img = np.transpose(images[0], (2, 1, 0))
    depth_img = depths[0]
    

注意：HDF5中存储的图像通常是通道优先（Channel-first）格式，因此需要使用np.transpose调整维度顺序。

三、可视化RGB与深度图像

为了验证数据是否正确加载，可以使用OpenCV或Matplotlib进行可视化。以下为使用Matplotlib的示例：

import matplotlib.pyplot as plt

plt.figure(figsize=(10, 5))

# 显示RGB图像
plt.subplot(1, 2, 1)
plt.imshow(rgb_img.astype(np.uint8))
plt.title('RGB Image')
plt.axis('off')

# 显示深度图（伪彩色）
plt.subplot(1, 2, 2)
plt.imshow(depth_img, cmap='jet')
plt.colorbar()
plt.title('Depth Map')
plt.axis('off')

plt.show()
    

深度图通常为单通道浮点型数值，代表每个像素点距离摄像头的距离（单位：米）。使用伪彩色映射有助于直观理解深度分布。

四、训练集与测试集划分

根据官方文档，NYU Depth V2 提供了预定义的训练/测试划分索引。若未提供，则可手动划分数据集：

from sklearn.model_selection import train_test_split

# 假设我们已将所有图像和深度图加载到numpy数组中
X_train, X_test, y_train, y_test = train_test_split(
    all_images, all_depths, test_size=0.2, random_state=42
)
    

建议保存划分后的索引，便于后续复现实验结果。

五、内存管理与数据增强策略

由于NYU Depth V2数据量较大，直接一次性加载全部数据可能引发内存溢出问题。以下是几种优化策略：

1. 按需加载（Lazy Loading）：使用生成器逐批读取数据。
2. 数据分片（Chunking）：将大文件分割成多个小块进行处理。
3. 使用HDF5的Dataset切片功能：仅加载当前批次所需数据。

此外，可在训练过程中引入数据增强技术提升模型泛化能力，例如：

操作	描述
随机裁剪	从原始图像中随机裁剪区域，保持输入尺寸一致。
色彩抖动	对RGB图像应用亮度、对比度等扰动。
深度归一化	将深度值归一化到[0,1]范围以适应神经网络输入。

六、总结流程图

整个数据处理流程可通过以下Mermaid流程图展示：

graph TD
            A[开始] --> B{检查.mat版本}
            B -- v4~v7 --> C[使用scipy.io.loadmat()]
            B -- v7.3+ --> D[使用h5py.File()]
            D --> E[提取RGB和深度图像]
            E --> F[调整图像维度]
            F --> G[可视化图像]
            G --> H[划分训练/测试集]
            H --> I[内存管理与数据增强]
            I --> J[准备送入模型训练]