reorganize

utils/dataset.py → dataset.py

@@ -15,7 +15,7 @@ from sklearn.model_selection import StratifiedShuffleSplit
 if not os.path.exists('README.md'):  # set pwd to root
     os.chdir('..')
 
-from utils.splitting import split_img
+from utils.dataset_utils import split_img, list_images_in_dir, load_image
 from utils.base import np2torch, torch2np, b2_download_folder
 
 IMAGE_FILE_TYPES = ['dng', 'png', 'tif', 'tiff']
@@ -41,24 +41,6 @@ def get_dataset(name, I_ratio=1.0):
     raise ValueError(name)
 
 
-def load_image(path):
-    file_type = path.split('.')[-1].lower()
-    if file_type == 'dng':
-        img = rawpy.imread(path).raw_image_visible
-    elif file_type == 'tiff' or file_type == 'tif':
-        img = np.array(tiff.imread(path), dtype=np.float32)
-    else:
-        img = np.array(Image.open(path), dtype=np.float32)
-    return img
-
-
-def list_images_in_dir(path):
-    image_list = [os.path.join(path, img_name)
-                  for img_name in sorted(os.listdir(path))
-                  if img_name.split('.')[-1].lower() in IMAGE_FILE_TYPES]
-    return image_list
-
-
 class ImageFolderDataset(Dataset):
     """Creates a dataset of images in img_dir and corresponding masks in mask_dir.
     Corresponding mask files need to contain the filename of the image.
@@ -166,18 +148,20 @@ class ImageFolderDatasetSegmentation(Dataset):
 
         return img, mask
 
+
 class MultiIntensity(Dataset):
     """Wrap datasets with different intesities
 
     Args:
         datasets (list): list of datasets to wrap
     """
+
     def __init__(self, datasets):
         self.dataset = datasets[0]
 
-        for d in range(1,len(datasets)):
-            self.dataset.images = self.dataset.images+datasets[d].images
-            self.dataset.labels = self.dataset.labels+datasets[d].labels
+        for d in range(1, len(datasets)):
+            self.dataset.images = self.dataset.images + datasets[d].images
+            self.dataset.labels = self.dataset.labels + datasets[d].labels
 
     def __len__(self):
         return len(self.dataset)
@@ -191,6 +175,7 @@ class MultiIntensity(Dataset):
             x = self.transform(x)
         return x, y
 
+
 class Subset(Dataset):
     """Define a subset of a dataset by only selecting given indices.
 
@@ -228,8 +213,8 @@ class DroneDatasetSegmentationFull(ImageFolderDatasetSegmentation):
     camera_parameters = black_level, white_balance, colour_matrix
 
     def __init__(self, I_ratio=1.0, transform=None, force_download=False, bits=16):
-                
-        assert I_ratio in [0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 1.0] 
+
+        assert I_ratio in [0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 1.0]
 
         img_dir = f'data/drone/images_full/raw_scale{int(I_ratio*100):03d}'
         mask_dir = 'data/drone/masks_full'
@@ -411,7 +396,7 @@ def download_microscopy_dataset(force_download):
 
 
 def unzip_microscopy_images():
-            
+
     if os.path.isfile('data/microscopy/labels/.bzEmpty'):
         os.remove('data/microscopy/labels/.bzEmpty')
 
@@ -421,6 +406,7 @@ def unzip_microscopy_images():
             zip.extractall('data/microscopy/images')
             os.remove(os.path.join('data/microscopy/images', file))
 
+
 def unzip_drone_images():
 
     if os.path.isfile('data/drone/masks_full/.bzEmpty'):
@@ -585,38 +571,3 @@ def check_image_folder_consistency(images, masks):
             f"image file {img_file} file type mismatch. Shoule be: {file_type_images}"
         assert mask_file.split('.')[-1].lower() == file_type_masks, \
             f"image file {mask_file} file type mismatch. Should be: {file_type_masks}"
-
-
-def k_fold(dataset, n_splits: int, seed: int, train_size: float):
-    """Split dataset in subsets for cross-validation 
-
-       Args:
-            dataset (class): dataset to split
-            n_split (int): Number of re-shuffling & splitting iterations.
-            seed (int): seed for k_fold splitting
-            train_size (float): should be between 0.0 and 1.0 and represent the proportion of the dataset to include in the train split.
-       Returns:
-           idxs (list): indeces for splitting the dataset. The list contain n_split pair of train/test indeces.
-    """
-    if hasattr(dataset, 'labels'):
-        x = dataset.images
-        y = dataset.labels
-    elif hasattr(dataset, 'masks'):
-        x = dataset.images
-        y = dataset.masks
-
-    idxs = []
-
-    if dataset.task == 'classification':
-        sss = StratifiedShuffleSplit(n_splits=n_splits, train_size=train_size, random_state=seed)
-
-        for idxs_train, idxs_test in sss.split(x, y):
-            idxs.append((idxs_train.tolist(), idxs_test.tolist()))
-
-    elif dataset.task == 'segmentation':
-        for n in range(n_splits):
-            split_idx = int(len(dataset) * train_size)
-            indices = np.random.permutation(len(dataset))
-            idxs.append((indices[:split_idx].tolist(), indices[split_idx:].tolist()))
-
-    return idxs

ABtesting.py → figures/ABtesting.py

@@ -8,11 +8,11 @@ from torch.utils.data import DataLoader
 from torchvision.transforms import Compose, Normalize
 import torch.nn.functional as F
 
-from utils.dataset import get_dataset, Subset
+from dataset import get_dataset, Subset
 from utils.base import get_mlflow_model_by_name, SmartFormatter
-from processingpipeline.pipeline import RawProcessingPipeline
+from processing.pipeline_numpy import RawProcessingPipeline
 
-from utils.Cperturb import Distortions
+from utils.hendrycks_robustness import Distortions
 
 import segmentation_models_pytorch as smp
 
@@ -20,94 +20,104 @@ import matplotlib.pyplot as plt
 
 parser = argparse.ArgumentParser(description="AB testing, Show Results", formatter_class=SmartFormatter)
 
-#Select experiment
-parser.add_argument("--mode", type=str, default="ABShowImages", choices=('ABMakeTable', 'ABShowTable', 'ABShowImages', 'ABShowAllImages', 'CMakeTable', 'CShowTable', 'CShowImages', 'CShowAllImages'), 
+# Select experiment
+parser.add_argument("--mode", type=str, default="ABShowImages", choices=('ABMakeTable', 'ABShowTable', 'ABShowImages', 'ABShowAllImages', 'CMakeTable', 'CShowTable', 'CShowImages', 'CShowAllImages'),
                     help='R|Choose operation to compute. \n'
-                        'A) Lens2Logit image generation: \n  '
-                        'ABMakeTable: Compute cross-validation metrics results \n  '
-                        'ABShowTable: Plot cross-validation results on a table \n  '
-                        'ABShowImages: Choose a training and testing image to compare different pipelines \n  '
-                        'ABShowAllImages: Plot all possible pipelines \n'
-                        'B) Hendrycks Perturbations, C-type dataset: \n  '
-                        'CMakeTable: For each pipeline, it computes cross-validation metrics for different perturbations  \n  '
-                        'CShowTable: Plot metrics for different pipelines and perturbations \n  '
-                        'CShowImages: Plot an image with a selected a pipeline and perturbation\n  '
-                        'CShowAllImages: Plot all possible perturbations for a fixed pipeline' )
-
-parser.add_argument("--dataset_name", type=str, default='Microscopy', choices=['Microscopy', 'Drone', 'DroneSegmentation'], help='Choose dataset')
-parser.add_argument("--augmentation", type=str, default='weak', choices=['none','weak','strong'], help='Choose augmentation')
+                    'A) Lens2Logit image generation: \n  '
+                    'ABMakeTable: Compute cross-validation metrics results \n  '
+                    'ABShowTable: Plot cross-validation results on a table \n  '
+                    'ABShowImages: Choose a training and testing image to compare different pipelines \n  '
+                    'ABShowAllImages: Plot all possible pipelines \n'
+                    'B) Hendrycks Perturbations, C-type dataset: \n  '
+                    'CMakeTable: For each pipeline, it computes cross-validation metrics for different perturbations  \n  '
+                    'CShowTable: Plot metrics for different pipelines and perturbations \n  '
+                    'CShowImages: Plot an image with a selected a pipeline and perturbation\n  '
+                    'CShowAllImages: Plot all possible perturbations for a fixed pipeline')
+
+parser.add_argument("--dataset_name", type=str, default='Microscopy',
+                    choices=['Microscopy', 'Drone', 'DroneSegmentation'], help='Choose dataset')
+parser.add_argument("--augmentation", type=str, default='weak',
+                    choices=['none', 'weak', 'strong'], help='Choose augmentation')
 parser.add_argument("--N_runs", type=int, default=5, help='Number of k-fold splitting used in the training')
 parser.add_argument("--download_model", default=False, action='store_true', help='Download Models in cache')
 
-#Select pipelines
-parser.add_argument("--dm_train", type=str, default='bilinear', choices= ('bilinear', 'malvar2004', 'menon2007'), help='Choose demosaicing for training processing model')
-parser.add_argument("--s_train",  type=str, default='sharpening_filter', choices= ('sharpening_filter', 'unsharp_masking'), help='Choose sharpening for training processing model')
-parser.add_argument("--dn_train", type=str, default='gaussian_denoising', choices= ('gaussian_denoising', 'median_denoising'), help='Choose denoising for training processing model')
-parser.add_argument("--dm_test", type=str, default='bilinear', choices= ('bilinear', 'malvar2004', 'menon2007'), help='Choose demosaicing for testing processing model')
-parser.add_argument("--s_test", type=str, default='sharpening_filter', choices= ('sharpening_filter', 'unsharp_masking'), help='Choose sharpening for testing processing model')
-parser.add_argument("--dn_test", type=str, default='gaussian_denoising', choices= ('gaussian_denoising', 'median_denoising'), help='Choose denoising for testing processing model')
-
-#Select Ctest parameters
-parser.add_argument("--transform", type=str,  default='identity', choices= ('identity','gaussian_noise', 'shot_noise', 'impulse_noise', 'speckle_noise', 
-                    'gaussian_blur', 'zoom_blur', 'contrast', 'brightness', 'saturate', 'elastic_transform'), help='Choose transformation to show for Ctesting')
-parser.add_argument("--severity", type=int, default=1, choices= (1,2,3,4,5), help='Choose severity for Ctesting')
+# Select pipelines
+parser.add_argument("--dm_train", type=str, default='bilinear', choices=('bilinear', 'malvar2004',
+                                                                         'menon2007'), help='Choose demosaicing for training processing model')
+parser.add_argument("--s_train", type=str, default='sharpening_filter', choices=('sharpening_filter',
+                                                                                 'unsharp_masking'), help='Choose sharpening for training processing model')
+parser.add_argument("--dn_train", type=str, default='gaussian_denoising', choices=('gaussian_denoising',
+                                                                                   'median_denoising'), help='Choose denoising for training processing model')
+parser.add_argument("--dm_test", type=str, default='bilinear', choices=('bilinear', 'malvar2004',
+                                                                        'menon2007'), help='Choose demosaicing for testing processing model')
+parser.add_argument("--s_test", type=str, default='sharpening_filter', choices=('sharpening_filter',
+                                                                                'unsharp_masking'), help='Choose sharpening for testing processing model')
+parser.add_argument("--dn_test", type=str, default='gaussian_denoising', choices=('gaussian_denoising',
+                                                                                  'median_denoising'), help='Choose denoising for testing processing model')
+
+# Select Ctest parameters
+parser.add_argument("--transform", type=str, default='identity', choices=('identity', 'gaussian_noise', 'shot_noise', 'impulse_noise', 'speckle_noise',
+                                                                          'gaussian_blur', 'zoom_blur', 'contrast', 'brightness', 'saturate', 'elastic_transform'), help='Choose transformation to show for Ctesting')
+parser.add_argument("--severity", type=int, default=1, choices=(1, 2, 3, 4, 5), help='Choose severity for Ctesting')
 
 args = parser.parse_args()
 
+
 class metrics:
     def __init__(self, confusion_matrix):
         self.cm = confusion_matrix
         self.N_classes = len(confusion_matrix)
 
     def accuracy(self):
-        Tp = torch.diagonal(self.cm,0).sum()
+        Tp = torch.diagonal(self.cm, 0).sum()
         N_elements = torch.sum(self.cm)
-        return Tp/N_elements
-    
+        return Tp / N_elements
+
     def precision(self):
         Tp_Fp = torch.sum(self.cm, 1)
         Tp_Fp[Tp_Fp == 0] = 1
-        return torch.diagonal(self.cm,0) / Tp_Fp
+        return torch.diagonal(self.cm, 0) / Tp_Fp
 
     def recall(self):
         Tp_Fn = torch.sum(self.cm, 0)
         Tp_Fn[Tp_Fn == 0] = 1
-        return torch.diagonal(self.cm,0) / Tp_Fn
+        return torch.diagonal(self.cm, 0) / Tp_Fn
 
     def f1_score(self):
-        prod = (self.precision()*self.recall())
+        prod = (self.precision() * self.recall())
         sum = (self.precision() + self.recall())
         sum[sum == 0.] = 1.
-        return 2*( prod / sum )
+        return 2 * (prod / sum)
 
     def over_N_runs(ms, N_runs):
-        m, m2 = 0, 0 
+        m, m2 = 0, 0
 
         for i in ms:
-            m += i 
-        mu = m/N_runs
-        
+            m += i
+        mu = m / N_runs
+
         for i in ms:
-            m2 += (i-mu)**2
+            m2 += (i - mu)**2
+
+        sigma = torch.sqrt(m2 / (N_runs - 1))
 
-        sigma = torch.sqrt( m2 / (N_runs-1) ) 
-        
         return mu.tolist(), sigma.tolist()
 
+
 class ABtesting:
