第 10 章 · MAE 遮挡重建代码实验¶

# 使用真实建筑照片，按 ViT 常见设置切成 16x16 图块。
raw_photo = load_sample_image("china.jpg")
vit_image = np.asarray(Image.fromarray(raw_photo).resize((224, 224))) / 255.0
patch_size = 16
patch_grid = vit_image.shape[0] // patch_size
patches = vit_image.reshape(patch_grid, patch_size, patch_grid, patch_size, 3).swapaxes(1, 2)
patch_tokens = patches.reshape(-1, patch_size * patch_size * 3)

patch_summary = []
for patch_id, patch in enumerate(patches.reshape(-1, patch_size, patch_size, 3)):
    row, col = divmod(patch_id, patch_grid)
    patch_summary.append({
        "图块编号": patch_id,
        "行": row,
        "列": col,
        "向量维度": patch_tokens.shape[1],
        "R均值": patch[:, :, 0].mean(),
        "G均值": patch[:, :, 1].mean(),
        "B均值": patch[:, :, 2].mean(),
        "亮度标准差": patch.mean(axis=2).std(),
    })

patch_df = pd.DataFrame(patch_summary)
display(patch_df.head(12).round(3))

# MAE 核心目标：随机遮挡大部分图块，再根据可见图块重建遮挡区域。
rng = np.random.default_rng(4)
mae_image = vit_image.astype(float)
mae_patch = patch_size
mae_grid = patch_grid
num_patches = len(patch_tokens)
patch_vectors = patches.reshape(num_patches, mae_patch * mae_patch * 3)
coords = np.array([divmod(idx, mae_grid) for idx in range(num_patches)], dtype=float)
coords[:, 0] = coords[:, 0] / (mae_grid - 1)
coords[:, 1] = coords[:, 1] / (mae_grid - 1)

mae_model_name = "facebook/vit-mae-base"
mae_source = "ViT-MAE 预训练模型"
mae_error = ""

try:
    mae_packages = {"torch": "torch>=2.2", "transformers": "transformers>=4.40"}
    missing = [package for module, package in mae_packages.items() if importlib.util.find_spec(module) is None]
    install_packages(missing)

    import torch
    from transformers import AutoImageProcessor, ViTMAEForPreTraining

    with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
        mae_processor = AutoImageProcessor.from_pretrained(mae_model_name)
        mae_model = ViTMAEForPreTraining.from_pretrained(mae_model_name)
    mae_model.eval()
    mae_inputs = mae_processor(images=Image.fromarray(raw_photo).resize((224, 224)), return_tensors="pt")
    torch.manual_seed(4)
    with torch.no_grad():
        mae_outputs = mae_model(**mae_inputs)

    patch_size_mae = mae_model.config.patch_size
    x_tensor = mae_inputs["pixel_values"]
    y_tensor = mae_model.unpatchify(mae_outputs.logits)
    mask_tensor = mae_outputs.mask.unsqueeze(-1).repeat(1, 1, patch_size_mae ** 2 * 3)
    mask_tensor = mae_model.unpatchify(mask_tensor)

    mean = np.array(mae_processor.image_mean)
    std = np.array(mae_processor.image_std)
    def denorm_image(array):
        return np.clip(array * std + mean, 0, 1)

    x_image = x_tensor[0].permute(1, 2, 0).numpy()
    y_image = y_tensor[0].permute(1, 2, 0).numpy()
    mask_pixels = mask_tensor[0].permute(1, 2, 0).numpy()
    mae_image = denorm_image(x_image)
    pred_image = denorm_image(y_image)
    mae_reconstruction = denorm_image(x_image * (1 - mask_pixels) + y_image * mask_pixels)
    pixel_mask = mask_pixels[..., 0].astype(bool)
    mae_masked_image = mae_image.copy()
    mae_masked_image[pixel_mask] = 0.72
    mae_mask = mae_outputs.mask[0].cpu().numpy().astype(bool)
    mae_mask_map = mae_mask.reshape(mae_grid, mae_grid)
except Exception as exc:
    mae_source = "本地图块先验回退"
    mae_error = str(exc)[:160]
    mask_ratio = 0.50
    visible_ids = np.sort(rng.choice(num_patches, size=int(num_patches * (1 - mask_ratio)), replace=False))
    mae_mask = np.ones(num_patches, dtype=bool)
    mae_mask[visible_ids] = False
    mae_mask_map = mae_mask.reshape(mae_grid, mae_grid)

    def fourier_position_features(points):
        features = [points]
        for freq in [1, 2, 3, 4, 6, 8, 10]:
            features.append(np.sin(2 * np.pi * freq * points))
            features.append(np.cos(2 * np.pi * freq * points))
        return np.hstack(features)

    decoder_features = fourier_position_features(coords)
    reconstructor = make_pipeline(StandardScaler(), Ridge(alpha=0.2))
    reconstructor.fit(decoder_features, patch_vectors)
    pred_vectors = np.clip(reconstructor.predict(decoder_features), 0, 1)
    pred_image = pred_vectors.reshape(mae_grid, mae_grid, mae_patch, mae_patch, 3).swapaxes(1, 2).reshape(224, 224, 3)
    mae_masked_image = mae_image.copy()
    mae_reconstruction = mae_image.copy()
    for patch_id in np.flatnonzero(mae_mask):
        row, col = divmod(patch_id, mae_grid)
        r0, r1 = row * mae_patch, (row + 1) * mae_patch
        c0, c1 = col * mae_patch, (col + 1) * mae_patch
        mae_masked_image[r0:r1, c0:c1] = 0.72
        mae_reconstruction[r0:r1, c0:c1] = pred_image[r0:r1, c0:c1]
    pixel_mask = np.kron(mae_mask_map, np.ones((mae_patch, mae_patch), dtype=bool))

masked_ids = np.flatnonzero(mae_mask)
visible_ids = np.flatnonzero(~mae_mask)
masked_positions = [{"图块编号": idx, "行": idx // mae_grid, "列": idx % mae_grid} for idx in masked_ids]
mae_summary = pd.DataFrame({
    "指标": ["重建来源", "总图块数", "可见图块数", "遮挡图块数", "遮挡比例", "遮挡区域均方误差", "回退信息"],
    "值": [
        mae_source,
        num_patches,
        len(visible_ids),
        len(masked_ids),
        round(float(mae_mask.mean()), 3),
        round(float(mean_squared_error(mae_image[pixel_mask].reshape(-1), mae_reconstruction[pixel_mask].reshape(-1))), 4),
        mae_error,
    ],
})
display(mae_summary)
display(pd.DataFrame(masked_positions).head(12))

# 绘制可见输入、预测图块和重建图像。
fig, axes = plt.subplots(2, 3, figsize=(10.8, 7.0))
plot_items = [
    (mae_image, "原图", None),
    (mae_mask_map, "遮挡位置", "Greys"),
    (mae_masked_image, "可见输入", None),
    (pred_image, "根据可见图块预测", None),
    (mae_reconstruction, "遮挡处填回预测", None),
    (np.abs(mae_image - mae_reconstruction) * 3, "重建误差 x3", None),
]
for ax, (data, title, cmap) in zip(axes.ravel(), plot_items):
    ax.imshow(np.clip(data, 0, 1), cmap=cmap, vmin=0, vmax=1)
    ax.set_title(title, fontweight="bold")
    ax.set_xticks([])
    ax.set_yticks([])
fig.suptitle("MAE 核心目标：只看可见图块，重建被遮挡区域", x=0.08, ha="left", fontsize=14, fontweight="bold", color="#0f172a")
plt.tight_layout()
plt.show()