保姆级教程：用Python解析ScanNet数据集中的2D图像与深度图（附避坑指南）

当你第一次打开ScanNet数据集时，面对密密麻麻的 .sens 、 .ply 和 .json 文件，是否感到无从下手？作为计算机视觉领域最常用的RGB-D数据集之一，ScanNet包含了超过1500个场景的250万帧数据，但如何高效提取其中的2D图像和深度信息，却是许多研究者面临的第一个实战难题。本文将手把手带你用Python拆解这个"数据黑箱"，从环境配置到完整可视化流程，避开我踩过的所有坑。

1. 环境准备：避开Python版本与依赖的深坑

在开始解析前，我们需要特别注意ScanNet官方工具对Python版本的敏感度。虽然官方示例多用Python 2.7，但通过几个关键调整，完全可以适配Python 3环境。

1.1 必备工具包清单

首先通过 pip 安装以下核心依赖（建议使用虚拟环境）：

pip install numpy imageio tqdm opencv-python pillow

特别注意 imageio 的版本选择——这是第一个大坑。新版本默认不包含 freeimage 插件，而深度图的解析恰恰需要它。有两种解决方案：

• 方案A ：安装指定旧版本（推荐）

  pip install imageio==2.27.0
  pip install imageio[freeimage]
  

• 方案B ：手动下载DLL文件

  1. 从 imageio-binaries 下载对应系统的 freeimage.dll
  2. 放入Python安装目录的 Lib/site-packages/imageio/plugins 文件夹

提示：如果遇到 DLL load failed 错误，通常是因为系统缺少Visual C++ Redistributable，需安装VS 2015-2022的运行库。

1.2 Python 3兼容性改造

官方提供的 SensorData.py 脚本存在Python 2到3的兼容问题，主要集中在字符串处理上。需要修改以下关键点：

# 原代码（Python 2）
header = f.read(28)
# 修改为（Python 3）
header = f.read(28).decode('utf-8')

对于 struct.unpack 调用，需要显式指定字节顺序：

# 原代码
data = struct.unpack('I', f.read(4))[0]
# 修改为
data = struct.unpack('<I', f.read(4))[0]  # 小端序

2. 深度解析.sens文件结构

ScanNet的核心数据存储在 .sens 二进制文件中，理解其结构是高效提取的前提。文件格式如下表所示：

偏移量	数据类型	描述
0x00	char[4]	魔数"RGBD"
0x04	int32	版本号(当前为2)
0x08	int64	总帧数
0x10	float32	色彩图像宽度
0x14	float32	色彩图像高度
...	...	其他元数据
变长	帧数据块	连续存储的帧数据

每个帧数据块包含：

• 4字节：色彩图像数据长度
• 色彩图像JPEG数据
• 4字节：深度图像数据长度
• 深度图像PNG数据
• 64字节：相机位姿矩阵（4x4 float32）

3. 实战：提取与可视化RGB-D数据

3.1 使用改进版解析脚本

以下是增强版的解析代码，添加了进度显示和异常处理：

from sensor_data import SensorData
from tqdm import tqdm
import os

def export_scan(scan_file, output_folder, skip_frames=50):
    try:
        sd = SensorData(scan_file)
        os.makedirs(output_folder, exist_ok=True)
        
        # 创建各子目录
        subdirs = ['color', 'depth', 'pose']
        for d in subdirs:
            os.makedirs(os.path.join(output_folder, d), exist_ok=True)

        # 带进度条的导出过程
        for idx in tqdm(range(0, len(sd.frames), skip_frames),
                       desc=f"Processing {os.path.basename(scan_file)}"):
            frame = sd.frames[idx]
            
            # 导出彩色图
            color_file = os.path.join(output_folder, 'color', f"{idx:06d}.jpg")
            with open(color_file, 'wb') as f:
                f.write(frame.color_data)
            
            # 导出深度图（单位转换为mm）
            depth_file = os.path.join(output_folder, 'depth', f"{idx:06d}.png")
            frame.save_depth(depth_file, depth_shift=1000)
            
            # 导出位姿
            pose_file = os.path.join(output_folder, 'pose', f"{idx:06d}.txt")
            np.savetxt(pose_file, frame.camera_to_world)
            
    except Exception as e:
        print(f"Error processing {scan_file}: {str(e)}")

3.2 深度图可视化技巧

原始深度图是16位PNG，直接打开可能显示全黑。推荐使用以下OpenCV方法增强可视化：

import cv2

