BIWI数据集实战指南：从数据加载到3D人脸关键点检测模型训练

1次阅读

没有评论

共计 3970 个字符，预计需要花费 10 分钟才能阅读完成。

BIWI 数据集是苏黎世联邦理工学院计算机视觉实验室发布的 3D 人脸分析基准数据集，包含超过 15,000 张高精度 3D 面部扫描数据。该数据集通过结构光扫描仪采集，每帧包含：

深度图像（640×480 分辨率）
对应的 RGB 彩色图像
预标注的 3D 面部关键点（30 个解剖学标记点）
头部姿态参数（欧拉角表示）

典型应用场景包括：

3D 人脸关键点检测
头部姿态估计
面部表情分析
3D 人脸重建

原始数据采用专有二进制格式存储，主要挑战在于：

多模态数据对齐：深度图与 RGB 图需要严格同步
二进制解析 ：需要处理.pose 和.depth文件的字节流
坐标系转换：世界坐标系、相机坐标系和图像坐标系的转换
缺失值处理：深度图中的无效像素（值为 0）

import os
import numpy as np
import torch
from torch.utils.data import Dataset
import cv2

class BIWIDataset(Dataset):
    """BIWI 3D 人脸关键点检测数据集"""

    def __init__(self, root_dir, transform=None):
        """
        参数:
            root_dir: 数据集根目录
            transform: 数据增强函数
        """
        self.root_dir = root_dir
        self.transform = transform
        self.samples = self._load_samples()

    def _load_samples(self):
        samples = []
        for seq_dir in os.listdir(self.root_dir):
            seq_path = os.path.join(self.root_dir, seq_dir)
            if not os.path.isdir(seq_path):
                continue

            # 解析每帧数据
            for frame in os.listdir(os.path.join(seq_path, 'depth')):
                if not frame.endswith('.depth'):
                    continue

                frame_id = frame.split('.')[0]
                sample = {'depth_path': os.path.join(seq_path, 'depth', frame),
                    'rgb_path': os.path.join(seq_path, 'rgb', f'{frame_id}.png'),
                    'pose_path': os.path.join(seq_path, 'pose', f'{frame_id}.pose')
                }
                samples.append(sample)
        return samples

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        sample = self.samples[idx]

        # 加载深度图（16 位无符号整型）depth = np.fromfile(sample['depth_path'], dtype='<u2').reshape(480, 640)
        depth = depth.astype(np.float32) / 1000.0  # 转换为米单位

        # 加载 RGB 图像
        rgb = cv2.imread(sample['rgb_path'])
        rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)

        # 加载头部姿态
        with open(sample['pose_path'], 'r') as f:
            pose = np.array([float(x) for x in f.read().split()])

        # 数据预处理
        if self.transform:
            rgb, depth = self.transform(rgb, depth)

        return {'rgb': torch.FloatTensor(rgb).permute(2, 0, 1),
            'depth': torch.FloatTensor(depth).unsqueeze(0),
            'pose': torch.FloatTensor(pose)
        }

建议的预处理流程：

深度图归一化 ：将深度值缩放到[0,1] 范围
RGB 归一化：使用 ImageNet 均值标准差
数据增强：
随机水平翻转（需同步调整关键点）
小角度旋转（±5 度）
颜色抖动

class BIWITransform:
    def __init__(self, augment=True):
        self.augment = augment

    def __call__(self, rgb, depth):
        # 归一化
        depth = np.clip(depth / 4.0, 0, 1)  # 假设最大深度 4 米
        rgb = rgb.astype(np.float32) / 255.0

        if self.augment:
            # 随机水平翻转
            if np.random.rand() > 0.5:
                rgb = np.fliplr(rgb)
                depth = np.fliplr(depth)

            # 随机旋转
            angle = np.random.uniform(-5, 5)
            rgb = self._rotate_image(rgb, angle)
            depth = self._rotate_image(depth, angle)

        return rgb, depth

    def _rotate_image(self, img, angle):
        """辅助旋转函数"""
        rows, cols = img.shape[:2]
        M = cv2.getRotationMatrix2D((cols/2, rows/2), angle, 1)
        return cv2.warpAffine(img, M, (cols, rows))

import torch.nn as nn

class Simple3DKeypointModel(nn.Module):
    def __init__(self, num_keypoints=30):
        super().__init__()

        # 特征提取主干网络
        self.backbone = nn.Sequential(nn.Conv2d(4, 32, kernel_size=3, padding=1),  # 输入通道 4 (RGB+D)
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )

        # 回归头
        self.regressor = nn.Sequential(nn.Linear(128 * 60 * 80, 512),  # 假设输入 640x480，经过 3 次下采样
            nn.ReLU(),
            nn.Linear(512, num_keypoints * 3)  # 输出 x,y,z 坐标
        )

    def forward(self, rgb, depth):
        # 拼接 RGB 和深度
        x = torch.cat([rgb, depth], dim=1)

        # 特征提取
        features = self.backbone(x)
        features = features.view(features.size(0), -1)

        # 回归关键点
        keypoints = self.regressor(features)
        return keypoints.view(-1, num_keypoints, 3)

关键训练参数：

损失函数：Smooth L1 Loss（对异常值鲁棒）
评估指标：平均关键点误差（Mean Per Joint Position Error）
优化器：AdamW (lr=1e-4)
Batch Size：16（取决于显存）

model = Simple3DKeypointModel().cuda()
criterion = nn.SmoothL1Loss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

# 训练循环示例
for epoch in range(100):
    for batch in dataloader:
        rgb = batch['rgb'].cuda()
        depth = batch['depth'].cuda()
        gt_keypoints = load_keypoints(batch)  # 需实现关键点加载

        # 前向传播
        pred_keypoints = model(rgb, depth)

        # 计算损失
        loss = criterion(pred_keypoints, gt_keypoints)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()