Update app.py

app.py

@@ -1,69 +1,36 @@
 import torch
 import numpy as np
-from transformers import ViTMAEForPreTraining, ViTFeatureExtractor
+from transformers import ViTMAEForPreTraining, ViTImageProcessor
 from PIL import Image
 import gradio as gr
 import matplotlib.pyplot as plt
 
-# 加载模型和处理器
 model = ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base")
-feature_extractor = ViTFeatureExtractor.from_pretrained("facebook/vit-mae-base")
+feature_extractor = ViTImageProcessor.from_pretrained("facebook/vit-mae-base")
 
 def visualize_mae(image):
-    # 预处理图像
+    # 调整图像大小并预处理
+    image = image.resize((224, 224))
     inputs = feature_extractor(images=image, return_tensors="pt")
-
-    # 模型推理（启用掩码）
+    
+    # 模型推理
     with torch.no_grad():
         outputs = model(**inputs)
-
-    # 获取掩码和重建的像素
-    mask = outputs.mask  # [1, 196]
-    reconstructed_pixels = outputs.logits  # [1, 196, 768]
-
-    # 将掩码应用到原始图像（模拟掩码效果）
-    patch_size = model.config.patch_size
-    image_np = np.array(image)
-    h, w = image_np.shape[0], image_np.shape[1]
-    num_patches_h = h // patch_size
-    num_patches_w = w // patch_size
-
-    # 创建一个掩码图像（黑色表示被掩码的部分）
-    mask_image = np.zeros_like(image_np)
-    mask = mask[0].reshape(num_patches_h, num_patches_w)
-    for i in range(num_patches_h):
-        for j in range(num_patches_w):
-            if mask[i, j] == 1:  # 被掩码的patch
-                mask_image[
-                    i * patch_size : (i + 1) * patch_size,
-                    j * patch_size : (j + 1) * patch_size,
-                ] = 0
-            else:
-                mask_image[
-                    i * patch_size : (i + 1) * patch_size,
-                    j * patch_size : (j + 1) * patch_size,
-                ] = image_np[
-                    i * patch_size : (i + 1) * patch_size,
-                    j * patch_size : (j + 1) * patch_size,
-                ]
-
-    # 可视化结果（原始图像 + 掩码图像）
-    fig, axes = plt.subplots(1, 2, figsize=(10, 5))
-    axes[0].imshow(image_np)
-    axes[0].set_title("Original Image")
-    axes[0].axis("off")
-    axes[1].imshow(mask_image)
-    axes[1].set_title("Masked Image (MAE Input)")
-    axes[1].axis("off")
-    plt.tight_layout()
-    return fig
+    
+    # 获取掩码（14x14）
+    mask = outputs.mask[0].reshape(14, 14).cpu().numpy()
+    
+    # 可视化
+    plt.imshow(mask, cmap="gray")
+    plt.title("MAE Mask (14x14)")
+    plt.axis("off")
+    return plt.gcf()
 
 iface = gr.Interface(
     fn=visualize_mae,
     inputs=gr.Image(type="pil"),
     outputs="plot",
-    title="ViT-MAE Masked Image Visualization",
-    description="Upload an image to see how MAE masks patches.",
+    title="ViT-MAE Mask Visualization",
+    description="Upload an image to see the MAE masking pattern.",
 )
-
 iface.launch()