如何展示HuggingFace ViTMAEModel生成的重建图像？

你的代码目前仅获取了模型的最后隐藏层状态，但ViT-MAE自编码器的输出包含重建后的像素张量，只需提取该张量并做格式转换，就能得到可显示的重建图像。以下是完整实现代码：

from transformers import AutoImageProcessor, ViTMAEModel
from PIL import Image
import requests
import torch

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

image_processor = AutoImageProcessor.from_pretrained("facebook/vit-mae-base")
model = ViTMAEModel.from_pretrained("facebook/vit-mae-base")

# 预处理图像并输入模型
inputs = image_processor(images=image, return_tensors="pt")
outputs = model(**inputs)

# 提取重建后的像素张量
reconstructed_tensor = outputs.reconstruction

# 反归一化，恢复像素值原始范围
mean = torch.tensor(image_processor.image_mean).view(1, 3, 1, 1)
std = torch.tensor(image_processor.image_std).view(1, 3, 1, 1)
reconstructed_tensor = reconstructed_tensor * std + mean

# 调整张量维度并转换为0-255的uint8格式（符合PIL图像要求）
reconstructed_array = (reconstructed_tensor.clamp(0, 1) * 255).permute(0, 2, 3, 1).squeeze().cpu().numpy().astype("uint8")

# 转换为PIL图像并显示
reconstructed_image = Image.fromarray(reconstructed_array)
reconstructed_image.show()

# 可选：同时显示原图像做对比
image.show()

关键步骤说明 #

• outputs.reconstruction：ViTMAEModel的核心输出字段，对应模型重建后的像素张量，形状为(batch_size, 3, height, width)
• 反归一化：图像预处理阶段做了减均值、除以标准差的归一化操作，这里需要反向计算恢复像素值的原始范围
• 维度调整：PyTorch张量默认通道维度在前，而PIL图像要求通道维度在最后，通过permute调整顺序后，用squeeze去掉批量维度
• 格式转换：将张量值限制在0-1区间后乘以255，转为uint8格式的numpy数组，再用Image.fromarray生成可显示的图像对象

内容的提问来源于stack exchange，提问作者Hadar D