FFaceNeRF模块

论文《FFaceNeRF: Few-shot Face Editing in Neural Radiance Fields》

深度交流Q裙：1051849847
全网同名【大嘴带你水论文】 B站定时发布详细讲解视频

视频地址，点击查看论文详细讲解，每日更新：
https://b23.tv/zdapaC3

详细代码见文章最后

在这里插入图片描述

1、作用

FFaceNeRF旨在解决现有基于NeRF的3D人脸编辑方法严重依赖固定布局的预训练分割蒙版、导致用户控制能力有限的问题。它使用户能够根据特定的编辑需求（如虚拟试妆、医疗整形预览等）自由定义和使用新的蒙版布局，而无需收集和标注大规模数据集，极大地提升了3D人脸编辑的灵活性和实用性。

在这里插入图片描述

2、机制

几何适配器 (Geometry Adapter) : 在预训练的几何解码器（用于生成固定布局的分割图）之后，添加一个轻量级的MLP网络作为几何适配器。该适配器负责将固定布局的输出调整为用户期望的、任意布局的分割蒙版。
特征注入 (Feature Injection) : 为了在适应新蒙版时保留丰富的几何细节，模型将预训练模型中的三平面特征（tri-plane feature）和视角方向（view direction）直接注入到几何适配器中，弥补了预训练解码器可能丢失的信息。
三平面增强的潜在混合 (LMTA) : 这是一种为小样本学习设计的数据增强策略。通过在生成器的潜在空间中混合不同层的潜在编码，可以在保持核心语义信息（如人脸结构）不变的同时，生成多样化的训练样本（如改变色调、饱和度），有效避免了在仅有10个样本的情况下发生的过拟合。
基于重叠的优化 (Overlap-based Optimization) : 在训练和推理过程中，除了使用传统的交叉熵损失外，还引入了基于DICE系数的重叠损失（overlap loss）。这种损失函数对小区域的变化更敏感，确保了即使在编辑精细区域（如瞳孔、鼻翼）时也能实现精确对齐和稳定生成。

3、独特优势

小样本高效学习 : 最显著的优势是其小样本（Few-shot）能力，仅需约10个带有自定义蒙版的样本即可完成训练，快速适应新的编辑任务，极大降低了数据和时间成本。
高度的灵活性和控制力 : 用户不再受限于固定的分割类别，可以为任何感兴趣的面部区域创建蒙版并进行编辑，实现了前所未有的控制自由度。
精细区域的精准编辑 : 通过基于重叠的优化策略，模型能够精确地处理和编辑微小面部特征，解决了传统方法在小区域编辑上容易失败的痛点。
保持身份和非编辑区域 : 在编辑特定区域时，能够很好地保持人脸的身份特征以及未编辑区域的图像内容，生成结果自然且保真度高。

4、代码实现