-    def __init__(self, 
-                    dataset_name: str, 
-                    augmentation: str, 
-                    dm_train: str, 
-                    s_train: str, 
-                    dn_train: str, 
-                    dm_test: str, 
-                    s_test: str, 
-                    dn_test: str, 
-                    N_runs: int, 
-                    severity=1, 
-                    transform='identity', 
-                    download_model=False):
+    def __init__(self,
+                 dataset_name: str,
+                 augmentation: str,
+                 dm_train: str,
+                 s_train: str,
+                 dn_train: str,
+                 dm_test: str,
+                 s_test: str,
+                 dn_test: str,
+                 N_runs: int,
+                 severity=1,
+                 transform='identity',
+                 download_model=False):
         self.experiment_name = 'ABtesting'
         self.dataset_name = dataset_name
         self.augmentation = augmentation
@@ -129,12 +139,12 @@ class ABtesting:
         if sharpening == None:
             sharpening = self.s_test
         if denoising == None:
-            denoising = self.dn_test   
+            denoising = self.dn_test
         if severity == None:
-            severity = self.severity   
+            severity = self.severity
         if transform == None:
             transform = self.transform
-        
+
         dataset = get_dataset(self.dataset_name)
 
         if self.dataset_name == "Drone" or self.dataset_name == "DroneSegmentation":
@@ -146,40 +156,41 @@ class ABtesting:
 
         if not plot_mode:
             dataset.transform = Compose([RawProcessingPipeline(
-                    camera_parameters=dataset.camera_parameters,
-                    debayer=debayer,
-                    sharpening=sharpening,
-                    denoising=denoising,
-                    ), Distortions(severity=severity, transform=transform), 
-                    Normalize(mean, std)])
+                camera_parameters=dataset.camera_parameters,
+                debayer=debayer,
+                sharpening=sharpening,
+                denoising=denoising,
+            ), Distortions(severity=severity, transform=transform),
+                Normalize(mean, std)])
         else:
-             dataset.transform = Compose([RawProcessingPipeline(
-                    camera_parameters=dataset.camera_parameters,
-                    debayer=debayer,
-                    sharpening=sharpening,
-                    denoising=denoising,
-                    ), Distortions(severity=severity, transform=transform)])
+            dataset.transform = Compose([RawProcessingPipeline(
+                camera_parameters=dataset.camera_parameters,
+                debayer=debayer,
+                sharpening=sharpening,
+                denoising=denoising,
+            ), Distortions(severity=severity, transform=transform)])
 
         return dataset
 
     def ABclassification(self):
-                
+
         DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
         parent_run_name = f"{self.dataset_name}_{self.dm_train}_{self.s_train}_{self.dn_train}_{self.augmentation}"
 
-        print(f'\nTraining pipeline:\n Dataset: {self.dataset_name}, Augmentation: {self.augmentation} \n Debayer: {self.dm_train}, Sharpening: {self.s_train}, Denoiser: {self.dn_train} \n')
-        print(f'\nTesting pipeline:\n Dataset: {self.dataset_name}, Augmentation: {self.augmentation} \n Debayer: {self.dm_test}, Sharpening: {self.s_test}, Denoiser: {self.dn_test} \n Transform: {self.transform}, Severity: {self.severity}\n')     
+        print(
+            f'\nTraining pipeline:\n Dataset: {self.dataset_name}, Augmentation: {self.augmentation} \n Debayer: {self.dm_train}, Sharpening: {self.s_train}, Denoiser: {self.dn_train} \n')
+        print(f'\nTesting pipeline:\n Dataset: {self.dataset_name}, Augmentation: {self.augmentation} \n Debayer: {self.dm_test}, Sharpening: {self.s_test}, Denoiser: {self.dn_test} \n Transform: {self.transform}, Severity: {self.severity}\n')
 
-        accuracies, precisions, recalls, f1_scores = [],[],[],[]
+        accuracies, precisions, recalls, f1_scores = [], [], [], []
 
         os.system('rm -r /tmp/py*')
 
         for N_run in range(self.N_runs):
 
             print(f"Evaluating Run {N_run}")
-                    
-            run_name = parent_run_name+'_'+str(N_run)
+
+            run_name = parent_run_name + '_' + str(N_run)
 
             state_dict, model = get_mlflow_model_by_name(self.experiment_name, run_name,
                                                          download_model=self.download_model)
@@ -190,18 +201,18 @@ class ABtesting:
 
             model.eval()
 
-            len_classes = len(dataset.classes)   
+            len_classes = len(dataset.classes)
             confusion_matrix = torch.zeros((len_classes, len_classes))
 
             for img, label in valid_loader:
-                
+
                 prediction = model(img.to(DEVICE)).detach().cpu()
-                prediction  = torch.argmax(prediction, dim=1)
-                confusion_matrix[label,prediction] += 1 # Real value rows, Declared columns
+                prediction = torch.argmax(prediction, dim=1)
+                confusion_matrix[label, prediction] += 1  # Real value rows, Declared columns
 
             m = metrics(confusion_matrix)
 
-            accuracies.append(m.accuracy()) 
+            accuracies.append(m.accuracy())
             precisions.append(m.precision())
             recalls.append(m.recall())
             f1_scores.append(m.f1_score())
@@ -213,15 +224,16 @@ class ABtesting:
         recall = metrics.over_N_runs(recalls, self.N_runs)
         f1_score = metrics.over_N_runs(f1_scores, self.N_runs)
         return dataset.classes, accuracy, precision, recall, f1_score
-        
+
     def ABsegmentation(self):
-                        
+
         DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
         parent_run_name = f"{self.dataset_name}_{self.dm_train}_{self.s_train}_{self.dn_train}_{self.augmentation}"
 
-        print(f'\nTraining pipeline:\n Dataset: {self.dataset_name}, Augmentation: {self.augmentation} \n Debayer: {self.dm_train}, Sharpening: {self.s_train}, Denoiser: {self.dn_train} \n')
-        print(f'\nTesting pipeline:\n Dataset: {self.dataset_name}, Augmentation: {self.augmentation} \n Debayer: {self.dm_test}, Sharpening: {self.s_test}, Denoiser: {self.dn_test} \n Transform: {self.transform}, Severity: {self.severity}\n')     
+        print(
+            f'\nTraining pipeline:\n Dataset: {self.dataset_name}, Augmentation: {self.augmentation} \n Debayer: {self.dm_train}, Sharpening: {self.s_train}, Denoiser: {self.dn_train} \n')
+        print(f'\nTesting pipeline:\n Dataset: {self.dataset_name}, Augmentation: {self.augmentation} \n Debayer: {self.dm_test}, Sharpening: {self.s_test}, Denoiser: {self.dn_test} \n Transform: {self.transform}, Severity: {self.severity}\n')
 
         IoUs = []
 
@@ -230,34 +242,34 @@ class ABtesting:
         for N_run in range(self.N_runs):
 
             print(f"Evaluating Run {N_run}")
-                    
-            run_name = parent_run_name+'_'+str(N_run)
 
-            state_dict, model = get_mlflow_model_by_name(self.experiment_name, run_name, 
-                                                        download_model=self.download_model)
+            run_name = parent_run_name + '_' + str(N_run)
+
+            state_dict, model = get_mlflow_model_by_name(self.experiment_name, run_name,
+                                                         download_model=self.download_model)
 
             dataset = self.static_pip_val()
-            
+
             valid_set = Subset(dataset, indices=state_dict['valid_indices'])
             valid_loader = DataLoader(valid_set, batch_size=1, num_workers=16, shuffle=False)
 
             model.eval()
 
-            IoU=0
+            IoU = 0
 
             for img, label in valid_loader:
-                
+
                 prediction = model(img.to(DEVICE)).detach().cpu()
                 prediction = F.logsigmoid(prediction).exp().squeeze()
-                IoU += smp.utils.metrics.IoU()(prediction,label)  
+                IoU += smp.utils.metrics.IoU()(prediction, label)
 
-            IoU = IoU/len(valid_loader)
+            IoU = IoU / len(valid_loader)
             IoUs.append(IoU.item())
 
             os.system('rm -r /tmp/t*')
 
         IoU = metrics.over_N_runs(torch.tensor(IoUs), self.N_runs)
-        return IoU           
+        return IoU
 
     def ABShowImages(self):
 
@@ -265,164 +277,169 @@ class ABtesting:
         if not os.path.exists(path):
             os.makedirs(path)
 
-        path = os.path.join(path, f'{self.dataset_name}_{self.augmentation}_{self.dm_train[:2]}{self.s_train[0]}{self.dn_train[:2]}_{self.dm_test[:2]}{self.s_test[0]}{self.dn_test[:2]}')
+        path = os.path.join(
+            path, f'{self.dataset_name}_{self.augmentation}_{self.dm_train[:2]}{self.s_train[0]}{self.dn_train[:2]}_{self.dm_test[:2]}{self.s_test[0]}{self.dn_test[:2]}')
 
         if not os.path.exists(path):
             os.makedirs(path)
 
-        run_name = f"{self.dataset_name}_{self.dm_train}_{self.s_train}_{self.dn_train}_{self.augmentation}"+'_'+str(0)
+        run_name = f"{self.dataset_name}_{self.dm_train}_{self.s_train}_{self.dn_train}_{self.augmentation}" + \
+            '_' + str(0)
 
         state_dict, model = get_mlflow_model_by_name(self.experiment_name, run_name, download_model=self.download_model)
 
         model.augmentation = None
-                
-        for t in ([self.dm_train, self.s_train, self.dn_train, 'train_img'], 
-                    [self.dm_test, self.s_test, self.dn_test, 'test_img']):
-            
+
+        for t in ([self.dm_train, self.s_train, self.dn_train, 'train_img'],
+                  [self.dm_test, self.s_test, self.dn_test, 'test_img']):
+
             debayer, sharpening, denoising, img_type = t[0], t[1], t[2], t[3]
 
             dataset = self.static_pip_val(debayer=debayer, sharpening=sharpening, denoising=denoising, plot_mode=True)
             valid_set = Subset(dataset, indices=state_dict['valid_indices'])
-            
+
             img, _ = next(iter(valid_set))
 
             plt.figure()
-            plt.imshow(img.permute(1,2,0))
+            plt.imshow(img.permute(1, 2, 0))
             if img_type == 'train_img':
                 plt.title('Train Image')
                 plt.savefig(os.path.join(path, f'img_train.png'))
                 imgA = img
             else:
                 plt.title('Test Image')
-                plt.savefig(os.path.join(path,f'img_test.png'))
-            
-                for c, color in enumerate(['Red','Green','Blue']):
-                    diff = torch.abs(imgA-img)
+                plt.savefig(os.path.join(path, f'img_test.png'))
+
+                for c, color in enumerate(['Red', 'Green', 'Blue']):
+                    diff = torch.abs(imgA - img)
                     plt.figure()
                     # plt.imshow(diff.permute(1,2,0))
-                    plt.imshow(diff[c,50:200,50:200], cmap=f'{color}s')
+                    plt.imshow(diff[c, 50:200, 50:200], cmap=f'{color}s')
                     plt.title(f'|Train Image - Test Image| - {color}')
                     plt.colorbar()
                     plt.savefig(os.path.join(path, f'diff_{color}.png'))
                     plt.figure()
-                    diff[diff == 0.]= 1e-5
+                    diff[diff == 0.] = 1e-5
                     # plt.imshow(torch.log(diff.permute(1,2,0)))
                     plt.imshow(torch.log(diff)[c])
                     plt.title(f'log(|Train Image - Test Image|) - color')
                     plt.colorbar()
                     plt.savefig(os.path.join(path, f'logdiff_{color}.png'))
-                
+
             if self.dataset_name == 'DroneSegmentation':
                 plt.figure()
                 plt.imshow(model(img[None].cuda()).detach().cpu().squeeze())
                 if img_type == 'train_img':
                     plt.savefig(os.path.join(path, f'mask_train.png'))
                 else:
-                    plt.savefig(os.path.join(path,f'mask_test.png'))
+                    plt.savefig(os.path.join(path, f'mask_test.png'))
 
     def ABShowAllImages(self):
         if not os.path.exists('results/ABtesting'):
             os.makedirs('results/ABtesting')
 
-        demosaicings=['bilinear','malvar2004', 'menon2007']
-        sharpenings=['sharpening_filter', 'unsharp_masking']
-        denoisings=['median_denoising', 'gaussian_denoising']
+        demosaicings = ['bilinear', 'malvar2004', 'menon2007']
+        sharpenings = ['sharpening_filter', 'unsharp_masking']
+        denoisings = ['median_denoising', 'gaussian_denoising']
 
         fig = plt.figure()
-        columns=4
-        rows=3
+        columns = 4
+        rows = 3
 
-        i=1
+        i = 1
 
         for dm in demosaicings:
             for s in sharpenings:
                 for dn in denoisings:
-                            
-                    dataset = self.static_pip_val(self.dm_test, self.s_test, 
-                                                self.dn_test, plot_mode=True)
 
-                    img,_ = dataset[0]
-                            
+                    dataset = self.static_pip_val(self.dm_test, self.s_test,
+                                                  self.dn_test, plot_mode=True)
+
+                    img, _ = dataset[0]
+
                     fig.add_subplot(rows, columns, i)
-                    plt.imshow(img.permute(1,2,0))
+                    plt.imshow(img.permute(1, 2, 0))
                     plt.title(f'{dm}\n{s}\n{dn}', fontsize=8)
                     plt.xticks([])
                     plt.yticks([])
                     plt.tight_layout()
 
-                    i+=1
+                    i += 1
 
         plt.show()
         plt.savefig(f'results/ABtesting/ABpipelines.png')
 
     def CShowImages(self):
-        
+
         path = 'results/Ctesting/imgs/'
         if not os.path.exists(path):
             os.makedirs(path)
 
