add attention

app.py

@@ -87,6 +87,8 @@ if file is not None:
     # initiliaze the smoothing parameters
     smooth_scale = st.sidebar.slider("Smooth scale", min_value=0.1, max_value =1.0, step = 0.1)
     
+    smooth_thres = st.sidebar.slider("Smooth thres", min_value=0.01, max_value =1.0, step = 0.01)
+    
     smooth_size = st.sidebar.slider("Smooth size", min_value=1, max_value =10)
     
     smooth_iter = st.sidebar.slider("Smooth iter", min_value=1, max_value =10)
@@ -129,7 +131,7 @@ if file is not None:
             attention = outputs.attentions[-1][0]
             
             # Let us recuperate the attention image
-            attention_image = get_attention(image, attention, size = (224, 224), patch_size = (14, 14), scale = scale, head = head, smooth_scale = smooth_scale, smooth_size = smooth_size, smooth_iter = smooth_iter)
+            attention_image = get_attention(image, attention, size = (224, 224), patch_size = (14, 14), scale = scale, head = head, smooth_scale = smooth_scale, smooth_thres=smooth_thres, smooth_size = smooth_size, smooth_iter = smooth_iter)
 
             # Let us transform the attention image to a opencv image
             attention_image = cv2.cvtColor(attention_image.astype('float32'), cv2.COLOR_RGB2BGR)

fake_face_detection/metrics/make_predictions.py

@@ -17,7 +17,7 @@ import numpy as np
 import torch
 import os
 
-def get_attention(image: Union[str, JpegImageFile], attention: torch.Tensor, size: tuple, patch_size: tuple, scale: int = 50, head: int = 1, smooth_iter: int = 2, smooth_scale: float = 0.2, smooth_size = 5):
+def get_attention(image: Union[str, JpegImageFile], attention: torch.Tensor, size: tuple, patch_size: tuple, scale: int = 50, head: int = 1, smooth_iter: int = 2, smooth_thres: float = 0.01, smooth_scale: float = 0.2, smooth_size = 5):
     
     # recuperate the image as a numpy array
     if isinstance(image, str):
@@ -49,7 +49,7 @@ def get_attention(image: Union[str, JpegImageFile], attention: torch.Tensor, siz
     attention = attention.reshape(size[0], size[1], 1)
     
     # add the smoothest attention
-    attention = smooth_attention(attention, smooth_iter, smooth_scale, smooth_size)
+    attention = smooth_attention(attention, smooth_iter, smooth_thres, smooth_scale, smooth_size)
     
     # recuperate the result
     attention_image = img / 255 * attention.numpy() * scale
@@ -69,6 +69,7 @@ def make_predictions(test_dataset: FakeFaceDetectionDataset,
                      show: bool = True, 
                      head: int = 1,
                      smooth_iter: int = 2,
+                     smooth_thres: float = 0.01,
                      smooth_scale: float = 0.2,
                      smooth_size = 5):
     """Make predictions with a vision transformer model
@@ -86,6 +87,7 @@ def make_predictions(test_dataset: FakeFaceDetectionDataset,
         show (bool, optional): A boolean value indicating if we want to recuperate the figure. Defaults to True.
         head (int, optional): The head number. Defaults to 1.
         smooth_iter (int, optional): The number of iterations for the smoothest attention. Defaults to 2.
+        smooth_thres (float, optional): The threshold for the smoothest attention. Defaults to 0.01.
         smooth_scale (float, optional): The scale for the smoothest attention. Defaults to 0.2.
         smooth_size ([type], optional): The size for the smoothest attention. Defaults to 5.
 
@@ -153,7 +155,7 @@ def make_predictions(test_dataset: FakeFaceDetectionDataset,
         
         for i in range(len(images)):
             
-            attention_image = get_attention(images[i], predictions['attentions'][i], size, patch_size)
+            attention_image = get_attention(images[i], predictions['attentions'][i], size, patch_size, attention_scale, head, smooth_iter, smooth_thres, smooth_scale, smooth_size)
         
             axes[i].imshow(attention_image)