从零构建3D实例分割训练集：Scannet数据预处理实战指南

当你第一次打开Scannet数据集时，可能会被数百个散乱的文件目录和复杂的文件命名规则搞得晕头转向。作为3D视觉领域最常用的基准数据集之一，Scannet为3D语义分割和实例分割任务提供了丰富的标注数据，但原始数据的组织方式并不直接适配大多数算法的训练流程。本文将以PointGroup算法为例，带你一步步完成从原始数据到训练就绪格式的完整转换过程。

1. Scannet数据集深度解析

Scannet v2数据集包含1,613个扫描场景，每个场景平均包含40,000个点云数据点。原始数据采用以下目录结构存储：

code复制scans/
    scene0000_00/
        scene0000_00_vh_clean_2.ply
        scene0000_00_vh_clean_2.labels.ply
        scene0000_00_vh_clean_2.0.010000.segs.json
        scene0000_00.aggregation.json
        ...
    scene0001_00/
    ...

关键文件类型说明：

在准备PointGroup等实例分割算法的训练数据时，我们需要特别注意以下几点：

• 测试集仅包含点云文件(_vh_clean_2.ply)，不包含任何标注
• 验证集和训练集需要完整的四种文件类型
• 不同算法可能需要不同的预处理方式

2. 数据目录重构实战

大多数3D分割算法要求数据按以下结构组织：

code复制scannetv2/
    train/
        scene0000_00_vh_clean_2.ply
        scene0000_00_vh_clean_2.labels.ply
        ...
    val/
        scene0700_00_vh_clean_2.ply
        ...
    test/
        scene1200_00_vh_clean_2.ply
        ...

下面是一个完整的Python脚本，用于将原始Scannet数据转换为上述结构：

python复制import os
import shutil
from tqdm import tqdm

def reorganize_scannet(base_dir, target_dir, split_files_dir):
    """
    重组Scannet数据目录结构
    
    参数:
        base_dir: 原始Scannet数据根目录(scans和scans_test所在目录)
        target_dir: 目标目录
        split_files_dir: 包含train.txt/val.txt/test.txt的目录
    """
    splits = ['train', 'val', 'test']
    test_split = False  # 标记是否处理测试集
    
    for split in splits:
        split_file = os.path.join(split_files_dir, f'{split}.txt')
        if not os.path.exists(split_file):
            continue
            
        with open(split_file) as f:
            scene_ids = [line.strip() for line in f]
        
        # 确定源目录(测试集单独处理)
        src_dir = os.path.join(base_dir, 'scans_test' if split == 'test' and test_split else 'scans')
        
        # 确定需要复制的文件类型
        file_types = [
            '_vh_clean_2.ply',
            '_vh_clean_2.labels.ply',
            '_vh_clean_2.0.010000.segs.json',
            '.aggregation.json'
        ] if split != 'test' else ['_vh_clean_2.ply']
        
        # 创建目标目录
        os.makedirs(os.path.join(target_dir, split), exist_ok=True)
        
        # 复制文件
        for scene_id in tqdm(scene_ids, desc=f'Processing {split} set'):
            for file_type in file_types:
                src_path = os.path.join(src_dir, scene_id, f'{scene_id}{file_type}')
                if not os.path.exists(src_path):
                    continue
                    
                dst_path = os.path.join(target_dir, split, f'{scene_id}{file_type}')
                shutil.copy2(src_path, dst_path)

使用示例：

bash复制python reorganize.py \
    --base_dir /path/to/scannet \
    --target_dir /path/to/scannetv2 \
    --split_files_dir /path/to/split_files

注意：官方提供的train.txt/val.txt/test.txt文件通常包含在算法代码仓库的dataset目录中，不同算法可能使用不同的划分方式。

3. 数据预处理进阶技巧

3.1 坐标系统一化

Scannet使用Metrical空间坐标系，而许多算法要求数据归一化到特定范围。以下是常用的坐标转换代码：

python复制import numpy as np
from plyfile import PlyData

