adding smooth attention

app.py

@@ -84,6 +84,13 @@ if file is not None:
     
     left.markdown("""---""")
     
+    # initiliaze the smoothing parameters
+    smooth_scale = st.sidebar.slider("Smooth scale", min_value=0.1, max_value =1.0, step = 0.1)
+    
+    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)
+    
     # add a side for the scaler and the head number
     scale = st.sidebar.slider("Attention scale", min_value=30, max_value =200)
 
@@ -122,7 +129,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)
+            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)
 
             # 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

@@ -1,6 +1,7 @@
 
 from fake_face_detection.data.fake_face_dataset import FakeFaceDetectionDataset
 from fake_face_detection.metrics.compute_metrics import compute_metrics
+from fake_face_detection.smoothest_attention import smooth_attention
 from torch.utils.tensorboard import SummaryWriter
 from PIL.JpegImagePlugin import JpegImageFile
 from torch.utils.data import DataLoader
@@ -16,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):
+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):
     
     # recuperate the image as a numpy array
     if isinstance(image, str):
@@ -47,6 +48,9 @@ def get_attention(image: Union[str, JpegImageFile], attention: torch.Tensor, siz
     # let us reshape the attention to the right size
     attention = attention.reshape(size[0], size[1], 1)
     
+    # add the smoothest attention
+    attention = smooth_attention(attention, smooth_iter, smooth_scale, smooth_size)
+    
     # recuperate the result
     attention_image = img / 255 * attention.numpy() * scale
     
@@ -63,7 +67,10 @@ def make_predictions(test_dataset: FakeFaceDetectionDataset,
                      figsize: tuple = (24, 24),
                      attention_scale: int = 50,
                      show: bool = True, 
-                     head: int = 1):
+                     head: int = 1,
+                     smooth_iter: int = 2,
+                     smooth_scale: float = 0.2,
+                     smooth_size = 5):
     """Make predictions with a vision transformer model
 
     Args:
@@ -78,6 +85,9 @@ def make_predictions(test_dataset: FakeFaceDetectionDataset,
         attention_scale (int, optional): The attention scale. Defaults to 50.
         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_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.
 
     Returns:
         Union[Tuple[pd.DataFrame, dict], Tuple[pd.DataFame, dict, figure]]: The return prediction and the metrics

fake_face_detection/utils/split_data.py

@@ -27,11 +27,11 @@ def split_data_from_dir(path: str, new_path: str, test_size: float = 0.2, valid_
             # let us recuperate the files' paths in it
             images = os.listdir(dir_path)
             
-            # let us split the data between training and test set
-            train_set, test_set = train_test_split(images, test_size = test_size)
+            # let us split the data between training and test + valid sets
+            train_set, test_valid_set = train_test_split(images, test_size = test_size + valid_size)
             
-            # let us split the training set between training and validation set
-            train_set, valid_set = train_test_split(train_set, test_size = valid_size)
+            # let us split the test + valid sets between test and valid sets
+            test_set, valid_set = train_test_split(test_valid_set, test_size = valid_size)
             
             # let us create the train test and valid directories
             if not os.path.exists(os.path.join(os.path.join(new_path, 'train'), dir_)) or\