init

app.py

@@ -0,0 +1,124 @@
+import matplotlib.pyplot as plt
+import torch
+import torchvision.transforms as T
+from PIL import Image
+import gradio as gr
+from featup.util import norm, unnorm, pca, remove_axes
+from pytorch_lightning import seed_everything
+import os
+import requests
+import csv
+import spaces
+
+
+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)], 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
+
+
+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 = "/tmp/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.")
+
+os.environ['TORCH_HOME'] = '/tmp/.cache'
+os.environ['GRADIO_EXAMPLES_CACHE'] = '/tmp/gradio_cache'
+csv.field_size_limit(100000000)
+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}
+
+
+@spaces.GPU
+def upsample_features(image, model_option):
+    # Image preprocessing
+    input_size = 224
+    transform = T.Compose([
+        T.Resize(input_size),
+        T.CenterCrop((input_size, input_size)),
+        T.ToTensor(),
+        norm
+    ])
+    image_tensor = transform(image).unsqueeze(0).cuda()
+
+    # Load the selected model
+    upsampler = models[model_option].cuda()
+    hr_feats = upsampler(image_tensor)
+    lr_feats = upsampler.model(image_tensor)
+    upsampler.cpu()
+
+    return plot_feats(unnorm(image_tensor)[0], lr_feats[0], hr_feats[0])
+
+
+demo = gr.Interface(fn=upsample_features,
+                    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.",
+                    examples=[
+                        ["/tmp/sample_images/skate.jpg", "dino16"],
+                        ["/tmp/sample_images/car.jpg", "dinov2"],
+                        ["/tmp/sample_images/plant.png", "dino16"],
+                    ]
+                    )
+
+demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)