def normalize_coords(ply_path):
    """将点云坐标归一化到[-1,1]范围"""
    plydata = PlyData.read(ply_path)
    xyz = np.vstack([plydata['vertex']['x'], 
                    plydata['vertex']['y'],
                    plydata['vertex']['z']]).T
    
    # 计算归一化参数
    center = (xyz.max(0) + xyz.min(0)) / 2
    scale = (xyz.max(0) - xyz.min(0)).max() / 2
    
    # 应用归一化
    xyz_normalized = (xyz - center) / scale
    
    # 更新PLY文件(此处省略写入代码)
    return xyz_normalized

3.2 标签映射处理

Scannet原始标签包含200+类别，通常需要映射到20个基准类别：

python复制SCANNET_LABEL_MAPPING = {
    1: 1,   # wall
    2: 2,   # floor
    3: 3,   # cabinet
    # ...其他映射规则
}

def remap_labels(label_path):
    plydata = PlyData.read(label_path)
    labels = plydata['vertex']['label']
    remapped = np.array([SCANNET_LABEL_MAPPING.get(l, 0) for l in labels])
    
    # 更新PLY文件(此处省略写入代码)
    return remapped

3.3 数据增强策略

针对3D点云数据的常用增强方法：

1. 随机旋转：绕z轴旋转0-360度
2. 随机缩放：尺度变化0.8-1.2倍
3. 随机位移：在xyz方向小范围平移
4. 点采样：随机下采样或上采样点云

实现示例：

python复制def augment_pointcloud(xyz, labels=None, rotation_range=360, scale_range=(0.8, 1.2)):
    # 随机旋转
    theta = np.random.uniform(0, rotation_range) * np.pi / 180
    rot_mat = np.array([
        [np.cos(theta), -np.sin(theta), 0],
        [np.sin(theta), np.cos(theta), 0],
        [0, 0, 1]
    ])
    xyz = xyz @ rot_mat
    
    # 随机缩放
    scale = np.random.uniform(*scale_range)
    xyz = xyz * scale
    
    # 随机位移
    offset = np.random.uniform(-0.1, 0.1, size=(1, 3))
    xyz = xyz + offset
    
    return xyz if labels is None else (xyz, labels)

4. 不同任务的数据准备差异

4.1 语义分割 vs 实例分割

4.2 PointGroup的特殊需求

PointGroup算法需要额外处理超点分割结果，将segs.json转换为算法所需的格式：

python复制import json

def process_segments(seg_file):
    with open(seg_file) as f:
        seg_data = json.load(f)
    
    seg_indices = seg_data['segIndices']
    segments = seg_data['segGroups']
    
    # 为每个点分配实例ID
    instance_ids = np.zeros_like(seg_indices)
    for seg in segments:
        for idx in seg['segments']:
            instance_ids[seg_indices == idx] = seg['id']
    
    return instance_ids

4.3 自定义数据集划分

当需要创建自定义训练/验证划分时，可以使用以下策略：

python复制import random

def create_custom_split(scene_ids, train_ratio=0.7, val_ratio=0.2):
    random.shuffle(scene_ids)
    n = len(scene_ids)
    train_end = int(n * train_ratio)
    val_end = train_end + int(n * val_ratio)
    
    return {
        'train': scene_ids[:train_end],
        'val': scene_ids[train_end:val_end],
        'test': scene_ids[val_end:]
    }

5. 常见问题排查

问题1：文件复制时报错"文件不存在"

• 检查原始文件路径是否正确
• 确认文件名后缀是否匹配（特别是_vh_clean_2与_vh_clean的区别）

问题2：训练时出现标签越界错误

• 验证标签映射是否正确
• 检查是否有未映射的原始标签

问题3：实例分割结果不理想

• 确认.aggregation.json文件是否正确处理
• 检查超点分割结果是否被正确利用

问题4：内存不足

• 考虑使用数据流式加载
• 减少同时加载的场景数量

python复制class StreamingDataset:
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.scene_list = os.listdir(data_dir)
    
    def __getitem__(self, idx):
        scene_id = self.scene_list[idx]
        ply_path = os.path.join(self.data_dir, scene_id)
        # 仅加载当前需要的场景数据
        return load_ply(ply_path)

在实际项目中，我发现最耗时的部分往往是数据验证阶段。建议在完整训练前，先对小批量数据进行端到端测试，确保数据管道的每个环节都正确无误。