import torch
import torch.nn as nn
import torch.nn.functional as Fclass MockTriPlaneGenerator(nn.Module):""" NeRF三平面生成器，用于生成基础的分割图。 """def __init__(self, num_classes=10):super().__init__()self.num_classes = num_classes# 一个简单的卷积层，模拟从三平面特征到分割图的解码过程self.decoder = nn.Conv2d(32, num_classes, kernel_size=1)def forward(self, triplane_features, view_direction):# 解码过程，返回一个随机的分割图batch_size = triplane_features.shape[0]# 返回一个随机的、固定布局的分割图return torch.randn(batch_size, self.num_classes, 64, 64)class MockDataset(torch.utils.data.Dataset):def __init__(self, num_samples=10, num_classes=12):self.num_samples = num_samplesself.num_classes = num_classesdef __len__(self):return self.num_samplesdef __getitem__(self, idx):# 生成模拟数据triplane_features = torch.randn(1, 32, 256, 256) # 模拟三平面特征view_direction = torch.randn(1, 3) # 模拟视角方向# 模拟用户自定义的目标蒙版target_mask = torch.randint(0, self.num_classes, (1, 64, 64), dtype=torch.long)return triplane_features, view_direction, target_mask# ----------------------------------------------------------------------------
# 核心代码实现 (Core Implementation)
# ----------------------------------------------------------------------------class GeometryAdapter(nn.Module):""" 几何适配器，将固定布局的分割图调整为用户自定义的布局。 """def __init__(self, input_channels, output_channels):super().__init__()# 一个轻量级的MLP，用于适配几何特征self.adapter = nn.Sequential(nn.Conv2d(input_channels, 64, kernel_size=1),nn.ReLU(),nn.Conv2d(64, output_channels, kernel_size=1))def forward(self, x):return self.adapter(x)class FFaceNeRF(nn.Module):""" FFaceNeRF 模块，集成了预训练生成器和几何适配器。 """def __init__(self, pretrained_generator, num_custom_classes):super().__init__()self.pretrained_generator = pretrained_generator# 几何适配器的输入通道数等于预训练生成器的输出类别数self.adapter = GeometryAdapter(pretrained_generator.num_classes, num_custom_classes)def forward(self, triplane_features, view_direction):# 首先，使用预训练生成器获取固定布局的分割图fixed_layout_seg = self.pretrained_generator(triplane_features, view_direction)# 然后，通过几何适配器将其调整为自定义布局custom_layout_seg = self.adapter(fixed_layout_seg.detach()) # detach以冻结预训练部分return custom_layout_segdef dice_loss(pred, target, smooth=1e-5):""" 计算DICE损失，对小区域优化更有效。 """pred = F.softmax(pred, dim=1)# 将target转换为one-hot编码target_one_hot = F.one_hot(target.squeeze(1), num_classes=pred.shape[1]).permute(0, 3, 1, 2).float()intersection = torch.sum(pred * target_one_hot, dim=(2, 3))union = torch.sum(pred, dim=(2, 3)) + torch.sum(target_one_hot, dim=(2, 3))dice = (2. * intersection + smooth) / (union + smooth)return 1 - dice.mean()# ----------------------------------------------------------------------------
# 运行示例 (Runnable Example)
# ----------------------------------------------------------------------------if __name__ == '__main__':# --- 1. 初始化参数和模型 ---num_fixed_classes = 10  # 预训练模型支持的固定类别数num_custom_classes = 12 # 用户自定义的类别数（例如，眉毛、上唇、下唇等）num_samples = 10        # 小样本数量epochs = 50             # 训练轮次# 初始化模拟的预训练生成器和FFaceNeRF模型pretrained_generator = MockTriPlaneGenerator(num_classes=num_fixed_classes)fface_nerf = FFaceNeRF(pretrained_generator, num_custom_classes)optimizer = torch.optim.Adam(fface_nerf.adapter.parameters(), lr=0.001)# --- 2. 创建模拟数据集 ---dataset = MockDataset(num_samples=num_samples, num_classes=num_custom_classes)dataloader = torch.utils.data.DataLoader(dataset, batch_size=2)# --- 3. 训练循环 ---for epoch in range(epochs):total_loss = 0for triplane_features, view_direction, target_mask in dataloader:optimizer.zero_grad()# 获取模型预测的自定义分割图predicted_seg = fface_nerf(triplane_features, view_direction)# 计算损失（交叉熵 + DICE损失）loss_ce = F.cross_entropy(predicted_seg, target_mask.squeeze(1))loss_dice = dice_loss(predicted_seg, target_mask)loss = loss_ce + 0.5 * loss_dice # 组合损失loss.backward()optimizer.step()total_loss += loss.item()if (epoch + 1) % 10 == 0:print(f"Epoch [{epoch+1}/{epochs}], 平均损失: {total_loss / len(dataloader):.4f}")# --- 4. 模拟推理 ---print("\n进行一次模拟推理...")with torch.no_grad():# 取一个样本进行测试test_features, test_view_dir, ground_truth_mask = next(iter(dataloader))predicted_mask_logits = fface_nerf(test_features, test_view_dir)predicted_mask = torch.argmax(predicted_mask_logits, dim=1)print(f"输入特征尺寸: {test_features.shape}")print(f"预测蒙版尺寸: {predicted_mask.shape}")print(f"预测蒙版中的类别: {torch.unique(predicted_mask)}")