noise option added

Author: luciana-munhos

interface.py

@@ -89,7 +89,9 @@ def __init__(
                     load_model_path: str = None,
                     custom_augment_figuratif=None,
                     custom_augment_abstrait=None,
-                    n_transforms_augmented = 2
+                    n_transforms_augmented = 2,
+                    noise: bool = False,
+                    noise_std: float = 0.1
                 ):
 
         # Profiler
@@ -151,6 +153,8 @@ def __init__(
         self.augmentation = augmentation
         self.input_size = input_size
         self.padding = padding
+        self.noise = noise
+        self.noise_std = noise_std
 
         # Initialize the dataset and dataloader
 
@@ -161,23 +165,29 @@ def __init__(
                                         image_input_size=self.input_size,
                                         custom_augment_abstrait=custom_augment_abstrait,
                                         custom_augment_figuratif=custom_augment_figuratif,
-                                        n_transforms_augmented=n_transforms_augmented)
-        
+                                        n_transforms_augmented=n_transforms_augmented,
+                                        noise=self.noise,
+                                        noise_std=self.noise_std)
+
         self.dataset_val = PaintingsDataset(self.data_path+'val/',
                                         augment=False, 
                                         transform=self.transform, 
                                         padding=self.padding, 
                                         image_input_size=self.input_size,
                                         custom_augment_abstrait=None,
-                                        custom_augment_figuratif=None,)
-        
+                                        custom_augment_figuratif=None,
+                                        noise=self.noise,
+                                        noise_std=self.noise_std)
+
         self.dataset_test = PaintingsDataset(self.data_path+'test/',
                                         augment=False, 
                                         transform=self.transform, 
                                         padding=self.padding, 
                                         image_input_size=self.input_size,
                                         custom_augment_abstrait=None,
-                                        custom_augment_figuratif=None,)
+                                        custom_augment_figuratif=None,
+                                        noise=self.noise,
+                                        noise_std=self.noise_std)
 
         n_figurative = self.dataset_train.len_figurative #only for training
         n_abstract = self.dataset_train.len_abstract

paintings_dataset.py

@@ -73,7 +73,7 @@ class PaintingsDataset(Dataset):
     In case of `transform` being True, the dataset will apply transformations to the images or 
     a custom transformation can be passed as a parameter named `custom_transform`(if None we will use custom).
     """
-    def __init__(self, data_path, augment= False, transform=False, custom_augment_figuratif=None,custom_augment_abstrait=None, padding: PaddingOptions = PaddingOptions.ZERO, image_input_size: int = 224,n_transforms_augmented=2):
+    def __init__(self, data_path, augment= False, transform=False, custom_augment_figuratif=None,custom_augment_abstrait=None, padding: PaddingOptions = PaddingOptions.ZERO, image_input_size: int = 224,n_transforms_augmented=2, noise: bool = False, noise_std: float = 0.1):
 
         # Path to the data directory
         self.data_path = data_path
@@ -106,6 +106,10 @@ def __init__(self, data_path, augment= False, transform=False, custom_augment_fi
         self.padding = padding
         self.image_input_size = image_input_size
 
+        # Noise configuration
+        self.noise = noise
+        self.noise_std = noise_std
+
 
     def __len__(self):
         return self.total_length
@@ -176,6 +180,17 @@ def __getitem__(self, idx: int)->Dict[str, Union[torch.Tensor, int]]:
                                                                 size=(max(imh, int(self.image_input_size / 2 + 1)), max(imw, int(self.image_input_size/2 + 1))), 
                                                                 mode='nearest').squeeze(0).clone()
         output['image'] = output['image'].float() / 255.0  # Normalize the image to [0, 1]
+
+        # Noise configuration
+        if self.noise:
+            noise = torch.randn_like(output['image']) * self.noise_std
+            output['image'] = torch.clamp(output['image'] + noise, 0.0, 1.0)
+
+        # Apply noise to transformed image (if it exists and noise enabled)
+        if self.noise and self.transform:
+            noise_transformed = torch.randn_like(output['transformed_image']) * self.noise_std
+            output['transformed_image'] = torch.clamp(output['transformed_image'] + noise_transformed, 0.0, 1.0)
+            
         return output
 
     def change_padding(self, padding: Literal['zero', 'mirror', 'replicate'], image_input_size: int = 224):