说明

本项目精简自 https://github.com/InterDigitalInc/FeatureStyleEncoder 原项目代码逻辑比较复杂，不便于我们理解论文

1. 工作逻辑

本项目对原始 FeatureStyleEncoder 代码库做了大幅精简，只保留两个核心推理功能：

功能	入口	说明
人脸重建 (inversion)	python infer.py reconstruct	输入一张人脸图像，编码后重建
人脸编辑 (editing)	python infer.py edit	输入图像 + 编辑方向，输出编辑后图像

2. 保留的文件结构

FeatureStyleEncoder-main/
├── arcface_iresnet.py                    # ArcFace IResNet-50 backbone
├── fs_encoder.py                         # FS Encoder (特征风格编码器)
├── infer.py                              # 推理入口 (reconstruct + edit)
|── stylegan2/
│     ├── model.py                       # StyleGAN2 Generator
│     └── op/                             # CUDA 算子 (upfirdn2d, fused_act)
├── edit_boundary/
│   ├── Eyeglasses_boundary.npy           # 戴眼镜方向
│   ├── Heavy_Makeup_boundary.npy         # 浓妆方向
│   └── Smiling_boundary.npy              # 微笑方向
├── pretrained_models/                    # 预训练权重（需自行下载）
│   ├── backbone.pth                      # ArcFace backbone
│   ├── 143_enc.pth                       # FS Encoder 权重
│   └── psp_ffhq_encode.pt                # StyleGAN2 / PSP 权重
└── ARCHITECTURE.md                       # 本文档

---

3. 模型架构

3.1 编码器：fs_encoder_v2

文件：fs_encoder.py

输入：256×256 人脸图像 (1, 3, 256, 256)
         │
    ┌────▼────┐
    │  conv   │  IResNet-50 前3层 (conv1 + bn1 + prelu)
    └────┬────┘
         │
    ┌────▼────┐
    │ block_1 │  IResNet layer1: (1, 64,  128, 128)
    └────┬────┘  → avg_pool(3×3) → feature_1 (1, 64,  3, 3)
         │
    ┌────▼────┐
    │ block_2 │  IResNet layer2: (1, 128, 64, 64)
    └────┬────┘  → avg_pool(3×3) → feature_2 (1, 128, 3, 3)
         │
    ┌────▼────┐
    │ block_3 │  IResNet layer3: (1, 256, 32, 32)
    └────┬────┘  → content_layer → fea (1, 512, H, W)   ← 空间内容特征
         │       → avg_pool(3×3) → feature_3 (1, 256, 3, 3)
         │
    ┌────▼────┐
    │ block_4 │  IResNet layer4: (1, 512, 16, 16)
    └────┬────┘  → avg_pool(3×3) → feature_4 (1, 512, 3, 3)
         │
    concat([feature_1, feature_2, feature_3, feature_4])
    reshape → (1, 960*9)
         │
    18× Linear(960*9, 512)
         │
    stack → w_delta (1, 18, 512)   ← W+ 潜码增量

输出：w_delta (1, 18, 512),  fea (1, 512, H, W)

关键点：

• stride=(2,2) 控制 content_layer 最后一层卷积的步长
• 最终 W+ 潜码 = w_delta + dlatent_avg（加上 StyleGAN 的平均潜码）

3.2 Backbone：arcface_iresnet.py

文件：arcface_iresnet.py

标准 IResNet-50（InsightFace 版本），用作 fs_encoder_v2 的特征提取骨干网络。 层结构：[3, 4, 14, 3]（IBasicBlock）

3.3 生成器：StyleGAN2 Generator

文件：pixel2style2pixel/models/stylegan2/model.py

输入：W+ 潜码 (1, 18, 512)  +  可选的 features_in 列表

解码器共 18 层 StyledConv（对应 4×4 到 1024×1024）
层索引 5 的空间分辨率：32×32（对应编码器 content feature 的分辨率）

features_in = [None, None, None, None, None, fea, None, ..., None]
                0     1     2     3     4     5    6          17

feature_scale=1.0 时，第 5 层输出 = fea（完全替换）

---

4. 推理流程

4.1 人脸重建 (Reconstruction)

输入图像 (任意尺寸)
    │
    ▼
Resize → 1024×1024, Normalize[-1,1]         img (1,3,1024,1024)
    │
    ▼
Downscale ×2 (bilinear)                     img_256 (1,3,256,256)
    │
    ▼
fs_encoder_v2(img_256)
    ├─── w_delta  (1,18,512)
    └─── fea      (1,512,H,W)
    │
    ▼
w = w_delta + dlatent_avg                   W+ 潜码 (1,18,512)
    │
    ▼
features_in = [None]*5 + [fea] + [None]*12
    │
    ▼
StyleGAN2 Generator([w], features_in=features_in, feature_scale=1.0)
    │
    ▼
x_recon (1,3,1024,1024) → Clip[0,1] → 保存

4.2 人脸编辑 (Boundary-based Editing)

输入图像
    │
    ├─── [同重建流程] ───► w (1,18,512), fea (1,512,H,W)
    │
    ▼
boundary = np.load("xxx_boundary.npy")      shape: (1, 9216)=(1, 18*512)
w_edit = (w.reshape(1,-1) + strength * boundary).reshape(1,18,512)
    │
    ├──► StyleGAN2([w],      return_features=True) → feats_orig
    └──► StyleGAN2([w_edit], return_features=True) → feats_edit
    │
    ▼
delta = feats_edit[5] - feats_orig[5]       层5的特征差
fea_new = fea + delta                       补偿编码器内容特征
    │
    ▼
features_in = [None]*5 + [fea_new] + [None]*12
    │
    ▼
StyleGAN2([w_edit], features_in=features_in, feature_scale=1.0)
    │
    ▼
x_edit → Clip[0,1] → 保存

为什么需要特征补偿 (delta)？

编码器输出的 fea 是从原始 W 潜码对应的图像空间提取的内容特征。 当我们将 W 潜码移动到 w_edit 后，StyleGAN 解码器在第 5 层会生成不同的中间特征。 直接注入原始 fea 会造成 W 空间与特征空间的不一致。 通过加上 delta = feats_edit[5] - feats_orig[5]，将内容特征迁移到编辑后的特征空间， 从而在保持身份信息的同时实现属性编辑。

---

6. 使用方法

人脸重建

python infer.py reconstruct \
    --input  test/00020.jpg \
    --output output/recon.jpg \
    --arcface_path  pretrained_models/backbone.pth \
    --enc_path      pretrained_models/143_enc.pth \
    --stylegan_path pixel2style2pixel/pretrained_models/psp_ffhq_encode.pt

人脸编辑

# 微笑
python infer.py edit \
    --input    test/00020.jpg \
    --boundary boundaries_ours/Smiling_boundary.npy \
    --output   output/edit_smiling.jpg \
    --strength 5.0

# 戴眼镜
python infer.py edit \
    --input    test/00020.jpg \
    --boundary boundaries_ours/Eyeglasses_boundary.npy \
    --output   output/edit_glasses.jpg \
    --strength 5.0

# 浓妆（反向编辑用负数 strength）
python infer.py edit \
    --input    test/00020.jpg \
    --boundary boundaries_ours/Heavy_Makeup_boundary.npy \
    --output   output/edit_makeup.jpg \
    --strength -5.0

7. 依赖环境

torch >= 1.8
torchvision
numpy
Pillow