fixed file saving issues

app.py

@@ -79,16 +79,12 @@ def predict_tensor(img_tensor):
         probs = torch.softmax(logits, dim=0)
         topk_prob, topk_id = torch.topk(probs, 3)
 
-with nopdb.capture_call(vision_transformer.blocks[5].attn.forward) as attn_call:
-    predict_tensor(img_transformed)
-
 random_image = PIL.Image.open('examples/TomatoYellowCurlVirus6.JPG').convert('RGB')
 
 def plot_attention(image, layer_num):
     """Given an input image, plot the average attention weight given to each image patch by each attention head."""
+    attention_map_outputs = []
     input_data = data_transforms(image)
-    print(input_data.shape)
-    print(layer_num)
     with nopdb.capture_call(vision_transformer.blocks[int(layer_num)-1].attn.forward) as attn_call:
         predict_tensor(img_transformed)
     attn = attn_call.locals['attn'][0]
@@ -99,9 +95,9 @@ def plot_attention(image, layer_num):
             h_weights = h_weights.mean(axis=-2)  # average over all attention keys
             h_weights = h_weights[1:]  # skip the [class] token
             attention_block_num += 1
-            plot_weights(input_data, h_weights, attention_block_num)
-    attention_maps = glob.glob('storage/*.png')
-    return attention_maps
+            output_img = plot_weights(input_data, h_weights, attention_block_num)
+            attention_map_outputs.append(output_img)
+    return attention_map_outputs
 
 def plot_weights(input_data, patch_weights, num_attention_block):
     """Display the image: Brighter the patch, higher is the attention."""
@@ -113,10 +109,9 @@ def plot_weights(input_data, patch_weights, num_attention_block):
         plot[:, y:y + 16, x:x + 16] *= patch_weights[i]
     attn_map_img = inv_transform(plot, normalize=False)
     attn_map_img = attn_map_img.resize((224, 224), Image.ANTIALIAS)
-    attn_map_img.save(f"storage/attention_map_{num_attention_block}.png", "PNG")
+    return attn_map_img
 
-DEFAULT_LAYER = "6"
-attention_maps = plot_attention(random_image, DEFAULT_LAYER)
+attention_maps = plot_attention(random_image, 6)
 
 title_classify = "Image Based Plant Disease Identification 🍃🤓"
 
@@ -147,7 +142,7 @@ attention_interface = gr.Interface(
     fn=plot_attention,
     inputs=[gr.Image(type="pil", label="Image"),
             gr.Dropdown(choices=["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12"], 
-                        value=DEFAULT_LAYER, label="Attention Layer")],
+                        label="Attention Layer", value="6")],
     outputs=gr.Gallery(value=attention_maps, label="Attention Maps").style(grid=(3, 4)),
     examples=example_list,
     title=title_attention,