fix app

Dockerfile

@@ -0,0 +1,24 @@
+# Use an NVIDIA CUDA base image with CUDA 11.8 and Ubuntu 20.04
+FROM mhamilton723/featup:latest
+
+# Set a working directory
+WORKDIR /app
+
+RUN pip3 install gradio
+
+# Copy your application files into the container
+COPY . /app
+
+# Expose the port Streamlit will run on
+EXPOSE 7860
+
+RUN mkdir -m 700 flagged
+
+ENV PYTHONUNBUFFERED=1 \
+	GRADIO_ALLOW_FLAGGING=never \
+	GRADIO_NUM_PORTS=1 \
+	GRADIO_SERVER_NAME=0.0.0.0 \
+	SYSTEM=spaces
+
+# Set the command to run your Streamlit app
+CMD ["python3", "app.py"]

README.md

@@ -1,13 +1,10 @@
 ---
 title: FeatUp
 emoji: 👣⬆️
-colorFrom: green
-colorTo: indigo
-sdk: streamlit
-sdk_version: 1.32.2
-app_file: app.py
+colorFrom: blue
+colorTo: purple
+sdk: docker
 pinned: false
-license: mit
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,46 +1,66 @@
-import streamlit as st
+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
 
-# Assuming the necessary packages (featup, clip, etc.) are installed and accessible
-from featup.util import norm, unnorm
-from featup.plotting import plot_feats
 
-# Setup - ensure the repository content is accessible in the environment
+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")
+    remove_axes(ax)
+    plt.tight_layout()
+    plt.close(fig)  # Close plt to avoid additional empty plots
+    return fig
 
-# Streamlit UI
-st.title("Feature Upsampling Demo")
 
-# File uploader
-uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
-if uploaded_file is not None:
-    image = Image.open(uploaded_file).convert("RGB")
 
-    # 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) # Assuming CUDA is available, .cuda()
+if __name__ == "__main__":
+    os.environ['TORCH_HOME'] = '/tmp/.cache'
 
-    # Model selection
-    model_option = st.selectbox(
-        'Choose a model for feature upsampling',
-        ('dino16', 'dinov2', 'clip', 'resnet50')
-    )
+    options = ['dino16','vit', 'dinov2', 'clip', 'resnet50']
+    image_input = gr.Image(label="Choose an image to featurize", type="pil", image_mode='RGB')
+    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}
+
+    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()
 
-    if st.button('Upsample Features'):
         # Load the selected model
-        upsampler = torch.hub.load("mhamilton723/FeatUp", model_option).cuda()
+        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.")
 
-        # Plotting - adjust the plot_feats function or find an alternative to display images in Streamlit
-        # This step will likely need customization to display within Streamlit's interface
-        plot_feats(unnorm(image_tensor)[0], lr_feats[0], hr_feats[0])
+    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)
 

pre-requirements.txt

@@ -1,15 +0,0 @@
--i https://download.pytorch.org/whl/cu118
---extra-index-url https://pypi.org/simple/
-torch
-torchvision
-torchaudio
-kornia
-omegaconf
-pytorch-lightning
-torchvision
-tqdm
-torchmetrics
-scikit-learn
-numpy
-matplotlib
-git+https://github.com/mhamilton723/CLIP.git

requirements.txt

@@ -1 +0,0 @@
-git+https://github.com/mhamilton723/FeatUp