update demo

.gitignore

@@ -1,2 +1,3 @@
 produce_attmaps.py
-__pycache__
+__pycache__
+flagged/

app.py

@@ -1,33 +1,16 @@
 import os
-import PIL
-import ast
-import cv2
-import json
-import torch
-import pickle
-import torchvision
 import numpy as np
 import gradio as gr
-from PIL import Image
-from typing import Tuple, Dict
 import matplotlib.pyplot as plt
-from timeit import default_timer as timer
-from torchvision import datasets, transforms
-
-import vision_transformer as vits
 
-'''
-import warnings
-warnings.filterwarnings('ignore')
-
-example_list = [["examples/" + example] for example in os.listdir("examples")]
+import PIL
+from PIL import Image
 
-with open('labels/imagenet1k-simple-labels.json') as f:
-    class_names = json.load(f)
+import torch
+import torchvision
+from torchvision import datasets, transforms
 
-from model import VisionTransformer
-from capture_weights import vit_weights
-'''
+import vision_transformer as vits
 
 arch = "vit_small"
 mode = "simpool"
@@ -40,8 +23,8 @@ checkpoint_key = "teacher"
 
 cm = plt.get_cmap('viridis')
 attn_map_size = 224
-width_display = 300
-height_display = 300
+width_display = 290
+height_display = 290
 
 example_dir = "examples/"
 example_list = [[example_dir + example] for example in os.listdir(example_dir)]
@@ -96,4 +79,4 @@ demo = gr.TabbedInterface([attention_interface],
                           ["Visualize Attention Maps"], title="SimPool Attention Map Visualizer 🌌")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)