-        run_name = f"{self.dataset_name}_{self.dm_test}_{self.s_test}_{self.dn_test}_{self.augmentation}"+'_'+str(0)
+        run_name = f"{self.dataset_name}_{self.dm_test}_{self.s_test}_{self.dn_test}_{self.augmentation}" + '_' + str(0)
 
         state_dict, model = get_mlflow_model_by_name(self.experiment_name, run_name, download_model=True)
 
         model.augmentation = None
-        
-        dataset = self.static_pip_val(self.dm_test, self.s_test, self.dn_test, self.severity, self.transform, plot_mode=True)
+
+        dataset = self.static_pip_val(self.dm_test, self.s_test, self.dn_test,
+                                      self.severity, self.transform, plot_mode=True)
         valid_set = Subset(dataset, indices=state_dict['valid_indices'])
-        
+
         img, _ = next(iter(valid_set))
 
         plt.figure()
-        plt.imshow(img.permute(1,2,0))
-        plt.savefig(os.path.join(path, f'{self.dataset_name}_{self.augmentation}_{self.dm_train[:2]}{self.s_train[0]}{self.dn_train[:2]}_{self.transform}_sev{self.severity}'))
-            
+        plt.imshow(img.permute(1, 2, 0))
+        plt.savefig(os.path.join(
+            path, f'{self.dataset_name}_{self.augmentation}_{self.dm_train[:2]}{self.s_train[0]}{self.dn_train[:2]}_{self.transform}_sev{self.severity}'))
+
     def CShowAllImages(self):
         if not os.path.exists('results/Cimages'):
             os.makedirs('results/Cimages')
 
-        transforms = ['identity','gaussian_noise', 'shot_noise', 'impulse_noise', 'speckle_noise', 
-                        'gaussian_blur', 'zoom_blur', 'contrast', 'brightness', 'saturate', 'elastic_transform']
+        transforms = ['identity', 'gaussian_noise', 'shot_noise', 'impulse_noise', 'speckle_noise',
+                      'gaussian_blur', 'zoom_blur', 'contrast', 'brightness', 'saturate', 'elastic_transform']
 
-        for i,t in enumerate(transforms):
-                    
-            fig = plt.figure(figsize=(10,6))
+        for i, t in enumerate(transforms):
+
+            fig = plt.figure(figsize=(10, 6))
             columns = 5
             rows = 1
 
-            for sev in range(1,6):
+            for sev in range(1, 6):
 
                 dataset = self.static_pip_val(severity=sev, transform=t, plot_mode=True)
 
-                img,_ = dataset[0]
-                        
+                img, _ = dataset[0]
+
                 fig.add_subplot(rows, columns, sev)
-                plt.imshow(img.permute(1,2,0))
+                plt.imshow(img.permute(1, 2, 0))
                 plt.title(f'Severity: {sev}')
                 plt.xticks([])
                 plt.yticks([])
                 plt.tight_layout()
-            
+
             if '_' in t:
-                t=t.replace('_', ' ')
-            t=t[0].upper()+t[1:]
+                t = t.replace('_', ' ')
+            t = t[0].upper() + t[1:]
 
             fig.suptitle(f'{t}', x=0.5, y=0.8, fontsize=24)
             plt.show()
             plt.savefig(f'results/Cimages/{i+1}_{t.lower()}.png')
 
-def ABMakeTable(dataset_name:str, augmentation: str, 
-            N_runs: int, download_model: bool):
 
-    demosaicings=['bilinear','malvar2004', 'menon2007']
-    sharpenings=['sharpening_filter', 'unsharp_masking']
-    denoisings=['median_denoising', 'gaussian_denoising']
+def ABMakeTable(dataset_name: str, augmentation: str,
+                N_runs: int, download_model: bool):
+
+    demosaicings = ['bilinear', 'malvar2004', 'menon2007']
+    sharpenings = ['sharpening_filter', 'unsharp_masking']
+    denoisings = ['median_denoising', 'gaussian_denoising']
 
-    path='results/ABtesting/tables'
+    path = 'results/ABtesting/tables'
     if not os.path.exists(path):
         os.makedirs(path)
 
-    runs={}
-    i=0
+    runs = {}
+    i = 0
 
     for dm_train in demosaicings:
         for s_train in sharpenings:
@@ -431,28 +448,28 @@ def ABMakeTable(dataset_name:str, augmentation: str,
                     for s_test in sharpenings:
                         for dn_test in denoisings:
                             train_pip = [dm_train, s_train, dn_train]
-                            test_pip = [dm_test, s_test, dn_test] 
+                            test_pip = [dm_test, s_test, dn_test]
                             runs[f'run{i}'] = {
-                            'dataset': dataset_name,
-                            'augmentation': augmentation,
-                            'train_pip': train_pip,
-                            'test_pip': test_pip,
-                            'N_runs': N_runs
+                                'dataset': dataset_name,
+                                'augmentation': augmentation,
+                                'train_pip': train_pip,
+                                'test_pip': test_pip,
+                                'N_runs': N_runs
                             }
                             ABclass = ABtesting(
-                                                dataset_name=dataset_name, 
-                                                augmentation=augmentation, 
-                                                dm_train = dm_train,
-                                                s_train = s_train,
-                                                dn_train = dn_train,
-                                                dm_test = dm_test,
-                                                s_test = s_test,
-                                                dn_test = dn_test,
-                                                N_runs=N_runs,
-                                                download_model=download_model
-                                            )
-
-                            if dataset_name == 'DroneSegmentation': 
+                                dataset_name=dataset_name,
+                                augmentation=augmentation,
+                                dm_train=dm_train,
+                                s_train=s_train,
+                                dn_train=dn_train,
+                                dm_test=dm_test,
+                                s_test=s_test,
+                                dn_test=dn_test,
+                                N_runs=N_runs,
+                                download_model=download_model
+                            )
+
+                            if dataset_name == 'DroneSegmentation':
                                 IoU = ABclass.ABsegmentation()
                                 runs[f'run{i}']['IoU'] = IoU
                             else:
@@ -462,15 +479,16 @@ def ABMakeTable(dataset_name:str, augmentation: str,
                                 runs[f'run{i}']['precision'] = precision
                                 runs[f'run{i}']['recall'] = recall
                                 runs[f'run{i}']['f1_score'] = f1_score
-                            
-                            with open(os.path.join(path,f'{dataset_name}_{augmentation}_runs.txt'), 'w') as outfile:
+
+                            with open(os.path.join(path, f'{dataset_name}_{augmentation}_runs.txt'), 'w') as outfile:
                                 json.dump(runs, outfile)
 
-                            i+=1
+                            i += 1
+
 
 def ABShowTable(dataset_name: str, augmentation: str):
-    
-    path='results/ABtesting/tables'
+
+    path = 'results/ABtesting/tables'
     assert os.path.exists(path), 'No tables to plot'
 
     json_file = os.path.join(path, f'{dataset_name}_{augmentation}_runs.txt')
@@ -478,14 +496,14 @@ def ABShowTable(dataset_name: str, augmentation: str):
     with open(json_file, 'r') as run_file:
         runs = json.load(run_file)
 
-        metrics=torch.zeros((2,12,12))
-        classes=[]
+        metrics = torch.zeros((2, 12, 12))
+        classes = []
 
-        i,j=0,0
+        i, j = 0, 0
 
         for r in range(len(runs)):
-                    
-            run = runs['run'+str(r)]
+
+            run = runs['run' + str(r)]
             if dataset_name == 'DroneSegmentation':
                 acc = run['IoU']
             else:
@@ -494,64 +512,66 @@ def ABShowTable(dataset_name: str, augmentation: str):
                 class_list = run['test_pip']
                 class_name = f'{class_list[0][:2]},{class_list[1][:1]},{class_list[2][:2]}'
                 classes.append(class_name)
-            mu,sigma = round(acc[0],4),round(acc[1],4)
+            mu, sigma = round(acc[0], 4), round(acc[1], 4)
+
+            metrics[0, j, i] = mu
+            metrics[1, j, i] = sigma
 
-            metrics[0,j,i] = mu
-            metrics[1,j,i] = sigma
-           
-            i+=1
+            i += 1
 
             if i == 12:
-                i=0
-                j+=1
+                i = 0
+                j += 1
 
     differences = torch.zeros_like(metrics)
 
-    diag_mu = torch.diagonal(metrics[0],0)
-    diag_sigma = torch.diagonal(metrics[1],0)
-    
+    diag_mu = torch.diagonal(metrics[0], 0)
+    diag_sigma = torch.diagonal(metrics[1], 0)
+
     for r in range(len(metrics[0])):
-        differences[0,r] = diag_mu[r] - metrics[0,r]
-        differences[1,r] = torch.sqrt(metrics[1,r]**2 + diag_sigma[r]**2)
+        differences[0, r] = diag_mu[r] - metrics[0, r]
+        differences[1, r] = torch.sqrt(metrics[1, r]**2 + diag_sigma[r]**2)
 
     # Plot with scatter
-        
-    for i,img in enumerate([metrics, differences]):
+
+    for i, img in enumerate([metrics, differences]):
 
         x, y = torch.arange(12), torch.arange(12)
         x, y = torch.meshgrid(x, y)
 
         if i == 0:
-            vmin = max(0.65, round(img[0].min().item(),2))
-            vmax = round(img[0].max().item(),2)
+            vmin = max(0.65, round(img[0].min().item(), 2))
+            vmax = round(img[0].max().item(), 2)
             step = 0.02
         elif i == 1:
-            vmin = round(img[0].min().item(),2)
+            vmin = round(img[0].min().item(), 2)
             if augmentation == 'none':
-                vmax = min(0.15, round(img[0].max().item(),2))
+                vmax = min(0.15, round(img[0].max().item(), 2))
             if augmentation == 'weak':
-                vmax = min(0.08, round(img[0].max().item(),2))
+                vmax = min(0.08, round(img[0].max().item(), 2))
             if augmentation == 'strong':
-                vmax = min(0.05, round(img[0].max().item(),2))
+                vmax = min(0.05, round(img[0].max().item(), 2))
             step = 0.01
-        
-        vmin = int(vmin/step)*step
-        vmax = int(vmax/step)*step
 
-        fig = plt.figure(figsize=(10,6.2))
+        vmin = int(vmin / step) * step
+        vmax = int(vmax / step) * step
+
+        fig = plt.figure(figsize=(10, 6.2))
         ax = fig.add_axes([0.1, 0.1, 0.8, 0.8])
-        marker_size=350  
-        plt.scatter(x, y, c=torch.rot90(img[1][x,y],-1,[0,1]), vmin = 0., vmax = img[1].max(), cmap='viridis', s=marker_size*2, marker='s')
-        ticks = torch.arange(0.,img[1].max(),0.03).tolist()
-        ticks = [round(tick,2) for tick in ticks]
+        marker_size = 350
+        plt.scatter(x, y, c=torch.rot90(img[1][x, y], -1, [0, 1]), vmin=0.,
+                    vmax=img[1].max(), cmap='viridis', s=marker_size * 2, marker='s')
+        ticks = torch.arange(0., img[1].max(), 0.03).tolist()
+        ticks = [round(tick, 2) for tick in ticks]
         cba = plt.colorbar(pad=0.06)
         cba.set_ticks(ticks)
         cba.ax.set_yticklabels(ticks)
         # cmap = plt.cm.get_cmap('tab20c').reversed()
         cmap = plt.cm.get_cmap('Reds')
-        plt.scatter(x,y, c=torch.rot90(img[0][x,y],-1,[0,1]), vmin = vmin, vmax = vmax, cmap=cmap, s=marker_size, marker='s')
+        plt.scatter(x, y, c=torch.rot90(img[0][x, y], -1, [0, 1]), vmin=vmin,
+                    vmax=vmax, cmap=cmap, s=marker_size, marker='s')
         ticks = torch.arange(vmin, vmax, step).tolist()
-        ticks = [round(tick,2) for tick in ticks]
+        ticks = [round(tick, 2) for tick in ticks]
         if ticks[-1] != vmax:
             ticks.append(vmax)
         cbb = plt.colorbar(pad=0.06)
@@ -563,26 +583,26 @@ def ABShowTable(dataset_name: str, augmentation: str):
         cbb.ax.set_yticklabels(ticks)
         for x in range(12):
             for y in range(12):
-                txt = round(torch.rot90(img[0],-1,[0,1])[x,y].item(),2)
+                txt = round(torch.rot90(img[0], -1, [0, 1])[x, y].item(), 2)
                 if str(txt) == '-0.0':
                     txt = '0.00'
                 elif str(txt) == '0.0':
                     txt = '0.00'
                 elif len(str(txt)) == 3:
-                    txt = str(txt)+'0'
+                    txt = str(txt) + '0'
                 else:
                     txt = str(txt)
-                    
-                plt.text(x-0.25,y-0.1,txt, color='black', fontsize='x-small')
 
-        ax.set_xticks(torch.linspace(0,11,12))
+                plt.text(x - 0.25, y - 0.1, txt, color='black', fontsize='x-small')
+
+        ax.set_xticks(torch.linspace(0, 11, 12))
         ax.set_xticklabels(classes)
-        ax.set_yticks(torch.linspace(0,11,12))
+        ax.set_yticks(torch.linspace(0, 11, 12))
         classes.reverse()
         ax.set_yticklabels(classes)
         classes.reverse()
-        plt.xticks(rotation = 45)
-        plt.yticks(rotation = 45)
+        plt.xticks(rotation=45)
+        plt.yticks(rotation=45)
         cba.set_label('Standard Deviation')
         plt.xlabel("Test pipelines")
         plt.ylabel("Train pipelines")
@@ -590,62 +610,63 @@ def ABShowTable(dataset_name: str, augmentation: str):
         if i == 0:
             if dataset_name == 'DroneSegmentation':
                 cbb.set_label('IoU')
-                plt.savefig(os.path.join(path,f"{dataset_name}_{augmentation}_IoU.png"))
+                plt.savefig(os.path.join(path, f"{dataset_name}_{augmentation}_IoU.png"))
             else:
                 cbb.set_label('Accuracy')
-                plt.savefig(os.path.join(path,f"{dataset_name}_{augmentation}_accuracies.png"))
+                plt.savefig(os.path.join(path, f"{dataset_name}_{augmentation}_accuracies.png"))
         elif i == 1:
             if dataset_name == 'DroneSegmentation':
                 cbb.set_label('IoU_d-IoU')
             else:
                 cbb.set_label('Accuracy_d - Accuracy')
-            plt.savefig(os.path.join(path,f"{dataset_name}_{augmentation}_differences.png"))
+            plt.savefig(os.path.join(path, f"{dataset_name}_{augmentation}_differences.png"))
+
 
 def CMakeTable(dataset_name: str, augmentation: str, severity: int, N_runs: int, download_model: bool):
-    
-    path='results/Ctesting/tables'
+
+    path = 'results/Ctesting/tables'
     if not os.path.exists(path):
         os.makedirs(path)
-    
-    demosaicings=['bilinear','malvar2004', 'menon2007']
-    sharpenings=['sharpening_filter', 'unsharp_masking']
-    denoisings=['median_denoising', 'gaussian_denoising']
 
-    transformations = ['identity','gaussian_noise', 'shot_noise', 'impulse_noise', 'speckle_noise', 
-                    'gaussian_blur', 'zoom_blur', 'contrast', 'brightness', 'saturate', 'elastic_transform']
+    demosaicings = ['bilinear', 'malvar2004', 'menon2007']
+    sharpenings = ['sharpening_filter', 'unsharp_masking']
+    denoisings = ['median_denoising', 'gaussian_denoising']
+
+    transformations = ['identity', 'gaussian_noise', 'shot_noise', 'impulse_noise', 'speckle_noise',
+                       'gaussian_blur', 'zoom_blur', 'contrast', 'brightness', 'saturate', 'elastic_transform']
 
-    runs={}
-    i=0
+    runs = {}
+    i = 0
 
     for dm in demosaicings:
         for s in sharpenings:
             for dn in denoisings:
                 for t in transformations:
-                    pip = [dm,s,dn]
+                    pip = [dm, s, dn]
                     runs[f'run{i}'] = {
-                    'dataset': dataset_name,
-                    'augmentation': augmentation,
-                    'pipeline': pip,
-                    'N_runs': N_runs,
-                    'transform': t,
-                    'severity': severity,
+                        'dataset': dataset_name,
+                        'augmentation': augmentation,
+                        'pipeline': pip,
+                        'N_runs': N_runs,
+                        'transform': t,
+                        'severity': severity,
                     }
                     ABclass = ABtesting(
-                                        dataset_name=dataset_name, 
-                                        augmentation=augmentation, 
-                                        dm_train = dm,
-                                        s_train = s,
-                                        dn_train = dn,
-                                        dm_test = dm,
-                                        s_test = s,
-                                        dn_test = dn,
-                                        severity=severity,
-                                        transform=t,
-                                        N_runs=N_runs,
-                                        download_model=download_model
-                                    )
-
-                    if dataset_name == 'DroneSegmentation': 
+                        dataset_name=dataset_name,
+                        augmentation=augmentation,
+                        dm_train=dm,
+                        s_train=s,
+                        dn_train=dn,
+                        dm_test=dm,
+                        s_test=s,
+                        dn_test=dn,
+                        severity=severity,
+                        transform=t,
+                        N_runs=N_runs,
+                        download_model=download_model
+                    )
+
+                    if dataset_name == 'DroneSegmentation':
                         IoU = ABclass.ABsegmentation()
                         runs[f'run{i}']['IoU'] = IoU
                     else:
