bottom 25 model added to visualize in seperate tab

app.py

@@ -11,12 +11,17 @@ from torchvision.transforms import v2
 from model import MAE_ViT, MAE_Encoder, MAE_Decoder, MAE_Encoder_FeatureExtractor
 
 path  = [['images/cat.jpg'], ['images/dog.jpg'], ['images/horse.jpg'], ['images/airplane.jpg'], ['images/truck.jpg']]
-model_name = "vit-t-mae-pretrain.pt"
-model = torch.load(model_name, map_location='cpu')
-
-model.eval()
 device = torch.device("cpu")
-model.to(device)
+
+model_name = "model/no_mode/vit-t-mae-pretrain.pt"
+model_no_mode = torch.load(model_name, map_location='cpu')
+model_no_mode.eval()
+model_no_mode.to(device)
+
+model_name = "model/bottom_256/vit-t-mae-pretrain.pt"
+model_pca_mode = torch.load(model_name, map_location='cpu')
+model_pca_mode.eval()
+model_pca_mode.to(device)
 
 transform = v2.Compose([
         v2.Resize((32, 32)),
@@ -39,13 +44,12 @@ def show_image(img, title):
     plt.title(title)
 
 # Visualize a Single Image
-def visualize_single_image(image_path):
+def visualize_single_image_no_mode(image_path):
     img = load_image(image_path, transform).to(device)
     
     # Run inference
-    model.eval()
     with torch.no_grad():
-        predicted_img, mask = model(img)
+        predicted_img, mask = model_no_mode(img)
     
     # Convert the tensor back to a displayable image
     # masked image
@@ -82,6 +86,47 @@ def visualize_single_image(image_path):
 
     return np.array(plt.imread("output.png"))
 
+def visualize_single_image_pca_mode(image_path):
+    img = load_image(image_path, transform).to(device)
+    
+    # Run inference
+    with torch.no_grad():
+        predicted_img, mask = model_pca_mode(img)
+    
+    # Convert the tensor back to a displayable image
+    # masked image
+    im_masked = img * (1 - mask)
+
+    # MAE reconstruction pasted with visible patches
+    im_paste = img * (1 - mask) + predicted_img * mask
+
+    # remove the batch dimension
+    im_masked = im_masked[0]
+    predicted_img = predicted_img[0]
+    im_paste = im_paste[0]
+    
+    # make the plt figure larger
+    plt.figure(figsize=(18, 8))
+
+    plt.subplot(1, 4, 1)
+    show_image(img, "original")
+
+    plt.subplot(1, 4, 2)
+    show_image(im_masked, "masked")
+
+    plt.subplot(1, 4, 3)
+    show_image(predicted_img, "reconstruction")
+
+    plt.subplot(1, 4, 4)
+    show_image(im_paste, "reconstruction + visible")
+
+    plt.tight_layout()
+
+    # convert the plt figure to a numpy array
+    plt.savefig("output.png")
+
+    return np.array(plt.imread("output.png"))
+
 inputs_image = [
     gr.components.Image(type="filepath", label="Input Image"),
 ]
@@ -90,11 +135,26 @@ outputs_image = [
     gr.components.Image(type="numpy", label="Output Image"),
 ]
 
-gr.Interface(
-    fn=visualize_single_image,
+inference_no_mode = gr.Interface(
+    fn=visualize_single_image_no_mode,
     inputs=inputs_image,
     outputs=outputs_image,
     examples=path,
+    cache_examples = False,
     title="MAE-ViT Image Reconstruction",
     description="This is a demo of the MAE-ViT model for image reconstruction.",
-).launch()
+)
+
+inference_pca_mode = gr.Interface(
+    fn=visualize_single_image_pca_mode,
+    inputs=inputs_image,
+    outputs=outputs_image,
+    examples=path,
+    title="MAE-ViT Image Reconstruction",
+    description="This is a demo of the MAE-ViT model for image reconstruction.",
+)
+
+gr.TabbedInterface(
+    [inference_no_mode, inference_pca_mode],
+    tab_names=['Normal Mode', 'PCA Mode']
+).queue().launch()

model/bottom_25/vit-t-mae-pretrain.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e00a227576d93ed464e783ba281866b7441a37123c649a55e648d6e5553b66b0
+size 28973864