fix app

README.md

@@ -1,8 +1,8 @@
 ---
 title: FeatUp
 emoji: 👣⬆️
-colorFrom: blue
-colorTo: purple
+colorFrom: pink
+colorTo: yellow
 sdk: docker
 pinned: false
 ---

app.py

@@ -11,30 +11,80 @@ import os
 def plot_feats(image, lr, hr):
     assert len(image.shape) == len(lr.shape) == len(hr.shape) == 3
     seed_everything(0)
-    [lr_feats_pca, hr_feats_pca], _ = pca([lr.unsqueeze(0), hr.unsqueeze(0)])
-    fig, ax = plt.subplots(1, 3, figsize=(15, 5))
-    ax[0].imshow(image.permute(1, 2, 0).detach().cpu())
-    ax[0].set_title("Image")
-    ax[1].imshow(lr_feats_pca[0].permute(1, 2, 0).detach().cpu())
-    ax[1].set_title("Original Features")
-    ax[2].imshow(hr_feats_pca[0].permute(1, 2, 0).detach().cpu())
-    ax[2].set_title("Upsampled Features")
+    [lr_feats_pca, hr_feats_pca], _ = pca([lr.unsqueeze(0), hr.unsqueeze(0)], dim=9)
+    fig, ax = plt.subplots(3, 3, figsize=(15, 15))
+    ax[0, 0].imshow(image.permute(1, 2, 0).detach().cpu())
+    ax[1, 0].imshow(image.permute(1, 2, 0).detach().cpu())
+    ax[2, 0].imshow(image.permute(1, 2, 0).detach().cpu())
+
+    ax[0, 0].set_title("Image", fontsize=22)
+    ax[0, 1].set_title("Original", fontsize=22)
+    ax[0, 2].set_title("Upsampled Features", fontsize=22)
+
+    ax[0, 1].imshow(lr_feats_pca[0, :3].permute(1, 2, 0).detach().cpu())
+    ax[0, 0].set_ylabel("PCA Components 1-3", fontsize=22)
+    ax[0, 2].imshow(hr_feats_pca[0, :3].permute(1, 2, 0).detach().cpu())
+
+    ax[1, 1].imshow(lr_feats_pca[0, 3:6].permute(1, 2, 0).detach().cpu())
+    ax[1, 0].set_ylabel("PCA Components 4-6", fontsize=22)
+    ax[1, 2].imshow(hr_feats_pca[0, 3:6].permute(1, 2, 0).detach().cpu())
+
+    ax[2, 1].imshow(lr_feats_pca[0, 6:9].permute(1, 2, 0).detach().cpu())
+    ax[2, 0].set_ylabel("PCA Components 7-9", fontsize=22)
+    ax[2, 2].imshow(hr_feats_pca[0, 6:9].permute(1, 2, 0).detach().cpu())
+
     remove_axes(ax)
     plt.tight_layout()
     plt.close(fig)  # Close plt to avoid additional empty plots
     return fig
 
 
+if __name__ == "__main__":
+    import requests
+    import os
+
 
+    def download_image(url, save_path):
+        response = requests.get(url)
+        with open(save_path, 'wb') as file:
+            file.write(response.content)
+
+    base_url = "https://marhamilresearch4.blob.core.windows.net/feature-upsampling-public/sample_images/"
+    sample_images_urls = {
+        "skate.jpg": base_url + "skate.jpg",
+        "car.jpg": base_url + "car.jpg",
+        "plant.png": base_url + "plant.png",
+    }
+
+    sample_images_dir = "sample_images"
+
+    # Ensure the directory for sample images exists
+    os.makedirs(sample_images_dir, exist_ok=True)
+
+    # Download each sample image
+    for filename, url in sample_images_urls.items():
+        save_path = os.path.join(sample_images_dir, filename)
+        # Download the image if it doesn't already exist
+        if not os.path.exists(save_path):
+            print(f"Downloading {filename}...")
+            download_image(url, save_path)
+        else:
+            print(f"{filename} already exists. Skipping download.")
 
-if __name__ == "__main__":
     os.environ['TORCH_HOME'] = '/tmp/.cache'
 
-    options = ['dino16','vit', 'dinov2', 'clip', 'resnet50']
-    image_input = gr.Image(label="Choose an image to featurize", type="pil", image_mode='RGB')
+    options = ['dino16', 'vit', 'dinov2', 'clip', 'resnet50']
+
+    image_input = gr.Image(label="Choose an image to featurize",
+                           height=480,
+                           type="pil",
+                           image_mode='RGB',
+                           sources=['upload', 'webcam', 'clipboard']
+                           )
     model_option = gr.Radio(options, value="dino16", label='Choose a backbone to upsample')
 
-    models = {o:torch.hub.load("mhamilton723/FeatUp", o) for o in options}
+    models = {o: torch.hub.load("mhamilton723/FeatUp", o) for o in options}
+
 
     def upsample_features(image, model_option):
         # Image preprocessing
@@ -60,7 +110,13 @@ if __name__ == "__main__":
                         inputs=[image_input, model_option],
                         outputs="plot",
                         title="Feature Upsampling Demo",
-                        description="This demo allows you to upsample features of an image using selected models.")
+                        description="This demo allows you to upsample features of an image using selected models.",
+                        examples=[
+                            ["sample_images/skate.jpg", "dino16"],
+                            ["sample_images/car.jpg", "dinov2"],
+                            ["sample_images/plant.png", "dino16"],
+                        ]
 
-    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)
+                        )
 
+    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)