chore: Add top 75 PCA mode to image reconstruction interface

app.py

@@ -12,6 +12,7 @@ from model import MAE_ViT, MAE_Encoder, MAE_Decoder, MAE_Encoder_FeatureExtracto
 
 path_1  = [['images/cat.jpg'], ['images/dog.jpg'], ['images/horse.jpg'], ['images/airplane.jpg'], ['images/truck.jpg']]
 path_2  = [['images/cat.jpg'], ['images/dog.jpg'], ['images/horse.jpg'], ['images/airplane.jpg'], ['images/truck.jpg']]
+path_3  = [['images/cat.jpg'], ['images/dog.jpg'], ['images/horse.jpg'], ['images/airplane.jpg'], ['images/truck.jpg']]
 device = torch.device("cpu")
 
 model_name = "model/no_mode/vit-t-mae-pretrain.pt"
@@ -20,9 +21,14 @@ model_no_mode.eval()
 model_no_mode.to(device)
 
 model_name = "model/bottom_25/vit-t-mae-pretrain.pt"
-model_pca_mode = torch.load(model_name, map_location='cpu')
-model_pca_mode.eval()
-model_pca_mode.to(device)
+model_pca_mode_bottom = torch.load(model_name, map_location='cpu')
+model_pca_mode_bottom.eval()
+model_pca_mode_bottom.to(device)
+
+model_name = "model/top_75/vit-t-mae-pretrain.pt"
+model_pca_mode_top = torch.load(model_name, map_location='cpu')
+model_pca_mode_top.eval()
+model_pca_mode_top.to(device)
 
 transform = v2.Compose([
         v2.Resize((96, 96)),
@@ -87,12 +93,54 @@ def visualize_single_image_no_mode(image_path):
 
     return np.array(plt.imread("output.png"))
 
-def visualize_single_image_pca_mode(image_path):
+def visualize_single_image_pca_mode_bottom(image_path):
+    img = load_image(image_path, transform).to(device)
+    
+    # Run inference
+    with torch.no_grad():
+        predicted_img, mask = model_pca_mode_bottom(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
+    img = img[0]
+    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"))
+
+def visualize_single_image_pca_mode_top(image_path):
     img = load_image(image_path, transform).to(device)
     
     # Run inference
     with torch.no_grad():
-        predicted_img, mask = model_pca_mode(img)
+        predicted_img, mask = model_pca_mode_top(img)
     
     # Convert the tensor back to a displayable image
     # masked image
@@ -145,6 +193,15 @@ outputs_image_2 = [
     gr.components.Image(type="numpy", label="Output Image"),
 ]
 
+inputs_image_3 = [
+    gr.components.Image(type="filepath", label="Input Image"),
+]
+
+outputs_image_3 = [
+    gr.components.Image(type="numpy", label="Output Image"),
+]
+
+
 inference_no_mode = gr.Interface(
     fn=visualize_single_image_no_mode,
     inputs=inputs_image_1,
@@ -155,8 +212,8 @@ inference_no_mode = gr.Interface(
     description="This is a demo of the MAE-ViT model for image reconstruction.",
 )
 
-inference_pca_mode = gr.Interface(
-    fn=visualize_single_image_pca_mode,
+inference_pca_mode_bottom = gr.Interface(
+    fn=visualize_single_image_pca_mode_bottom,
     inputs=inputs_image_2,
     outputs=outputs_image_2,
     examples=path_2,
@@ -164,7 +221,16 @@ inference_pca_mode = gr.Interface(
     description="This is a demo of the MAE-ViT model for image reconstruction.",
 )
 
+inference_pca_mode_top = gr.Interface(
+    fn=visualize_single_image_pca_mode_top,
+    inputs=inputs_image_3,
+    outputs=outputs_image_3,
+    examples=path_3,
+    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],
+    [inference_no_mode, inference_pca_mode_bottom, inference_pca_mode_top],
     tab_names=['Normal Mode', 'PCA Mode']
 ).queue().launch()

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

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