colormap, slider

app.py

@@ -64,7 +64,7 @@ def download_bin(model):
         os.rename(model_bin, bin_path)
 
 
-def run_attn(input_img, model="vc1-large",fusion="min"):
+def run_attn(input_img, model="vc1-large",fusion="min",slider=0):
     download_bin(model)
     model, embedding_dim, transform, metadata = get_model(model)
     if input_img.shape[0] != 3:
@@ -76,7 +76,7 @@ def run_attn(input_img, model="vc1-large",fusion="min"):
     input_img = resize_transform(input_img)
     x = transform(input_img)
 
-    attention_rollout = VITAttentionGradRollout(model,head_fusion=fusion)
+    attention_rollout = VITAttentionGradRollout(model,head_fusion=fusion,discard_ratio=slider)
     
     y = model(x)
     mask = attention_rollout.get_attn_mask()
@@ -96,12 +96,12 @@ input_img = gr.Image(shape=(250,250))
 input_button = gr.Radio(["min", "max", "mean"], value="min",label="Attention Head Fusion", info="How to combine the last layer attention across all 12 heads of the transformer.")
 output_img = gr.Image(shape=(250,250))
 output_plot = gr.Plot()
-css = ".output-image, .input-image, .image-preview {height: 600px !important}"
-
+css = "#component-3, .input-image, .image-preview {height: 240px !important}"
+slider = gr.Slider(0, 1)
 markdown ="This is a demo for the Visual Cortex models. When passed an image input, it displays the attention of the last layer of the transformer.\n \
          The user can decide how the attention heads will be combined. \
          Along with the attention heatmap, it also displays the embedding values reshaped to a 16x48 or 16x64 grid."
 demo = gr.Interface(fn=run_attn, title="Visual Cortex Large Model", description=markdown,
-                    examples=[[os.path.join('./imgs',x),None,None]for x in os.listdir(os.path.join(os.getcwd(),'imgs')) if 'jpg' in x],
-                    inputs=[input_img,model_type,input_button],outputs=[output_img,output_plot],css=css)
+                    examples=[[os.path.join('./imgs',x),None,None,None]for x in os.listdir(os.path.join(os.getcwd(),'imgs')) if 'jpg' in x],
+                    inputs=[input_img,model_type,input_button,slider],outputs=[output_img,output_plot],css=css)
 demo.launch()

attn_helper.py

@@ -9,7 +9,7 @@ def overlay_attn(original_image,mask):
     # Colormap and alpha for attention mask
     # COLORMAP_OCEAN
     # COLORMAP_OCEAN
-    colormap_attn, alpha_attn = cv2.COLORMAP_JET, 1 #0.85
+    colormap_attn, alpha_attn = cv2.COLORMAP_VIRIDIS, 1 #0.85
     
     # Resize mask to original image size
     w, h = original_image.shape[0], original_image.shape[1]