def visualize_depth(depth_path):
    depth = cv2.imread(depth_path, cv2.IMREAD_ANYDEPTH)
    depth = depth.astype(np.float32) / 1000.0  # 转换为米
    
    # 归一化到0-255
    depth_normalized = cv2.normalize(depth, None, 0, 255, cv2.NORM_MINMAX)
    depth_colormap = cv2.applyColorMap(
        depth_normalized.astype(np.uint8), 
        cv2.COLORMAP_JET
    )
    
    # 并排显示RGB和深度
    rgb = cv2.imread(rgb_path)
    rgb_resized = cv2.resize(rgb, (depth_colormap.shape[1], depth_colormap.shape[0]))
    combined = np.hstack((rgb_resized, depth_colormap))
    
    cv2.imshow('RGB-D', combined)
    cv2.waitKey(0)

4. 构建PyTorch数据管道

将原始数据转换为适合深度学习训练的格式：

from torch.utils.data import Dataset
from torchvision import transforms

class ScanNetDataset(Dataset):
    def __init__(self, root_dir, transform=None, target_size=(240, 320)):
        self.root = root_dir
        self.transform = transform or transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                               std=[0.229, 0.224, 0.225])
        ])
        self.target_size = target_size
        self.samples = self._load_samples()
    
    def _load_samples(self):
        samples = []
        scenes = [d for d in os.listdir(self.root) 
                 if os.path.isdir(os.path.join(self.root, d))]
        
        for scene in scenes:
            color_dir = os.path.join(self.root, scene, 'color')
            depth_dir = os.path.join(self.root, scene, 'depth')
            
            for img_name in os.listdir(color_dir):
                if img_name.endswith('.jpg'):
                    base_name = os.path.splitext(img_name)[0]
                    samples.append((
                        os.path.join(color_dir, img_name),
                        os.path.join(depth_dir, f"{base_name}.png")
                    ))
        return samples
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        color_path, depth_path = self.samples[idx]
        
        # 加载彩色图
        color = Image.open(color_path).convert('RGB')
        color = color.resize(self.target_size)
        
        # 加载深度图（单位：米）
        depth = cv2.imread(depth_path, cv2.IMREAD_ANYDEPTH)
        depth = cv2.resize(depth, self.target_size[::-1])
        depth = depth.astype(np.float32) / 1000.0
        
        if self.transform:
            color = self.transform(color)
            
        return color, torch.from_numpy(depth).unsqueeze(0)

5. 高级技巧与性能优化

5.1 多线程加速解析

对于大规模数据，可以使用Python的 multiprocessing 模块：

from multiprocessing import Pool

def process_scene(scene_path):
    try:
        export_scan(scene_path, os.path.join(output_root, os.path.basename(scene_path)))
    except Exception as e:
        print(f"Failed on {scene_path}: {e}")

if __name__ == '__main__':
    scenes = [f for f in os.listdir(input_dir) if f.endswith('.sens')]
    with Pool(processes=4) as pool:
        pool.map(process_scene, scenes)

5.2 数据预处理流水线

建议预处理阶段完成以下操作并存储为HDF5文件：

• 图像降采样到统一分辨率
• 深度图无效值填充
• 相机参数归一化
• 生成边缘mask

import h5py

def create_hdf5_dataset(src_dir, h5_path):
    with h5py.File(h5_path, 'w') as hf:
        for scene in os.listdir(src_dir):
            scene_group = hf.create_group(scene)
            
            # 存储全局属性
            scene_group.attrs['frame_count'] = len(os.listdir(f"{src_dir}/{scene}/color"))
            
            # 创建可扩展的dataset
            max_shape = (None, 240, 320)
            color_dset = scene_group.create_dataset(
                'color', shape=(0, 240, 320, 3),
                maxshape=max_shape, dtype='uint8', chunks=True)
            
            depth_dset = scene_group.create_dataset(
                'depth', shape=(0, 240, 320),
                maxshape=max_shape, dtype='float32', chunks=True)
            
            # 逐帧填充数据
            for i, img_name in enumerate(sorted(os.listdir(f"{src_dir}/{scene}/color"))):
                color = cv2.imread(f"{src_dir}/{scene}/color/{img_name}")
                color = cv2.resize(color, (320, 240))
                
                # 扩展dataset并写入
                color_dset.resize(i+1, axis=0)
                color_dset[i] = color
                
                # 相同操作处理深度图...

在实际项目中，我发现将数据预处理为HDF5格式后，训练时的数据加载速度可以提升3-5倍，特别当使用SSD存储时效果更明显。对于超大规模数据，可以考虑使用更专业的存储方案如TFRecords或WebDataset。