@@ -656,26 +677,27 @@ def CMakeTable(dataset_name: str, augmentation: str, severity: int, N_runs: int,
                         runs[f'run{i}']['recall'] = recall
                         runs[f'run{i}']['f1_score'] = f1_score
 
-                    with open(os.path.join(path,f'{dataset_name}_{augmentation}_runs.json'), 'w') as outfile:
+                    with open(os.path.join(path, f'{dataset_name}_{augmentation}_runs.json'), 'w') as outfile:
                         json.dump(runs, outfile)
 
-                    i+=1
+                    i += 1
+
 
 def CShowTable(dataset_name, augmentation):
 
-    path='results/Ctesting/tables'
+    path = 'results/Ctesting/tables'
     assert os.path.exists(path), 'No tables to plot'
 
     json_file = os.path.join(path, f'{dataset_name}_{augmentation}_runs.txt')
 
-    transforms = ['identity','gauss_noise', 'shot', 'impulse', 'speckle', 
-                    'gauss_blur', 'zoom', 'contrast', 'brightness', 'saturate', 'elastic']
+    transforms = ['identity', 'gauss_noise', 'shot', 'impulse', 'speckle',
+                  'gauss_blur', 'zoom', 'contrast', 'brightness', 'saturate', 'elastic']
 
     pip = []
-    
-    demosaicings=['bilinear','malvar2004', 'menon2007']
-    sharpenings=['sharpening_filter', 'unsharp_masking']
-    denoisings=['median_denoising', 'gaussian_denoising']
+
+    demosaicings = ['bilinear', 'malvar2004', 'menon2007']
+    sharpenings = ['sharpening_filter', 'unsharp_masking']
+    denoisings = ['median_denoising', 'gaussian_denoising']
 
     for dm in demosaicings:
         for s in sharpenings:
@@ -685,52 +707,54 @@ def CShowTable(dataset_name, augmentation):
     with open(json_file, 'r') as run_file:
         runs = json.load(run_file)
 
-        metrics=torch.zeros((2,len(pip),len(transforms)))
+        metrics = torch.zeros((2, len(pip), len(transforms)))
 
-        i,j=0,0
+        i, j = 0, 0
 
         for r in range(len(runs)):
-                    
-            run = runs['run'+str(r)]
+
+            run = runs['run' + str(r)]
             if dataset_name == 'DroneSegmentation':
                 acc = run['IoU']
             else:
                 acc = run['accuracy']
-            mu,sigma = round(acc[0],4),round(acc[1],4)
+            mu, sigma = round(acc[0], 4), round(acc[1], 4)
 
-            metrics[0,j,i] = mu
-            metrics[1,j,i] = sigma
-           
-            i+=1
+            metrics[0, j, i] = mu
+            metrics[1, j, i] = sigma
+
+            i += 1
 
             if i == len(transforms):
-                i=0
-                j+=1
+                i = 0
+                j += 1
 
         # Plot with scatter
 
         img = metrics
 
-        vmin=0.
-        vmax=1.
-        
+        vmin = 0.
+        vmax = 1.
+
         x, y = torch.arange(12), torch.arange(11)
         x, y = torch.meshgrid(x, y)
 
-        fig = plt.figure(figsize=(10,6.2))
+        fig = plt.figure(figsize=(10, 6.2))
         ax = fig.add_axes([0.1, 0.1, 0.8, 0.8])
-        marker_size=350  
-        plt.scatter(x, y, c=torch.rot90(img[1][x,y],-1,[0,1]), vmin = 0., vmax = img[1].max(), cmap='viridis', s=marker_size*2, marker='s')
-        ticks = torch.arange(0.,img[1].max(),0.03).tolist()
-        ticks = [round(tick,2) for tick in ticks]
+        marker_size = 350
+        plt.scatter(x, y, c=torch.rot90(img[1][x, y], -1, [0, 1]), vmin=0.,
+                    vmax=img[1].max(), cmap='viridis', s=marker_size * 2, marker='s')
+        ticks = torch.arange(0., img[1].max(), 0.03).tolist()
+        ticks = [round(tick, 2) for tick in ticks]
         cba = plt.colorbar(pad=0.06)
         cba.set_ticks(ticks)
         cba.ax.set_yticklabels(ticks)
         # cmap = plt.cm.get_cmap('tab20c').reversed()
         cmap = plt.cm.get_cmap('Reds')
-        plt.scatter(x,y, c=torch.rot90(img[0][x,y],-1,[0,1]), vmin=vmin, vmax=vmax, cmap=cmap, s=marker_size, marker='s')
+        plt.scatter(x, y, c=torch.rot90(img[0][x, y], -1, [0, 1]), vmin=vmin,
+                    vmax=vmax, cmap=cmap, s=marker_size, marker='s')
         ticks = torch.arange(vmin, vmax, step).tolist()
-        ticks = [round(tick,2) for tick in ticks]
+        ticks = [round(tick, 2) for tick in ticks]
         if ticks[-1] != vmax:
             ticks.append(vmax)
         cbb = plt.colorbar(pad=0.06)
@@ -742,65 +766,66 @@ def CShowTable(dataset_name, augmentation):
         cbb.ax.set_yticklabels(ticks)
         for x in range(12):
             for y in range(12):
-                txt = round(torch.rot90(img[0],-1,[0,1])[x,y].item(),2)
+                txt = round(torch.rot90(img[0], -1, [0, 1])[x, y].item(), 2)
                 if str(txt) == '-0.0':
                     txt = '0.00'
                 elif str(txt) == '0.0':
                     txt = '0.00'
                 elif len(str(txt)) == 3:
-                    txt = str(txt)+'0'
+                    txt = str(txt) + '0'
                 else:
                     txt = str(txt)
-                    
-                plt.text(x-0.25,y-0.1,txt, color='black', fontsize='x-small')
 
-        ax.set_xticks(torch.linspace(0,11,12))
+                plt.text(x - 0.25, y - 0.1, txt, color='black', fontsize='x-small')
+
+        ax.set_xticks(torch.linspace(0, 11, 12))
         ax.set_xticklabels(transforms)
-        ax.set_yticks(torch.linspace(0,11,12))
+        ax.set_yticks(torch.linspace(0, 11, 12))
         pip.reverse()
         ax.set_yticklabels(pip)
         pip.reverse()
-        plt.xticks(rotation = 45)
-        plt.yticks(rotation = 45)
+        plt.xticks(rotation=45)
+        plt.yticks(rotation=45)
         cba.set_label('Standard Deviation')
         plt.xlabel("Pipelines")
         plt.ylabel("Distortions")
         if dataset_name == 'DroneSegmentation':
             cbb.set_label('IoU')
-            plt.savefig(os.path.join(path,f"{dataset_name}_{augmentation}_IoU.png"))
+            plt.savefig(os.path.join(path, f"{dataset_name}_{augmentation}_IoU.png"))
         else:
             cbb.set_label('Accuracy')
-            plt.savefig(os.path.join(path,f"{dataset_name}_{augmentation}_accuracies.png"))
+            plt.savefig(os.path.join(path, f"{dataset_name}_{augmentation}_accuracies.png"))
+
 
 if __name__ == '__main__':
-     
-    if args.mode == 'ABMakeTable': 
+
+    if args.mode == 'ABMakeTable':
         ABMakeTable(args.dataset_name, args.augmentation, args.N_runs, args.download_model)
-    elif args.mode == 'ABShowTable': 
+    elif args.mode == 'ABShowTable':
         ABShowTable(args.dataset_name, args.augmentation)
-    elif args.mode == 'ABShowImages': 
-        ABclass = ABtesting(args.dataset_name, args.augmentation, args.dm_train, 
-                            args.s_train, args.dn_train, args.dm_test, args.s_test, 
+    elif args.mode == 'ABShowImages':
+        ABclass = ABtesting(args.dataset_name, args.augmentation, args.dm_train,
+                            args.s_train, args.dn_train, args.dm_test, args.s_test,
                             args.dn_test, args.N_runs, download_model=args.download_model)
         ABclass.ABShowImages()
-    elif args.mode == 'ABShowAllImages': 
-        ABclass = ABtesting(args.dataset_name, args.augmentation, args.dm_train, 
-                        args.s_train, args.dn_train, args.dm_test, args.s_test, 
-                        args.dn_test, args.N_runs, download_model=args.download_model)
+    elif args.mode == 'ABShowAllImages':
+        ABclass = ABtesting(args.dataset_name, args.augmentation, args.dm_train,
+                            args.s_train, args.dn_train, args.dm_test, args.s_test,
+                            args.dn_test, args.N_runs, download_model=args.download_model)
         ABclass.ABShowAllImages()
-    elif args.mode == 'CMakeTable': 
+    elif args.mode == 'CMakeTable':
         CMakeTable(args.dataset_name, args.augmentation, args.severity, args.N_runs, args.download_model)
-    elif args.mode == 'CShowTable': # TODO test it
+    elif args.mode == 'CShowTable':  # TODO test it
         CShowTable(args.dataset_name, args.augmentation, args.severity)
-    elif args.mode == 'CShowImages': 
-        ABclass = ABtesting(args.dataset_name, args.augmentation, args.dm_train, 
-                    args.s_train, args.dn_train, args.dm_test, args.s_test, 
-                    args.dn_test, args.N_runs, args.severity, args.transform, 
-                    download_model=args.download_model)
+    elif args.mode == 'CShowImages':
+        ABclass = ABtesting(args.dataset_name, args.augmentation, args.dm_train,
+                            args.s_train, args.dn_train, args.dm_test, args.s_test,
+                            args.dn_test, args.N_runs, args.severity, args.transform,
+                            download_model=args.download_model)
         ABclass.CShowImages()
-    elif args.mode == 'CShowAllImages': 
-        ABclass = ABtesting(args.dataset_name, args.augmentation, args.dm_train, 
-                    args.s_train, args.dn_train, args.dm_test, args.s_test, 
-                    args.dn_test, args.N_runs, args.severity, args.transform, 
-                    download_model=args.download_model)
+    elif args.mode == 'CShowAllImages':
+        ABclass = ABtesting(args.dataset_name, args.augmentation, args.dm_train,
+                            args.s_train, args.dn_train, args.dm_test, args.s_test,
+                            args.dn_test, args.N_runs, args.severity, args.transform,
+                            download_model=args.download_model)
         ABclass.CShowAllImages()

{processingpipeline → figures}/numpy_static_pipeline_show.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:64edef77495ab24143430e7a5d880b6f211568371f37eab03e1b32fb2f5b8015
-size 1906586
+oid sha256:9fd7f0e985d60a9d24fe74ef0faf0e8e10258416190be64d7b06908306f0e7fc
+size 1906578

sanity_checks_and_statistics.ipynb → figures/sanity_checks_and_statistics.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:53f62c6ce9a6656a31c3e0ae1deded2e4f9818cd891381dbe1030dd5edc5f278
-size 6103871
+oid sha256:446fa03594d28d8a23aaa4d721db22e0453724edd4199970418a631e2f4e7b18
+size 6103863

{utils → figures}/show_dataset.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cd09ffb969b9a0a5b414b892614b5b9e48fa32721ea9a14e9e0951160e8f92e4
-size 2115545
+oid sha256:11568132f27aee6b772327127e7dac8b2bd81bc219a2b091e74dd8d6514e419d
+size 2115521

models/classifier.py → model.py

@@ -29,6 +29,7 @@ class LitModel(pl.LightningModule):
                  weight_decay=0,
                  loss_aux=None,
                  adv_training=False,
+                 adv_parameters='all',
                  metrics=None,
                  processor=None,
                  augmentation=None,
@@ -57,11 +58,19 @@ class LitModel(pl.LightningModule):
         self.freeze_classifier = freeze_classifier
         self.freeze_processor = freeze_processor
 
+        self.unfreeze()
         if freeze_classifier:
             pl.LightningModule.freeze(self.classifier)
         if freeze_processor:
             pl.LightningModule.freeze(self.processor)
 
+        if adv_training and adv_parameters != 'all':
+            if adv_parameters != 'all':
+                pl.LightningModule.freeze(self.processor)
+                for name, p in self.processor.named_parameters():
+                    if adv_parameters in name:
+                        p.requires_grad = True
+
     def forward(self, x):
         x = self.processor(x)
         apply_augmentation_step = self.training or self.augmentation_on_eval
@@ -97,7 +106,6 @@ class LitModel(pl.LightningModule):
             y_hat = F.logsigmoid(logits).exp().squeeze()
         else:
             y_hat = torch.argmax(logits, dim=1)
-            
 
         if self.metrics is not None:
             for metric in self.metrics:
@@ -105,11 +113,15 @@ class LitModel(pl.LightningModule):
                 if metric_name == 'accuracy' or not self.training or self.metrics_on_training:
                     m = metric(y_hat.cpu().detach(), y.cpu())
                     self.log(f'{step_name}_{metric_name}', m, on_step=False, on_epoch=True,
-                             prog_bar=self.training or metric_name == 'accuracy')                
+                             prog_bar=self.training or metric_name == 'accuracy')
                 if metric_name == 'iou_score' or not self.training or self.metrics_on_training:
                     m = metric(y_hat.cpu().detach(), y.cpu())
                     self.log(f'{step_name}_{metric_name}', m, on_step=False, on_epoch=True,
                              prog_bar=self.training or metric_name == 'iou_score')
+                elif metric_name == 'accuracy' or not self.training or self.metrics_on_training:
+                    m = metric(y_hat.cpu().detach(), y.cpu())
+                    self.log(f'{step_name}_{metric_name}', m, on_step=False, on_epoch=True,
+                             prog_bar=self.training or metric_name == 'accuracy')
 
         return loss
 
@@ -124,25 +136,15 @@ class LitModel(pl.LightningModule):
 
     def train(self, mode=True):
         self.training = mode
-        # self.processor.train(False)
-        self.processor.train(mode=mode and not self.freeze_processor)
+
+        # don't update batchnorm in adversarial training
+        self.processor.train(mode=mode and not self.freeze_processor and not self.adv_training)
         self.classifier.train(mode=mode and not self.freeze_classifier)
-        if self.adv_training and self.processor.batch_norm is not None:  # don't update batchnorm in adversarial training
-            self.processor.batch_norm.track_running_stats = False
         return self
 
     def configure_optimizers(self):
         self.optimizer = torch.optim.Adam(self.parameters(), self.lr, weight_decay=self.weight_decay)
-        # parameters = [self.processor.additive_layer]
-        # self.optimizer = torch.optim.Adam(parameters, self.lr, weight_decay=self.weight_decay)
         return self.optimizer
-        # self.scheduler = {
-        #     'scheduler': torch.optim.lr_scheduler.ReduceLROnPlateau(
-        #         self.optimizer, mode='min', factor=0.2, patience=2, min_lr=1e-6, verbose=True,
-        #     ),
-        #     'monitor': 'val_loss',
-        # }
-        # return [self.optimizer], [self.scheduler]
 
     def get_progress_bar_dict(self):
         items = super().get_progress_bar_dict()
@@ -151,7 +153,7 @@ class LitModel(pl.LightningModule):
 
 
 class TrackImagesCallback(pl.callbacks.base.Callback):
-    def __init__(self, data_loader, track_every_epoch=False, track_processing=True, track_gradients=True, track_predictions=True, save_tensors=True):
+    def __init__(self, data_loader, reference_processor=None, track_every_epoch=False, track_processing=True, track_gradients=True, track_predictions=True, save_tensors=True):
         super().__init__()
         self.data_loader = data_loader
 
@@ -162,9 +164,12 @@ class TrackImagesCallback(pl.callbacks.base.Callback):
         self.track_predictions = track_predictions
         self.save_tensors = save_tensors
 
-    def callback_track_images(self, trainer, save_loc):
-        track_images(trainer.model,
+        self.reference_processor = reference_processor
+
+    def callback_track_images(self, model, save_loc):
+        track_images(model,
                      self.data_loader,
+                     reference_processor=self.reference_processor,
                      track_processing=self.track_processing,
                      track_gradients=self.track_gradients,
                      track_predictions=self.track_predictions,
@@ -175,12 +180,12 @@ class TrackImagesCallback(pl.callbacks.base.Callback):
     def on_fit_end(self, trainer, pl_module):
         if not self.track_every_epoch:
             save_loc = 'results'
-            self.callback_track_images(trainer, save_loc)
+            self.callback_track_images(trainer.model, save_loc)
 
     def on_train_epoch_end(self, trainer, pl_module, outputs):
         if self.track_every_epoch:
             save_loc = f'results/epoch_{trainer.current_epoch + 1:04d}'
-            self.callback_track_images(trainer, save_loc)
+            self.callback_track_images(trainer.model, save_loc)
 
 
 from utils.debug import debug
@@ -201,7 +206,7 @@ def log_tensor(batch, path, save_tensors=True, nrow=8):
     mlflow.log_artifact(img_path, os.path.dirname(path))
 
 
-def track_images(model, data_loader, track_processing=True, track_gradients=True, track_predictions=True, save_tensors=True, save_loc='results'):
+def track_images(model, data_loader, reference_processor=None, track_processing=True, track_gradients=True, track_predictions=True, save_tensors=True, save_loc='results'):
 
     device = model.device
     processor = model.processor
@@ -219,6 +224,9 @@ def track_images(model, data_loader, track_processing=True, track_gradients=True
     logits_full = []
     stages_full = defaultdict(list)
     grads_full = defaultdict(list)
+    diffs_full = defaultdict(list)
+
+    track_differences = reference_processor is not None
 
     for inputs, labels in data_loader:
 
@@ -227,6 +235,10 @@ def track_images(model, data_loader, track_processing=True, track_gradients=True
 
         processed_rgb = processor(inputs)
 
+        if track_differences:
+            # debug(processor)
+            processed_rgb_ref = reference_processor(inputs)
+
         if track_gradients or track_predictions:
             logits = classifier(processed_rgb)
 
@@ -241,6 +253,8 @@ def track_images(model, data_loader, track_processing=True, track_gradients=True
 
         for stage, batch in processor.stages.items():
             stages_full[stage].append(batch.cpu().detach())
+            if track_differences:
+                diffs_full[stage].append((reference_processor.stages[stage] - batch).cpu().detach())
             if track_gradients:
                 grads_full[stage].append(batch.grad.cpu().detach())
 
@@ -248,19 +262,29 @@ def track_images(model, data_loader, track_processing=True, track_gradients=True
 
         stages = stages_full
         grads = grads_full
+        diffs = diffs_full
 
         if track_processing:
-            for stage, batch in stages_full.items():
+            for stage, batch in stages.items():
                 stages[stage] = torch.cat(batch)
 
+        if track_differences:
+            for stage, batch in diffs.items():
+                diffs[stage] = torch.cat(batch)
+
         if track_gradients:
-            for stage, batch in grads_full.items():
+            for stage, batch in grads.items():
                 grads[stage] = torch.cat(batch)
 
         for stage_nr, stage_name in enumerate(stages):
             if track_processing:
                 batch = stages[stage_name]
                 log_tensor(batch, os.path.join(save_loc, f'processing_{stage_nr}_{stage_name}.pt'), save_tensors)
+
+            if track_differences:
+                batch = diffs[stage_name]
+                log_tensor(batch, os.path.join(save_loc, f'diffs_{stage_nr}_{stage_name}.pt'), False)
+
             if track_gradients:
                 batch_grad = grads[stage_name]
                 batch_grad = batch_grad.abs()
@@ -270,11 +294,11 @@ def track_images(model, data_loader, track_processing=True, track_gradients=True
 
         # inputs = torch.cat(inputs_full)
 
-        if track_predictions: #and model.is_segmentation_task:
+        if track_predictions:  # and model.is_segmentation_task:
             labels = torch.cat(labels_full)
             logits = torch.cat(logits_full)
             masks = labels.unsqueeze(1)
-            predictions = logits #torch.sigmoid(logits).unsqueeze(1)
+            predictions = logits  # torch.sigmoid(logits).unsqueeze(1)
             #mask_vis = torch.cat((masks, predictions, masks * predictions), dim=1)
             #log_tensor(mask_vis, os.path.join(save_loc, f'masks.pt'), save_tensors)
             log_tensor(masks, os.path.join(save_loc, f'targets.pt'), save_tensors)

processingpipeline/pipeline.py → processing/pipeline_numpy.py

@@ -23,7 +23,7 @@ from colour_demosaicing import (demosaicing_CFA_Bayer_bilinear,
 import torch
 import numpy as np
 
-from utils.dataset import Subset
+from dataset import Subset
 from torch.utils.data import DataLoader
 
 from colour_demosaicing import (demosaicing_CFA_Bayer_bilinear,

processingpipeline/torch_pipeline.py → processing/pipeline_torch.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 if not os.path.exists('README.md'):
     os.chdir('..')
 
-from processingpipeline.pipeline import processing as default_processing
+from processing.pipeline_numpy import processing as default_processing
 from utils.base import np2torch, torch2np
 
 import segmentation_models_pytorch as smp
@@ -83,7 +83,7 @@ class NNProcessing(nn.Module):
             in_channels=3,
             classes=3,
         )
-        self.batch_norm = None if not batch_norm_output else nn.BatchNorm2d(3)
+        self.batch_norm = None if not batch_norm_output else nn.BatchNorm2d(3, affine=False)
         self.normalize_mosaic = normalize_mosaic
 
     def forward(self, raw):
@@ -108,20 +108,23 @@ class NNProcessing(nn.Module):
         return rgb
 
 
+def add_additive_layer(processor):
+    processor.additive_layer = nn.Parameter(torch.zeros((1, 3, 256, 256)))
+    # processor.additive_layer = nn.Parameter(0.001 * torch.randn((1, 3, 256, 256)))
+
+
 class ParametrizedProcessing(nn.Module):
-    def __init__(self, camera_parameters, track_stages=False, batch_norm_output=True, noise_layer=False):
+    def __init__(self, camera_parameters, track_stages=False, batch_norm_output=True):
         super().__init__()
         self.stages = None
         self.buffer = None
         self.track_stages = track_stages
 
         black_level, white_balance, colour_matrix = camera_parameters
-        self.register_buffer('black_level', torch.as_tensor(black_level))
-        self.register_buffer('colour_correction',
-                             torch.as_tensor(white_balance).reshape(1, 3)
-                             * torch.as_tensor(colour_matrix).reshape(3, 3))
-        self.register_buffer('M_RGB_2_YUV', M_RGB_2_YUV.clone())
-        self.register_buffer('M_YUV_2_RGB', M_YUV_2_RGB.clone())
+
+        self.black_level = nn.Parameter(torch.as_tensor(black_level))
+        self.white_balance = nn.Parameter(torch.as_tensor(white_balance).reshape(1, 3))
+        self.colour_correction = nn.Parameter(torch.as_tensor(colour_matrix).reshape(3, 3))
 
         self.gamma_correct = nn.Parameter(torch.Tensor([2.2]))
 
@@ -133,14 +136,12 @@ class ParametrizedProcessing(nn.Module):
         self.gaussian_blur = nn.Conv2d(1, 1, kernel_size=5, padding=2, padding_mode='reflect', bias=False)
         self.gaussian_blur.weight.data[0][0] = K_BLUR.clone()
 
-        self.batch_norm = nn.BatchNorm2d(3) if batch_norm_output else None
+        self.batch_norm = nn.BatchNorm2d(3, affine=False) if batch_norm_output else None
 
-        # if noise_layer:
-        #     for param in self.parameters():
-        #         param.requires_grad = False
+        self.register_buffer('M_RGB_2_YUV', M_RGB_2_YUV.clone())
+        self.register_buffer('M_YUV_2_RGB', M_YUV_2_RGB.clone())
 
-        self.additive_layer = nn.Parameter(0.001 * torch.randn((1, 3, 256, 256))
-                                           ) if noise_layer else None  # XXX: can this be 0?
+        self.additive_layer = None  # this can be added in later
 
     def forward(self, raw):
         assert raw.ndim == 3, f"needs dims (B, H, W), got {raw.shape}"
@@ -157,6 +158,7 @@ class ParametrizedProcessing(nn.Module):
         rgb = self.debayer(rgb)
         # self.stages['debayer'] = rgb
 
+        rgb = torch.einsum('bchw,kc->bchw', rgb, self.white_balance).contiguous()
         rgb = torch.einsum('bchw,kc->bkhw', rgb, self.colour_correction).contiguous()
         self.stages['color_correct'] = rgb
 
@@ -179,7 +181,6 @@ class ParametrizedProcessing(nn.Module):
         self.stages['gamma_correct'] = rgb
 
         if self.additive_layer is not None:
-            # rgb = rgb + 0 * self.additive_layer
             rgb = rgb + self.additive_layer
             self.stages['noise'] = rgb
 
@@ -259,11 +260,11 @@ if __name__ == "__main__":
         os.chdir('..')
 
     import matplotlib.pyplot as plt
-    from utils.dataset import get_dataset
+    from dataset import get_dataset
     from utils.base import np2torch, torch2np
 
     from utils.debug import debug
-    from processingpipeline.pipeline import processing as default_processing
+    from processing.pipeline_numpy import processing as default_processing
 
     raw_dataset = get_dataset('DS')
     loader = torch.utils.data.DataLoader(raw_dataset, batch_size=1)

readme/init.md

@@ -1 +0,0 @@
-

train.py

@@ -4,6 +4,7 @@ import copy
 import argparse
 
 import torch
+from torch import optim
 import torch.nn as nn
 
 import mlflow.pytorch
@@ -16,17 +17,18 @@ from pytorch_lightning.callbacks import ModelCheckpoint
 
 from utils.base import display_mlflow_run_info, str2bool, fetch_from_mlflow, get_name, data_loader_mean_and_std
 from utils.debug import debug
+from utils.dataset_utils import k_fold
 from utils.augmentation import get_augmentation
-from utils.dataset import Subset, get_dataset, k_fold
+from dataset import Subset, get_dataset
 
-from processingpipeline.pipeline import RawProcessingPipeline
-from processingpipeline.torch_pipeline import raw2rgb, RawToRGB, ParametrizedProcessing, NNProcessing
+from processing.pipeline_numpy import RawProcessingPipeline
+from processing.pipeline_torch import add_additive_layer, raw2rgb, RawToRGB, ParametrizedProcessing, NNProcessing
 
-from models.classifier import log_tensor, resnet_model, LitModel, TrackImagesCallback
+from model import log_tensor, resnet_model, LitModel, TrackImagesCallback
 
 import segmentation_models_pytorch as smp
 
-from utils.pytorch_ssim import SSIM
+from utils.ssim import SSIM
 
 # args to set up task
 parser = argparse.ArgumentParser(description="classification_task")
@@ -106,41 +108,33 @@ parser.add_argument("--adv_training", action='store_true', help="Enable adversar
 parser.add_argument("--adv_aux_weight", type=float, default=1, help="Weighting of the adversarial auxilliary loss")
 parser.add_argument("--adv_aux_loss", type=str, default='ssim', choices=['l2', 'ssim'],
                     help="Type of adversarial auxilliary regularization loss")
+parser.add_argument("--adv_noise_layer", action='store_true', help="Adds an additive layer to Parametrized Processing")
+parser.add_argument("--adv_track_differences", action='store_true', help='Save difference to default pipeline')
+parser.add_argument('--adv_parameters', choices=['all', 'black_level', 'white_balance',
+                                                 'colour_correction', 'gamma_correct', 'sharpening_filter', 'gaussian_blur', 'additive_layer'])
 
 parser.add_argument("--cache_downloaded_models", type=str2bool, default=True)
 
 parser.add_argument('--test_run', action='store_true')
 
-if 'ipykernel_launcher' in sys.argv[0]:
-    args = parser.parse_args([
-        '--dataset=Microscopy',
-        '--epochs=100',
-        '--augmentation=strong',
-        '--lr=1e-5',
-        '--freeze_processor',
-        # '--track_processing',
-        # '--test_run',
-        # '--track_predictions',
-        # '--track_every_epoch',
-        # '--adv_training',
-        # '--adv_aux_weight=100',
-        # '--adv_aux_loss=l2',
-        # '--log_model=',
-    ])
-else:
-    args = parser.parse_args()
+
+args = parser.parse_args()
+
+os.makedirs('results', exist_ok=True)
 
 
 def run_train(args):
 
+    print(args)
+
     DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
     training_mode = 'adversarial' if args.adv_training else 'default'
 
     # set tracking uri, this is the address of the mlflow server where light experimental data will be stored
     mlflow.set_tracking_uri(args.tracking_uri)
     mlflow.set_experiment(args.experiment_name)
-    os.environ["AWS_ACCESS_KEY_ID"] = #TODO: add your AWS access key if you want to write your results to our collaborative lab server
-    os.environ["AWS_SECRET_ACCESS_KEY"] = #TODO: add your AWS seceret access key if you want to write your results to our collaborative lab server
+    os.environ["AWS_ACCESS_KEY_ID"] = "AKIAYYIYHFHKBIJHOJPA"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "eUSKKy+T+KzBKWvAw5PrM/MDwEgkE0LcpNWgnmir"
 
     # dataset
 
@@ -197,8 +191,9 @@ def run_train(args):
                 if args.processing_mode == 'parametrized':
                     processor = ParametrizedProcessing(
                         camera_parameters=dataset.camera_parameters, track_stages=track_stages, batch_norm_output=True,
-                        noise_layer=args.adv_training,   # XXX: Remove?
+                        # noise_layer=args.adv_noise_layer, # this has to be added manually afterwards for when a model is loaded that doesn't have one yet
                     )
+
                 elif args.processing_mode == 'neural_network':
                     processor = NNProcessing(track_stages=track_stages,
                                              normalize_mosaic=normalize_mosaic, batch_norm_output=True)
@@ -252,13 +247,19 @@ def run_train(args):
                 assert not args.freeze_processor, "Processor should not be frozen for adversarial training"
 
                 processor_default = copy.deepcopy(processor)
-                processor_default.track_stages = False
+                processor_default.track_stages = args.track_processing
                 processor_default.eval()
                 processor_default.to(DEVICE)
                 # debug(processor_default)
+                for p in processor_default.parameters():
+                    p.requires_grad = False
+
+                if args.adv_noise_layer:
+                    add_additive_layer(processor)
 
                 def l2_regularization(x, y):
-                    return (x - y).norm()
+                    return ((x - y) ** 2).sum()
+                    # return (x - y).norm()
 
                 if args.adv_aux_loss == 'l2':
                     regularization = l2_regularization
@@ -274,7 +275,8 @@ def run_train(args):
                         self.weight = weight
 
                     def forward(self, x):
-                        x_reference = processor_default(x)
+                        with torch.no_grad():
+                            x_reference = processor_default(x)
                         x_processed = processor.buffer['processed_rgb']
                         return self.weight * self.loss_aux(x_reference, x_processed)
 
@@ -290,7 +292,7 @@ def run_train(args):
                     def __repr__(self):
                         return f'{self.weight} * {get_name(self.loss)}'
 
-                loss = WeightedLoss(loss=nn.CrossEntropyLoss(), weight=-1)
+                loss = WeightedLoss(loss=loss, weight=-1)
                 # loss = WeightedLoss(loss=nn.CrossEntropyLoss(), weight=0)
                 loss_aux = AuxLoss(
                     loss_aux=regularization,
@@ -303,8 +305,10 @@ def run_train(args):
                 classifier=classifier,
                 processor=processor,
                 loss=loss,
+                lr=args.lr,
                 loss_aux=loss_aux,
                 adv_training=args.adv_training,
+                adv_parameters=args.adv_parameters,
                 metrics=metrics,
                 augmentation=augmentation,
                 is_segmentation_task=dataset.task == 'segmentation',
@@ -346,7 +350,7 @@ def run_train(args):
 
             with mlflow.start_run(run_name=f"{args.run_name}_{k_iter}", nested=True) as child_run:
 
-                #mlflow.pytorch.autolog(silent=True)
+                # mlflow.pytorch.autolog(silent=True)
 
                 if k_iter == 0:
                     display_mlflow_run_info(child_run)
@@ -373,19 +377,22 @@ def run_train(args):
                                                      tracking_uri=args.tracking_uri,)
                 mlf_logger._run_id = child_run.info.run_id
 
+                reference_processor = processor_default if args.adv_training and args.adv_track_differences else None
+
                 callbacks = []
                 if args.track_processing:
                     callbacks += [TrackImagesCallback(track_loader,
+                                                      reference_processor,
                                                       track_every_epoch=args.track_every_epoch,
                                                       track_processing=args.track_processing,
                                                       track_gradients=args.track_processing_gradients,
                                                       track_predictions=args.track_predictions,
                                                       save_tensors=args.track_save_tensors)]
 
-                #if True: #args.save_best:
+                # if True: #args.save_best:
                 #    if dataset.task == 'classification':
-                        #checkpoint_callback = ModelCheckpoint(pathmonitor="val_accuracy", mode='max')
-                #        checkpoint_callback = ModelCheckpoint(dirpath=args.tracking_uri, save_top_k=1, verbose=True, monitor="val_accuracy", mode="max") #dirpath=args.tracking_uri, 
+                    #checkpoint_callback = ModelCheckpoint(pathmonitor="val_accuracy", mode='max')
+                #        checkpoint_callback = ModelCheckpoint(dirpath=args.tracking_uri, save_top_k=1, verbose=True, monitor="val_accuracy", mode="max") #dirpath=args.tracking_uri,
                 #    else:
                 #        checkpoint_callback = ModelCheckpoint(monitor="val_iou_score")
                 #callbacks += [checkpoint_callback]
@@ -397,7 +404,7 @@ def run_train(args):
                     logger=mlf_logger,
                     callbacks=callbacks,
                     check_val_every_n_epoch=args.check_val_every_n_epoch,
-                    #checkpoint_callback=True,
+                    # checkpoint_callback=True,
                 )
 
                 if args.log_model:
@@ -410,9 +417,8 @@ def run_train(args):
                     val_dataloaders=valid_loader,
                 )
 
-                # if args.adv_training:
-                #     for (name, p1), p2 in zip(processor.named_parameters(), processor_default.cpu().parameters()):
-                #         print(f"param '{name}' diff: {p2 - p1}, l2: {(p2-p1).norm().item()}")
+    globals().update(locals())  # for convenient access
+
     return model
 
 

utils/augmentation.py

@@ -99,7 +99,7 @@ if __name__ == '__main__':
         os.chdir('..')
 
     # from utils.debug import debug
-    from utils.dataset import get_dataset
+    from dataset import get_dataset
     import matplotlib.pyplot as plt
 
     dataset = get_dataset('DS')  # drone segmentation

utils/{splitting.py → dataset_utils.py}

@@ -1,6 +1,3 @@
-"""
-Split images in blocks and vice versa
-"""
 
 import random
 import numpy as np
@@ -10,127 +7,183 @@ import torch
 from skimage.util.shape import view_as_windows
 
 
-def split_img(imgs, ROIs = (3,3) , step= (1,1)):
+def load_image(path):
+    file_type = path.split('.')[-1].lower()
+    if file_type == 'dng':
+        img = rawpy.imread(path).raw_image_visible
+    elif file_type == 'tiff' or file_type == 'tif':
+        img = np.array(tiff.imread(path), dtype=np.float32)
+    else:
+        img = np.array(Image.open(path), dtype=np.float32)
+    return img
+
+
+def list_images_in_dir(path):
+    image_list = [os.path.join(path, img_name)
+                  for img_name in sorted(os.listdir(path))
+                  if img_name.split('.')[-1].lower() in IMAGE_FILE_TYPES]
+    return image_list
+
+
+def k_fold(dataset, n_splits: int, seed: int, train_size: float):
+    """Split dataset in subsets for cross-validation 
+
+       Args:
+            dataset (class): dataset to split
+            n_split (int): Number of re-shuffling & splitting iterations.
+            seed (int): seed for k_fold splitting
+            train_size (float): should be between 0.0 and 1.0 and represent the proportion of the dataset to include in the train split.
+       Returns:
+           idxs (list): indeces for splitting the dataset. The list contain n_split pair of train/test indeces.
+    """
+    if hasattr(dataset, 'labels'):
+        x = dataset.images
+        y = dataset.labels
+    elif hasattr(dataset, 'masks'):
+        x = dataset.images
+        y = dataset.masks
+
+    idxs = []
+
+    if dataset.task == 'classification':
+        sss = StratifiedShuffleSplit(n_splits=n_splits, train_size=train_size, random_state=seed)
+
+        for idxs_train, idxs_test in sss.split(x, y):
+            idxs.append((idxs_train.tolist(), idxs_test.tolist()))
+
+    elif dataset.task == 'segmentation':
+        for n in range(n_splits):
+            split_idx = int(len(dataset) * train_size)
+            indices = np.random.permutation(len(dataset))
+            idxs.append((indices[:split_idx].tolist(), indices[split_idx:].tolist()))
+
+    return idxs
+
+
+def split_img(imgs, ROIs=(3, 3), step=(1, 1)):
     """Split the imgs in regions of size ROIs.
 
        Args:
             imgs (ndarray): images which you want to split 
             ROIs (tuple): size of sub-regions splitted (ROIs=region of interests)
             step (tuple): step path from one sub-region to the next one (in the x,y axis)
-       
+
         Returns:
             ndarray: splitted subimages. 
                      The size is (x_num_subROIs*y_num_subROIs, **) where:
                      x_num_subROIs = ( imgs.shape[1]-int(ROIs[1]/2)*2 )/step[1]
                      y_num_subROIs = ( imgs.shape[0]-int(ROIs[0]/2)*2 )/step[0]                     
-       
+
        Example:
             >>> from dataset_generator import split
             >>> imgs_splitted = split(imgs, ROI_size = (5,5), step=(2,3))
     """
-    
+
     if len(ROIs) > 2:
         return print("ROIs is a 2 element list")
-    
+
     if len(step) > 2:
         return print("step is a 2 element list")
-    
+
     if type(imgs) != type(np.array(1)):
         return print("imgs should be a ndarray")
-        
-    if len(imgs.shape) == 2: # Single image with one channel (HxW)   
-        splitted = view_as_windows(imgs, (ROIs[0],ROIs[1]), (step[0], step[1]))
+
+    if len(imgs.shape) == 2:  # Single image with one channel (HxW)
+        splitted = view_as_windows(imgs, (ROIs[0], ROIs[1]), (step[0], step[1]))
         return splitted.reshape((-1, ROIs[0], ROIs[1]))
-    
-    if len(imgs.shape) == 3: 
+
+    if len(imgs.shape) == 3:
         _, _, channels = imgs.shape
-        if channels <= 3: # Single image more channels (HxWxC)
-            splitted = view_as_windows(imgs, (ROIs[0],ROIs[1], channels), (step[0], step[1], channels))
+        if channels <= 3:  # Single image more channels (HxWxC)
+            splitted = view_as_windows(imgs, (ROIs[0], ROIs[1], channels), (step[0], step[1], channels))
             return splitted.reshape((-1, ROIs[0], ROIs[1], channels))
-        else: # More images with 1 channel
-            splitted = view_as_windows(imgs, (1, ROIs[0],ROIs[1]), (1, step[0], step[1]))
+        else:  # More images with 1 channel
+            splitted = view_as_windows(imgs, (1, ROIs[0], ROIs[1]), (1, step[0], step[1]))
             return splitted.reshape((-1, ROIs[0], ROIs[1]))
-    
-    if len(imgs.shape) == 4: # More images with more channels(BxHxWxC)
+
+    if len(imgs.shape) == 4:  # More images with more channels(BxHxWxC)
         _, _, _, channels = imgs.shape
-        splitted = view_as_windows(imgs, (1, ROIs[0],ROIs[1], channels), (1, step[0], step[1], channels))
+        splitted = view_as_windows(imgs, (1, ROIs[0], ROIs[1], channels), (1, step[0], step[1], channels))
         return splitted.reshape((-1, ROIs[0], ROIs[1], channels))
 
+
 def join_blocks(splitted, final_shape):
     """Join blocks to reobtain a splitted image
-    
+
     Attribute:
         splitted (tensor) = image splitted in blocks, size = (N_blocks, Channels, Height, Width)
         final_shape (tuple) = size of the final image reconstructed (Height, Width)
     Return:
         tensor: image restored from blocks. size = (Channels, Height, Width)
-    
+
     """
     n_blocks, channels, ROI_height, ROI_width = splitted.shape
-    
+
     rows = final_shape[0] // ROI_height
     columns = final_shape[1] // ROI_width
 
-    final_img = torch.empty(rows, channels, ROI_height, ROI_width*columns)
-    for r in range(rows):    
-        stackblocks = splitted[r*columns]
+    final_img = torch.empty(rows, channels, ROI_height, ROI_width * columns)
+    for r in range(rows):
+        stackblocks = splitted[r * columns]
         for c in range(1, columns):
-            stackblocks = torch.cat((stackblocks, splitted[r*columns+c]), axis=2)
+            stackblocks = torch.cat((stackblocks, splitted[r * columns + c]), axis=2)
         final_img[r] = stackblocks
-    
+
     joined_img = final_img[0]
-    
-    for i in np.arange(1, len(final_img)):    
-        joined_img = torch.cat((joined_img,final_img[i]), axis = 1)
-    
+
+    for i in np.arange(1, len(final_img)):
+        joined_img = torch.cat((joined_img, final_img[i]), axis=1)
+
     return joined_img
 
-def random_ROI(X, Y, ROIs = (512,512)):
+
+def random_ROI(X, Y, ROIs=(512, 512)):
     """ Return a random region for each input/target pair images of the dataset
         Args:
             Y (ndarray): target of your dataset --> size: (BxHxWxC)
             X (ndarray): input of your dataset --> size: (BxHxWxC)
             ROIs (tuple): size of random region (ROIs=region of interests)
-           
+
         Returns:
             For each pair images (input/target) of the dataset, return respectively random ROIs
             Y_cut (ndarray): target of your dataset --> size: (Batch,Channels,ROIs[0],ROIs[1])
             X_cut (ndarray): input of your dataset --> size: (Batch,Channels,ROIs[0],ROIs[1])
-            
+
         Example:
             >>> from dataset_generator import random_ROI
             >>> X,Y = random_ROI(X,Y, ROIs = (10,10))
-    """    
-    
+    """
+
     batch, channels, height, width = X.shape
-    
-    X_cut=np.empty((batch, ROIs[0], ROIs[1], channels))
-    Y_cut=np.empty((batch, ROIs[0], ROIs[1], channels))
-    
+
+    X_cut = np.empty((batch, ROIs[0], ROIs[1], channels))
+    Y_cut = np.empty((batch, ROIs[0], ROIs[1], channels))
+
     for i in np.arange(len(X)):
-        x_size=int(random.random()*(height-(ROIs[0]+1)))
-        y_size=int(random.random()*(width-(ROIs[1]+1)))
-        X_cut[i]=X[i, x_size:x_size+ROIs[0],y_size:y_size+ROIs[1], :]
-        Y_cut[i]=Y[i, x_size:x_size+ROIs[0],y_size:y_size+ROIs[1], :]
+        x_size = int(random.random() * (height - (ROIs[0] + 1)))
+        y_size = int(random.random() * (width - (ROIs[1] + 1)))
+        X_cut[i] = X[i, x_size:x_size + ROIs[0], y_size:y_size + ROIs[1], :]
+        Y_cut[i] = Y[i, x_size:x_size + ROIs[0], y_size:y_size + ROIs[1], :]
     return X_cut, Y_cut
 
-def one2many_random_ROI(X, Y, datasize=1000, ROIs = (512,512)):
+
+def one2many_random_ROI(X, Y, datasize=1000, ROIs=(512, 512)):
     """ Return a dataset of N subimages obtained from random regions of the same image
         Args:
             Y (ndarray): target of your dataset --> size: (1,H,W,C)
             X (ndarray): input of your dataset --> size: (1,H,W,C)
             datasize = number of random ROIs to generate
             ROIs (tuple): size of random region (ROIs=region of interests)
-           
+
         Returns:
             Y_cut (ndarray): target of your dataset --> size: (Datasize,ROIs[0],ROIs[1],Channels)
             X_cut (ndarray): input of your dataset --> size: (Datasize,ROIs[0],ROIs[1],Channels)
-    """   
+    """
 
     batch, channels, height, width = X.shape
-    
-    X_cut=np.empty((datasize, ROIs[0], ROIs[1], channels))
-    Y_cut=np.empty((datasize, ROIs[0], ROIs[1], channels))
+
+    X_cut = np.empty((datasize, ROIs[0], ROIs[1], channels))
+    Y_cut = np.empty((datasize, ROIs[0], ROIs[1], channels))
 
     for i in np.arange(datasize):
         X_cut[i], Y_cut[i] = random_ROI(X, Y, ROIs)

utils/debug.py

@@ -1,371 +0,0 @@
-import torch
-import numpy as np
-import inspect
-from functools import reduce, wraps
-from collections.abc import Iterable
-from IPython import embed
-
-try:
-    get_ipython()  # pylint: disable=undefined-variable
-    interactive_notebook = True
-except:
-    interactive_notebook = False
-
-_NONE = "__UNSET_VARIABLE__"
-
-
-def debug_init():
-    debug.disable = False
-    debug.silent = False
-    debug.verbose = 2
-    debug.expand_ignore = ["DataLoader", "Dataset", "Subset"]
-    debug.max_expand = 10
-    debug.show_tensor = False
-    debug.raise_exception = True
-    debug.full_stack = True
-    debug.restore_defaults_on_exception = not interactive_notebook
-    debug._indent = 0
-    debug._stack = ""
-
-    debug.embed = embed
-    debug.show = debug_show
-    debug.pause = debug_pause
-
-
-def debug_pause():
-    input("Press Enter to continue...")
-
-
-def debug(*args, assert_true=False):
-    """Decorator for debugging functions and tensors.
-    Will throw an exception as soon as a nan is encountered.
-    If used on iterables, these will be expanded and also searched for nans.
-    Usage:
-        debug(x)
-    Or:
-        @debug
-        def function():
-            ...
-    If used as a function wrapper, all arguments will be searched and printed.
-    """
-
-    single_arg = len(args) == 1
-
-    if debug.disable:
-        return args[0] if single_arg else None
-
-    try:
-        call_line = ''.join(inspect.stack()[1][4]).strip()
-    except:
-        call_line = '...'
-    used_as_wrapper = 'def ' == call_line[:4]
-    expect_return_arg = single_arg and 'debug' in call_line and call_line.split('debug')[0].strip() != ''
-    is_func = single_arg and hasattr(args[0], '__call__')
-
-    if is_func and (used_as_wrapper or expect_return_arg):
-        func = args[0]
-        sig_parameters = inspect.signature(func).parameters
-        sig_argnames = [p.name for p in sig_parameters.values()]
-        sig_defaults = {
-            k: v.default
-            for k, v in sig_parameters.items()
-            if v.default is not inspect.Parameter.empty
-        }
-
-        @wraps(func)
-        def _func(*args, **kwargs):
-            if debug.disable:
-                return func(*args, **kwargs)
-
-            if debug._indent == 0:
-                debug._stack = ""
-            stack_before = debug._stack
-            indent = ' ' * 4 * debug._indent
-            debug._indent += 1
-
-            args_kw = dict(zip(sig_argnames, args))
-            defaults = {k: v for k, v in sig_defaults.items()
-                        if k not in kwargs
-                        if k not in args_kw}
-            all_args = {**args_kw, **kwargs, **defaults}
-
-            func_name = None
-            if hasattr(func, '__name__'):
-                func_name = func.__name__
-            elif hasattr(func, '__class__'):
-                func_name = func.__class__.__name__
-
-            if func_name is None:
-                func_name = '... ' + call_line + '...'
-            else:
-                func_name = '@' + func_name + '()'
-
-            _debug_log('', indent=indent)
-            _debug_log(func_name, indent=indent)
-
-            debug._last_call = func
-            debug._last_args = all_args
-            debug._last_args_sig = sig_argnames
-
-            for argtype, params in [("args", args_kw.items()),
-                                    ("kwargs", kwargs.items()),
-                                    ("defaults", defaults.items())]:
-                if params:
-                    _debug_log(f"{argtype}:", indent=indent + ' ' * 6)
-                for argname, arg in params:
-                    if argname == 'self':
-                        # _debug_log(f"- self:  ...", indent=indent + ' ' * 8)
-                        pass
-                    else:
-                        _debug_log(f"- {argname}:  ", arg, indent + ' ' * 8, assert_true)
-            try:
-                out = func(*args, **kwargs)
-            except:
-                _debug_crash_save()
-                debug._stack = ""
-                debug._indent = 0
-                raise
-            debug.out = out
-            _debug_log("returned:  ", out, indent, assert_true)
-            _debug_log('', indent=indent)
-            debug._indent -= 1
-            if not debug.full_stack:
-                debug._stack = stack_before
-            return out
-        return _func
-    else:
-        if debug._indent == 0:
-            debug._stack = ""
-        argname = ')'.join('('.join(call_line.split('(')[1:]).split(')')[:-1])
-        if assert_true:
-            argname = ','.join(argname.split(',')[:-1])
-            _debug_log(f"assert{{{argname}}}  ", args[0], ' ' * 4 * debug._indent, assert_true)
-        else:
-            for arg in args:
-                _debug_log(f"{{{argname}}}  =  ", arg, ' ' * 4 * debug._indent, assert_true)
-        if expect_return_arg:
-            return args[0]
-        return
-
-
-def is_iterable(x):
-    return isinstance(x, Iterable) or hasattr(x, '__getitem__') and not isinstance(x, str)
-
-
-def ndarray_repr(t, assert_all=False):
-    exception_encountered = False
-    info = []
-    shape = tuple(t.shape)
-    single_entry = shape == () or shape == (1,)
-    if single_entry:
-        info.append(f"[{t.item():.4f}]")
-    else:
-        info.append(f"({', '.join(map(repr, shape))})")
-    invalid_sum = (~np.isfinite(t)).sum().item()
-    if invalid_sum:
-        info.append(
-            f"{invalid_sum} INVALID ENTR{'Y' if invalid_sum == 1 else 'IES'}")
-        exception_encountered = True
-    if debug.verbose > 1:
-        if not invalid_sum and not single_entry:
-            info.append(f"|x|={np.linalg.norm(t):.1f}")
-            if t.size:
-                info.append(f"x in [{t.min():.1f}, {t.max():.1f}]")
-    if debug.verbose and t.dtype != np.float:
-        info.append(f"dtype={str(t.dtype)}".replace("'", ''))
-    if assert_all:
-        assert_val = t.all()
-        if not assert_val:
-            exception_encountered = True
-    if assert_all and not exception_encountered:
-        output = "passed"
-    else:
-        if assert_all and not assert_val:
-            output = f"ndarray({info[0]})"
-        else:
-            output = f"ndarray({', '.join(info)})"
-    if exception_encountered and (not hasattr(debug, 'raise_exception') or debug.raise_exception):
-        if debug.restore_defaults_on_exception:
-            debug.raise_exception = False
-            debug.silent = False
-        debug.x = t
-        msg = output
-        debug._stack += output
-        if debug._stack and '\n' in debug._stack:
-            msg += '\nSTACK:  ' + debug._stack
-        if assert_all:
-            assert assert_val, "Assert did not pass on " + msg
-        raise Exception("Invalid entries encountered in " + msg)
-    return output
-
-
-def tensor_repr(t, assert_all=False):
-    exception_encountered = False
-    info = []
-    shape = tuple(t.shape)
-    single_entry = shape == () or shape == (1,)
-    if single_entry:
-        info.append(f"[{t.item():.3f}]")
-    else:
-        info.append(f"({', '.join(map(repr, shape))})")
-    invalid_sum = (~torch.isfinite(t)).sum().item()
-    if invalid_sum:
-        info.append(
-            f"{invalid_sum} INVALID ENTR{'Y' if invalid_sum == 1 else 'IES'}")
-        exception_encountered = True
-    if debug.verbose and t.requires_grad:
-        info.append('req_grad')
-        if debug.verbose > 2:
-            if t.is_leaf:
-                info.append('leaf')
-            if hasattr(t, 'retains_grad') and t.retains_grad:
-                info.append('retains_grad')
-    has_grad = (t.is_leaf or hasattr(t, 'retains_grad') and t.retains_grad) and t.grad is not None
-    if has_grad:
-        grad_invalid_sum = (~torch.isfinite(t.grad)).sum().item()
-        if grad_invalid_sum:
-            info.append(
-                f"GRAD {grad_invalid_sum} INVALID ENTR{'Y' if grad_invalid_sum == 1 else 'IES'}")
-            exception_encountered = True
-    if debug.verbose > 1:
-        if not invalid_sum and not single_entry:
-            info.append(f"|x|={t.float().norm():.1f}")
-            if t.numel():
-                info.append(f"x in [{t.min():.2f}, {t.max():.2f}]")
-        if has_grad and not grad_invalid_sum:
-            if single_entry:
-                info.append(f"grad={t.grad.float().item():.3f}")
-            else:
-                info.append(f"|grad|={t.grad.float().norm():.1f}")
-    if debug.verbose and t.dtype != torch.float:
-        info.append(f"dtype={str(t.dtype).split('.')[-1]}")
-    if debug.verbose and t.device.type != 'cpu':
-        info.append(f"device={t.device.type}")
-    if assert_all:
-        assert_val = t.all()
-        if not assert_val:
-            exception_encountered = True
-    if assert_all and not exception_encountered:
-        output = "passed"
-    else:
-        if assert_all and not assert_val:
-            output = f"tensor({info[0]})"
-        else:
-            output = f"tensor({', '.join(info)})"
-    if exception_encountered and (not hasattr(debug, 'raise_exception') or debug.raise_exception):
-        if debug.restore_defaults_on_exception:
-            debug.raise_exception = False
-            debug.silent = False
-        debug.x = t
-        msg = output
-        debug._stack += output
-        if debug._stack and '\n' in debug._stack:
-            msg += '\nSTACK:  ' + debug._stack
-        if assert_all:
-            assert assert_val, "Assert did not pass on " + msg
-        raise Exception("Invalid entries encountered in " + msg)
-    return output
-
-
-def _debug_crash_save():
-    if debug._indent:
-        debug.args = debug._last_args
-        debug.func = debug._last_call
-
-        @wraps(debug.func)
-        def _recall(*args, **kwargs):
-            call_args = {**debug.args, **kwargs, **dict(zip(debug._last_args_sig, args))}
-            return debug(debug.func)(**call_args)
-
-        def print_stack(stack=debug._stack):
-            print('\nSTACK:  ' + stack)
-        debug.stack = print_stack
-
-        debug.recall = _recall
-    debug._indent = 0
-
-
-def _debug_log(output, var=_NONE, indent='', assert_true=False, expand=True):
-    debug._stack += indent + output
-    if not debug.silent:
-        print(indent + output, end='')
-    if var is not _NONE:
-        type_str = type(var).__name__.lower()
-        if var is None:
-            _debug_log('None')
-        elif isinstance(var, str):
-            _debug_log(f"'{var}'")
-        elif type_str == 'ndarray':
-            _debug_log(ndarray_repr(var, assert_true))
-            if debug.show_tensor:
-                _debug_show_print(var, indent=indent + 4 * ' ')
-        # elif type_str in ('tensor', 'parameter'):
-        elif type_str == 'tensor':
-            _debug_log(tensor_repr(var, assert_true))
-            if debug.show_tensor:
-                _debug_show_print(var, indent=indent + 4 * ' ')
-        elif hasattr(var, 'named_parameters'):
-            _debug_log(type_str)
-            params = list(var.named_parameters())
-            _debug_log(f"{type_str}[{len(params)}] {{")
-            for k, v in params:
-                _debug_log(f"'{k}': ", v, indent + 6 * ' ')
-            _debug_log(indent + 4 * ' ' + '}')
-        elif is_iterable(var):
-            expand = debug.expand_ignore != '*' and expand
-            if expand:
-                if isinstance(debug.expand_ignore, str):
-                    if type_str == str(debug.expand_ignore).lower():
-                        expand = False
-                elif is_iterable(debug.expand_ignore):
-                    for ignore in debug.expand_ignore:
-                        if type_str == ignore.lower():
-                            expand = False
-            if hasattr(var, '__len__'):
-                length = len(var)
-            else:
-                var = list(var)
-                length = len(var)
-            if expand and length > 0:
-                _debug_log(f"{type_str}[{length}] {{")
-                if isinstance(var, dict):
-                    for k, v in var.items():
-                        _debug_log(f"'{k}': ", v, indent + 6 * ' ', assert_true)
-                else:
-                    i = 0
-                    for k, i in zip(var, range(debug.max_expand)):
-                        _debug_log('- ', k, indent + 6 * ' ', assert_true)
-                    if i < length - 1:
-                        _debug_log('- ' + ' ' * 6 + '...', indent=indent + 6 * ' ')
-                _debug_log(indent + 4 * ' ' + '}')
-            else:
-                _debug_log(f"{type_str}[{length}]")
-        else:
-            _debug_log(str(var))
-    else:
-        debug._stack += '\n'
-        if not debug.silent:
-            print()
-
-
-def debug_show(x):
-    assert is_iterable(x)
-    debug(x)
-    _debug_show_print(x, indent=' ' * 4 * debug._indent)
-
-
-def _debug_show_print(x, indent=''):
-    is_tensor = type(x).__name__ in ('Tensor', 'ndarray')
-    if is_tensor:
-        x = x.flatten()
-    if type(x).__name__ == 'Tensor' and x.dim() == 0:
-        return
-    n_samples = min(10, len(x))
-    di = len(x) // n_samples
-    var = list(x[i * di] for i in range(n_samples))
-    if is_tensor or type(var[0]) == float:
-        var = [round(float(v), 4) for v in var]
-    _debug_log('-->  ', str(var), indent, expand=False)
-
-
-debug_init()

utils/mutual_entropy.py

@@ -1,193 +0,0 @@
-import numpy as np
-from PIL import Image
-import matplotlib.pyplot as plt
-from scipy.signal import convolve2d
-
-def mse(x,y):
-    return ((x-y)**2).mean()
-
-def gaussian_noise_entropies(t1, bins=20):
-    all_MI= []
-    all_mse = []
-    for sigma in np.linspace(0,100,201):
-        t2 = np.random.normal(t1.copy(), scale=sigma, size = t1.shape)
-        hist_2d, x_edges, y_edges = np.histogram2d(
-                                             t1.ravel(),
-                                             t2.ravel(),
-                                             bins=bins)
-        all_mse.append(mse(t1,t2))
-        MI = mutual_information(hist_2d)
-        all_MI.append(MI)
-        
-    return np.array((all_MI)), np.array((all_mse))
-        
-def shifts_entropies(t1, bins=20):
-    all_MI=[]
-    all_mse=[]
-    for N in np.linspace(1,50,50):
-        N = int(N)
-        temp_t2 = t1[:-N].copy()
-        temp_t1 = t1[N:].copy()
-        hist_2d, x_edges, y_edges = np.histogram2d(
-                                     t1.ravel(),
-                                     t2.ravel(),
-                                     bins=bins)
-        MI = mutual_information(hist_2d)
-        
-        all_mse.append(mse(temp_t1,temp_t2))
-        all_MI.append(MI)
-        
-    return np.array((all_MI)), np.array((all_mse))
-
-def mutual_information(hgram):
-    """ Mutual information for joint histogram
-    """
-    # Convert bins counts to probability values
-    pxy = hgram / float(np.sum(hgram))
-    px = np.sum(pxy, axis=1) # marginal for x over y
-    py = np.sum(pxy, axis=0) # marginal for y over x
-    px_py = px[:, None] * py[None, :] # Broadcast to multiply marginals
-    # Now we can do the calculation using the pxy, px_py 2D arrays
-    nzs = pxy > 0 # Only non-zero pxy values contribute to the sum
-    return np.sum(pxy[nzs] * np.log(pxy[nzs] / px_py[nzs]))
-
-def entropy(image, bins=20):
-    image = image.ravel()   
-    hist, bin_edges = np.histogram(image, bins = bins)
-    hist = hist/hist.sum()
-    entropy_term = np.where(hist != 0, hist*np.log(hist), 0)
-    entropy = - np.sum(entropy_term)
-    
-    return entropy
-
-# Gray Image
-# t1 = np.array(Image.open("img.png"))[:,:,0].astype(float)
-
-# Colour Image
-t1 = np.array(Image.open("img.png").resize((255,255)))
-
-perturb = "gauss"
-show_image = True
-bins=20
-
-print(perturb)
-
-# Identity
-if perturb == "identity":
-    t2 = t1
-    title = "Identity"
-    image1 = "Clean"
-    image2 = "Clean"
-
-# Poisson Noise on t2
-if perturb == "poisson":
-    t2 = np.random.poisson(t1)
-    title = "Poisson Noise"
-    image1 = "Clean"
-    image2 = "Noisy"
-
-# Gaussian Noise on t2
-if perturb == "gauss":
-    print(np.shape(t1))
-    sigma = 50.0
-    t2 = np.random.normal(t1.copy(), scale=sigma, size = t1.shape)
-    if "grad" in locals():
-        title = f"Gaussian Noise, grad= True, sigma = {sigma:.2f}"
-    else:
-        title = f"Gaussian Noise, sigma = {sigma:.2f}"
-    image1 = "Clean"
-    image2 = "Noisy"
-
-if perturb == "box":
-    sigma = 50.0
-    mean = np.mean(t1)
-    print(np.shape(t1))
-    t2 = t1.copy()
-    t2[30:220,50:120,:] = mean
-    title = "Box with mean pixels"
-    image1 = "Clean"
-    image2 = "Noisy"
-
-
-# Shift t2 on y axis
-if perturb == "shift":
-    N=30
-    t2 = t1[:-N]
-    t1 = t1[N:]
-    title = "y shift"
-    image1 = "Clean"
-    image2 = "Shifted"
-
-t2 = np.clip(t2,0,255).astype(int)
-
-print("Correlation Coefficient: ",np.corrcoef(t1.ravel(), t2.ravel())[0, 1])
-
-# 2D Histogram
-hist_2d, x_edges, y_edges = np.histogram2d(
-     t1.ravel(),
-     t2.ravel(),
-     bins=bins)
-
-MI = mutual_information(hist_2d)
-
-print("Mutual Information", MI)
-print("Mean squared error:", mse(t1,t2))
-
-if show_image == True:
-    plt.figure()
-    plt.imshow(np.hstack((t2, t1)))
-    plt.title(title)
-    
-    plt.figure()
-    
-    plt.plot(t1.ravel(), t2.ravel(), '.')
-    plt.xlabel(image1)
-    plt.ylabel(image2)
-    plt.title('I1 vs I2')
-    
-    plt.figure()
-    plt.imshow((hist_2d.T)/hist_2d.max(), origin='lower')
-    plt.xlabel(image1)
-    plt.ylabel(image2)
-    plt.xticks(ticks=np.linspace(0,bins-1,10), labels=np.linspace(x_edges.min(),x_edges.max(),10).astype(int))
-    plt.yticks(ticks=np.linspace(0,bins-1,10), labels=np.linspace(y_edges.min(),y_edges.max(),10).astype(int))
-    plt.title('p(x,y)')
-    plt.colorbar()
-    
-    # Show log histogram, avoiding divide by 0
-    plt.figure(figsize=(4,4))
-    hist_2d_log = np.zeros(hist_2d.shape)
-    non_zeros = hist_2d != 0
-    hist_2d_log[non_zeros] = np.log(hist_2d[non_zeros])
-    plt.imshow((hist_2d_log.T)/hist_2d_log.max(), origin='lower')
-    plt.xlabel(image1)
-    plt.ylabel(image2)
-    plt.xticks(ticks=np.linspace(0,bins-1,10), labels=np.linspace(x_edges.min(),x_edges.max(),10).astype(int))
-    plt.yticks(ticks=np.linspace(0,bins-1,10), labels=np.linspace(y_edges.min(),y_edges.max(),10).astype(int))
-    plt.title('log(p(x,y))')
-    plt.colorbar()
-    plt.show()
-
-if perturb == "shift":
-    mi_array, mse_array = shifts_entropies(t1, bins=bins)
-    plt.figure()
-    plt.plot(np.linspace(0,50,50), mi_array)
-    plt.xlabel("y shift")
-    plt.ylabel("Mutual Information")
-    plt.figure()
-    plt.plot(np.linspace(0,50,50), mse_array)
-    plt.xlabel("y shift")
-    plt.ylabel("Mean Squared Error")
-    plt.show()
-        
-if perturb == "gauss":
-    mi_array, mse_array = gaussian_noise_entropies(t1, bins= bins)
-    plt.figure()
-    plt.plot(np.linspace(0,100,201), mi_array)
-    plt.xlabel("sigma")
-    plt.ylabel("Mutual Information")
-    plt.figure()
-    plt.plot(np.linspace(0,100,201), mse_array)
-    plt.xlabel("sigma")
-    plt.ylabel("Mean Squared Error")
-    plt.show()