fix aws access

.gitattributes

@@ -0,0 +1 @@
+*.ipynb filter=lfs diff=lfs merge=lfs -text

ABtesting.py

@@ -0,0 +1,806 @@
+import os
+import argparse
+import json
+from cv2 import transform
+
+import torch
+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 utils.base import get_mlflow_model_by_name, SmartFormatter
+from processingpipeline.pipeline import RawProcessingPipeline
+
+from utils.Cperturb import Distortions
+
+import segmentation_models_pytorch as smp
+
+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'), 
+                    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')
+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')
+
+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()
+        N_elements = torch.sum(self.cm)
+        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
+
+    def recall(self):
+        Tp_Fn = torch.sum(self.cm, 0)
+        Tp_Fn[Tp_Fn == 0] = 1
+        return torch.diagonal(self.cm,0) / Tp_Fn
+
+    def f1_score(self):
+        prod = (self.precision()*self.recall())
+        sum = (self.precision() + self.recall())
+        sum[sum == 0.] = 1.
+        return 2*( prod / sum )
+
+    def over_N_runs(ms, N_runs):
+        m, m2 = 0, 0 
+
+        for i in ms:
+            m += i 
+        mu = m/N_runs
+        
+        for i in ms:
+            m2 += (i-mu)**2
+
+        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):
+        self.experiment_name = 'ABtesting'
+        self.dataset_name = dataset_name
+        self.augmentation = augmentation
+        self.dm_train = dm_train
+        self.s_train = s_train
+        self.dn_train = dn_train
+        self.dm_test = dm_test
+        self.s_test = s_test
+        self.dn_test = dn_test
+        self.N_runs = N_runs
+        self.severity = severity
+        self.transform = transform
+        self.download_model = download_model
+
+    def static_pip_val(self, debayer=None, sharpening=None, denoising=None, severity=None, transform=None, plot_mode=False):
+
+        if debayer == None:
+            debayer = self.dm_test
+        if sharpening == None:
+            sharpening = self.s_test
+        if denoising == None:
+            denoising = self.dn_test   
+        if severity == None:
+            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":
+            mean = torch.tensor([0.35, 0.36, 0.35])
+            std = torch.tensor([0.12, 0.11, 0.12])
+        elif self.dataset_name == "Microscopy":
+            mean = torch.tensor([0.91, 0.84, 0.94])
+            std = torch.tensor([0.08, 0.12, 0.05])
+
+        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)])
+        else:
+             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')     
+
+        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)
+
+            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()
+
+            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
+
+            m = metrics(confusion_matrix)
+
+            accuracies.append(m.accuracy()) 
+            precisions.append(m.precision())
+            recalls.append(m.recall())
+            f1_scores.append(m.f1_score())
+
+            os.system('rm -r /tmp/t*')
+
+        accuracy = metrics.over_N_runs(accuracies, self.N_runs)
+        precision = metrics.over_N_runs(precisions, self.N_runs)
+        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')     
+
+        IoUs = []
+
+        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)
+
+            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
+
+            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 = 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           
+
+    def ABShowImages(self):
+
+        path = 'results/ABtesting/imgs/'
+        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]}')
+
+        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)
+
+        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']):
+            
+            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))
+            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.figure()
+                    # plt.imshow(diff.permute(1,2,0))
+                    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
+                    # 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'))
+
+    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']
+
+        fig = plt.figure()
+        columns=4
+        rows=3
+
+        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]
+                            
+                    fig.add_subplot(rows, columns, i)
+                    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
+
+        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)
+
+        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)
+        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}'))
+            
+    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']
+
+        for i,t in enumerate(transforms):
+                    
+            fig = plt.figure(figsize=(10,6))
+            columns = 5
+            rows = 1
+
+            for sev in range(1,6):
+
+                dataset = self.static_pip_val(severity=sev, transform=t, plot_mode=True)
+
+                img,_ = dataset[0]
+                        
+                fig.add_subplot(rows, columns, sev)
+                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:]
+
+            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']
+
+    path='results/ABtesting/tables'
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+    runs={}
+    i=0
+
+    for dm_train in demosaicings:
+        for s_train in sharpenings:
+            for dn_train in denoisings:
+                for dm_test in demosaicings:
+                    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] 
+                            runs[f'run{i}'] = {
+                            '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': 
+                                IoU = ABclass.ABsegmentation()
+                                runs[f'run{i}']['IoU'] = IoU
+                            else:
+                                classes, accuracy, precision, recall, f1_score = ABclass.ABclassification()
+                                runs[f'run{i}']['classes'] = classes
+                                runs[f'run{i}']['accuracy'] = accuracy
+                                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:
+                                json.dump(runs, outfile)
+
+                            i+=1
+
+def ABShowTable(dataset_name: str, augmentation: str):
+    
+    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')
+
+    with open(json_file, 'r') as run_file:
+        runs = json.load(run_file)
+
+        metrics=torch.zeros((2,12,12))
+        classes=[]
+
+        i,j=0,0
+
+        for r in range(len(runs)):
+                    
+            run = runs['run'+str(r)]
+            if dataset_name == 'DroneSegmentation':
+                acc = run['IoU']
+            else:
+                acc = run['accuracy']
+            if len(classes) < 12:
+                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)
+
+            metrics[0,j,i] = mu
+            metrics[1,j,i] = sigma
+           
+            i+=1
+
+            if i == 12:
+                i=0
+                j+=1
+
+    differences = torch.zeros_like(metrics)
+
+    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)
+
+    # Plot with scatter
+        
+    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)
+            step = 0.02
+        elif i == 1:
+            vmin = round(img[0].min().item(),2)
+            if augmentation == 'none':
+                vmax = min(0.15, round(img[0].max().item(),2))
+            if augmentation == 'weak':
+                vmax = min(0.08, round(img[0].max().item(),2))
+            if augmentation == 'strong':
+                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))
+        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]
+        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')
+        ticks = torch.arange(vmin, vmax, step).tolist()
+        ticks = [round(tick,2) for tick in ticks]
+        if ticks[-1] != vmax:
+            ticks.append(vmax)
+        cbb = plt.colorbar(pad=0.06)
+        cbb.set_ticks(ticks)
+        if i == 0:
+            ticks[0] = f'<{str(ticks[0])}'
+        elif i == 1:
+            ticks[-1] = f'>{str(ticks[-1])}'
+        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)
+                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'
+                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))
+        ax.set_xticklabels(classes)
+        ax.set_yticks(torch.linspace(0,11,12))
+        classes.reverse()
+        ax.set_yticklabels(classes)
+        classes.reverse()
+        plt.xticks(rotation = 45)
+        plt.yticks(rotation = 45)
+        cba.set_label('Standard Deviation')
+        plt.xlabel("Test pipelines")
+        plt.ylabel("Train pipelines")
+        plt.title(f'Dataset: {dataset_name}, Augmentation: {augmentation}')
+        if i == 0:
+            if dataset_name == 'DroneSegmentation':
+                cbb.set_label('IoU')
+                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"))
+        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"))
+
+def CMakeTable(dataset_name: str, augmentation: str, severity: int, N_runs: int, download_model: bool):
+    
+    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']
+
+    runs={}
+    i=0
+
+    for dm in demosaicings:
+        for s in sharpenings:
+            for dn in denoisings:
+                for t in transformations:
+                    pip = [dm,s,dn]
+                    runs[f'run{i}'] = {
+                    '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': 
+                        IoU = ABclass.ABsegmentation()
+                        runs[f'run{i}']['IoU'] = IoU
+                    else:
+                        classes, accuracy, precision, recall, f1_score = ABclass.ABclassification()
+                        runs[f'run{i}']['classes'] = classes
+                        runs[f'run{i}']['accuracy'] = accuracy
+                        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.json'), 'w') as outfile:
+                        json.dump(runs, outfile)
+
+                    i+=1
+
+def CShowTable(dataset_name, augmentation):
+
+    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']
+
+    pip = []
+    
+    demosaicings=['bilinear','malvar2004', 'menon2007']
+    sharpenings=['sharpening_filter', 'unsharp_masking']
+    denoisings=['median_denoising', 'gaussian_denoising']
+
+    for dm in demosaicings:
+        for s in sharpenings:
+            for dn in denoisings:
+                pip.append(f'{dm[:2]},{s[0]},{dn[2]}')
+
+    with open(json_file, 'r') as run_file:
+        runs = json.load(run_file)
+
+        metrics=torch.zeros((2,len(pip),len(transforms)))
+
+        i,j=0,0
+
+        for r in range(len(runs)):
+                    
+            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)
+
+            metrics[0,j,i] = mu
+            metrics[1,j,i] = sigma
+           
+            i+=1
+
+            if i == len(transforms):
+                i=0
+                j+=1
+
+        # Plot with scatter
+
+        img = metrics
+
+        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))
+        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]
+        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')
+        ticks = torch.arange(vmin, vmax, step).tolist()
+        ticks = [round(tick,2) for tick in ticks]
+        if ticks[-1] != vmax:
+            ticks.append(vmax)
+        cbb = plt.colorbar(pad=0.06)
+        cbb.set_ticks(ticks)
+        if i == 0:
+            ticks[0] = f'<{str(ticks[0])}'
+        elif i == 1:
+            ticks[-1] = f'>{str(ticks[-1])}'
+        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)
+                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'
+                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))
+        ax.set_xticklabels(transforms)
+        ax.set_yticks(torch.linspace(0,11,12))
+        pip.reverse()
+        ax.set_yticklabels(pip)
+        pip.reverse()
+        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"))
+        else:
+            cbb.set_label('Accuracy')
+            plt.savefig(os.path.join(path,f"{dataset_name}_{augmentation}_accuracies.png"))
+
+if __name__ == '__main__':
+     
+    if args.mode == 'ABMakeTable': 
+        ABMakeTable(args.dataset_name, args.augmentation, args.N_runs, args.download_model)
+    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, 
+                            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)
+        ABclass.ABShowAllImages()
+    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
+        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)
+        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)
+        ABclass.CShowAllImages()

README.md

@@ -0,0 +1,26 @@
+# Perturbed Minds
+
+## Conda environment and dependencies
+
+To make running this code easier you can install the latest conda environment for this project stored in `perturbed-environment.yml`.
+
+### Install environment from `perturbed-environment.yml`
+
+If you want to install the latest conda environment run
+
+`conda env create -f perturbed-environment.yml` 
+
+### Install segmentation_models_pytorch newest version
+
+PyPi version is not up-to-date with github version and lacks features
+
+`python -m pip install git+https://github.com/qubvel/segmentation_models.pytorch`
+
+### Update `perturbed-environment.yml`
+
+If you add code that requires new packages, inside your perturbed-minds conda environment run
+
+`conda env export > perturbed-environment.yml`
+
+## Walk-through
+Link to the repository structure we put down in miro: https://miro.com/app/board/o9J_lQdgyf8=/

figure1.sh

@@ -0,0 +1,7 @@
+python figures.py \
+--experiment_name track-test \
+--run_name track-all \
+--representation gradients \
+--step gamma_correct \
+--gif_name gradient \
+--output gif \

figure2.sh

@@ -0,0 +1,4 @@
+python figures.py \
+--experiment_name track-test \
+--run_name track-all \
+--output train_vs_val_loss \

figures.py

@@ -0,0 +1,92 @@
+import mlflow
+from mlflow.tracking import MlflowClient
+from mlflow.entities import ViewType
+import argparse
+#gif
+import os
+import pathlib
+import shutil
+import imageio
+#plot
+import matplotlib.pyplot as plt
+import numpy as np
+
+# -1. parse args
+parser = argparse.ArgumentParser(description="results_analysis")
+parser.add_argument("--tracking_uri", type=str,
+                    default="http://deplo-mlflo-1ssxo94f973sj-890390d809901dbf.elb.eu-central-1.amazonaws.com", help='URI of the mlflow server on AWS')
+parser.add_argument("--experiment_name", type=str, default=None,
+                    help='Name of the experiment on the mlflow server, e.g. "processing_comparison"')
+parser.add_argument("--run_name", type=str, default=None,
+                    help='Name of the run on the mlflow server, e.g. "proc_nn"')
+parser.add_argument("--representation", type=str, default=None,
+                    choices=["processing", "gradients"], help='The representation form you want retrieve("processing" or "gradients")')
+parser.add_argument("--step", type=str, default=None,
+                    choices=["pre_debayer", "demosaic", "color_correct", "sharpening", "gaussian", "clipped", "gamma_correct", "rgb"], 
+                    help='The processing step you want to track ("pre_debayer" or "rgb")') #TODO: include predictions and ground truths
+parser.add_argument("--gif_name", type=str, default=None,
+                    help='Name of the gif that will be saved. Note: .gif will be added later by script') #TODO: option to include filepath where result should be written
+                    #TODO: option to write results to existing run on mlflow
+parser.add_argument("--local_dir", type=str, default=None,
+                    help='Name of the local dir to be created to store mlflow data')
+parser.add_argument("--cleanup", type=bool, default=True,
+                    help='Whether to delete the local dir again after the script was run')
+parser.add_argument("--output", type=str, default=None,
+                    choices=["gif", "train_vs_val_loss"],
+                    help='Which output to generate') #TODO: make this cleaner, atm it is confusing because each figure may need different set of args and it is not clear how to manage that
+                    #TODO: idea -> fix the types of args for each figure which define the figure type but parametrize those things that can reasonably vary
+args = parser.parse_args()
+
+# 0. mlflow basics
+mlflow.set_tracking_uri(args.tracking_uri)
+
+# 1. specify experiment_name, run_name, representation and step
+#is done via parse_args
+
+# 2. use get_experiment_by_name to get experiment object
+experiment = mlflow.get_experiment_by_name(args.experiment_name)
+
+# 3. extract experiment_id
+#experiment.experiment_id
+
+# 4. use search_runs with experiment_id and run_name for string search query
+filter_string = "tags.mlflow.runName = '{}'".format(args.run_name) #create the filter string with using the runName tag to query mlflow
+runs = mlflow.search_runs(experiment.experiment_id, filter_string=filter_string) #returns a pandas data frame where each row is a run (if several exist under that name)
+client = MlflowClient() #TODO: look more into the options of client
+
+if args.output == "gif": #TODO: outsource these options to functions which are then loaded and can be called
+    # 5. extract run from list
+    #TODO: parent run and cv option for analysis
+    if args.local_dir:
+        local_dir = args.local_dir+"/artifacts"
+    else: #use the current working dir and make a subdir "artifacts" to store the data from mlflow
+        local_dir = str(pathlib.Path().resolve())+"/artifacts"
+    if not os.path.isdir('artifacts'):
+        os.mkdir(local_dir) #create the local_dir if it does not exist, yet #TODO: more advanced catching of existing files etc
+    dir = client.download_artifacts(runs["run_id"][0], "results", local_dir) #TODO: parametrize this number [0] so the right run is selected
+
+    # 6. get filenames in chronological sequence and write them to gif
+    dirs = [x[0] for x in os.walk(dir)]
+    dirs = sorted(dirs, key=str.lower)[1:] #sort chronologically and remove parent dir from list
+
+    with imageio.get_writer(args.gif_name+'.gif', mode='I') as writer: #https://imageio.readthedocs.io/en/stable/index.html#
+        for epoch in dirs: #extract the right file from each epoch
+            for _, _, files in os.walk(epoch): #
+                for name in files:
+                    if args.representation in name and args.step in name and "png" in name:
+                        image = imageio.imread(epoch+"/"+name)
+                        writer.append_data(image)
+
+    # 7. cleanup the downloaded artifacts from client file system
+    if args.cleanup:
+        shutil.rmtree(local_dir) #delete the files downloaded from mlflow
+
+elif args.output == "train_vs_val_loss":
+    train_loss = client.get_metric_history(runs["run_id"][0], "train_loss") #returns a list of metric entities https://www.mlflow.org/docs/latest/_modules/mlflow/entities/metric.html
+    val_loss = client.get_metric_history(runs["run_id"][0], "val_loss") #TODO: parametrize this number [0] so the right run is selected
+    train_loss = sorted(train_loss, key=lambda m: m.step) #sort the metric objects in list according to step property
+    val_loss = sorted(val_loss, key=lambda m: m.step)
+    plt.figure()
+    for m_train, m_val in zip(train_loss, val_loss):
+        plt.scatter(m_train.value, m_val.value, alpha=1/(m_train.step+1), color='blue')
+    plt.savefig("scatter.png") #TODO: parametrize filename

models/classifier.py

@@ -0,0 +1,281 @@
+import os
+from collections import defaultdict
+
+import torch
+import torch.optim
+from torchvision.models import resnet18
+from torchvision.utils import make_grid, save_image
+import torch.nn.functional as F
+
+import pytorch_lightning as pl
+
+import mlflow.pytorch
+
+
+def resnet_model(model=resnet18, pretrained=True, in_channels=3, fc_out_features=2):
+    resnet = model(pretrained=pretrained)
+    # if not pretrained:  # TODO: add case for in_channels=4
+    #     resnet.conv1 = torch.nn.Conv2d(channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
+    resnet.fc = torch.nn.Linear(in_features=512, out_features=fc_out_features, bias=True)
+    return resnet
+
+
+class LitModel(pl.LightningModule):
+
+    def __init__(self,
+                 classifier,
+                 loss,
+                 lr=1e-3,
+                 weight_decay=0,
+                 loss_aux=None,
+                 adv_training=False,
+                 metrics=None,
+                 processor=None,
+                 augmentation=None,
+                 is_segmentation_task=False,
+                 augmentation_on_eval=False,
+                 metrics_on_training=True,
+                 freeze_classifier=False,
+                 freeze_processor=False,
+                 ):
+        super().__init__()
+
+        self.classifier = classifier
+        self.processor = processor
+
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.loss_fn = loss
+        self.loss_aux_fn = loss_aux
+        self.adv_training = adv_training
+        self.metrics = metrics
+        self.augmentation = augmentation
+        self.is_segmentation_task = is_segmentation_task
+        self.augmentation_on_eval = augmentation_on_eval
+        self.metrics_on_training = metrics_on_training
+
+        self.freeze_classifier = freeze_classifier
+        self.freeze_processor = freeze_processor
+
+        if freeze_classifier:
+            pl.LightningModule.freeze(self.classifier)
+        if freeze_processor:
+            pl.LightningModule.freeze(self.processor)
+
+    def forward(self, x):
+        x = self.processor(x)
+        apply_augmentation_step = self.training or self.augmentation_on_eval
+        if self.augmentation is not None and apply_augmentation_step:
+            x = self.augmentation(x, retain_state=self.is_segmentation_task)
+        x = self.classifier(x)
+        return x
+
+    def update_step(self, batch, step_name):
+        x, y = batch
+        # debug(self.processor)
+        # debug(self.processor.parameters())
+        # debug.pause()
+        # print('type', type(self.processor).__name__)
+
+        logits = self(x)
+
+        apply_augmentation_mask = self.is_segmentation_task and (self.training or self.augmentation_on_eval)
+        if self.augmentation is not None and apply_augmentation_mask:
+            y = self.augmentation(y, mask_transform=True).contiguous()
+
+        loss = self.loss_fn(logits, y)
+
+        if self.loss_aux_fn is not None:
+            loss_aux = self.loss_aux_fn(x)
+            loss += loss_aux
+
+        self.log(f'{step_name}_loss', loss, on_step=False, on_epoch=True)
+        if self.loss_aux_fn is not None:
+            self.log(f'{step_name}_loss_aux', loss_aux, on_step=False, on_epoch=True)
+
+        if self.is_segmentation_task:
+            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:
+                metric_name = metric.__name__ if hasattr(metric, '__name__') else type(metric).__name__
+                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')                
+                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')
+
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        return self.update_step(batch, 'train')
+
+    def validation_step(self, batch, batch_idx):
+        return self.update_step(batch, 'val')
+
+    def test_step(self, batch, batch_idx):
+        return self.update_step(batch, 'test')
+
+    def train(self, mode=True):
+        self.training = mode
+        # self.processor.train(False)
+        self.processor.train(mode=mode and not self.freeze_processor)
+        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()
+        items.pop('v_num')
+        return items
+
+
+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):
+        super().__init__()
+        self.data_loader = data_loader
+
+        self.track_every_epoch = track_every_epoch
+
+        self.track_processing = track_processing
+        self.track_gradients = track_gradients
+        self.track_predictions = track_predictions
+        self.save_tensors = save_tensors
+
+    def callback_track_images(self, trainer, save_loc):
+        track_images(trainer.model,
+                     self.data_loader,
+                     track_processing=self.track_processing,
+                     track_gradients=self.track_gradients,
+                     track_predictions=self.track_predictions,
+                     save_tensors=self.save_tensors,
+                     save_loc=save_loc,
+                     )
+
+    def on_fit_end(self, trainer, pl_module):
+        if not self.track_every_epoch:
+            save_loc = 'results'
+            self.callback_track_images(trainer, 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)
+
+
+from utils.debug import debug
+
+
+# @debug
+def log_tensor(batch, path, save_tensors=True, nrow=8):
+    if save_tensors:
+        torch.save(batch, path)
+        mlflow.log_artifact(path, os.path.dirname(path))
+
+    img_path = path.replace('.pt', '.png')
+    split = img_path.split('/')
+    img_path = '/'.join(split[:-1]) + '/img_' + split[-1]  # insert 'img_'; make it easier to find in mlflow
+
+    grid = make_grid(batch, nrow=nrow).squeeze()
+    save_image(grid, img_path)
+    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'):
+
+    device = model.device
+    processor = model.processor
+    classifier = model.classifier
+
+    if not hasattr(processor, 'stages'):    # 'static' or 'none' pipeline
+        return
+
+    os.makedirs(save_loc, exist_ok=True)
+
+    # TODO: implement track_predictions
+
+    # inputs_full = []
+    labels_full = []
+    logits_full = []
+    stages_full = defaultdict(list)
+    grads_full = defaultdict(list)
+
+    for inputs, labels in data_loader:
+
+        inputs, labels = inputs.to(device), labels.to(device)
+        inputs.requires_grad = True
+
+        processed_rgb = processor(inputs)
+
+        if track_gradients or track_predictions:
+            logits = classifier(processed_rgb)
+
+            # NOTE: should zero grads for good measure
+            loss = model.loss_fn(logits, labels)
+            loss.backward()
+
+            if track_predictions:
+                labels_full.append(labels.cpu().detach())
+                logits_full.append(logits.cpu().detach())
+        # inputs_full.append(inputs.cpu().detach())
+
+        for stage, batch in processor.stages.items():
+            stages_full[stage].append(batch.cpu().detach())
+            if track_gradients:
+                grads_full[stage].append(batch.grad.cpu().detach())
+
+    with torch.no_grad():
+
+        stages = stages_full
+        grads = grads_full
+
+        if track_processing:
+            for stage, batch in stages_full.items():
+                stages[stage] = torch.cat(batch)
+
+        if track_gradients:
+            for stage, batch in grads_full.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_gradients:
+                batch_grad = grads[stage_name]
+                batch_grad = batch_grad.abs()
+                batch_grad = (batch_grad - batch_grad.min()) / (batch_grad.max() - batch_grad.min())
+                log_tensor(batch_grad, os.path.join(
+                    save_loc, f'gradients_{stage_nr}_{stage_name}.pt'), save_tensors)
+
+        # inputs = torch.cat(inputs_full)
+
+        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)
+            #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)
+            log_tensor(predictions, os.path.join(save_loc, f'preds.pt'), save_tensors)

perturbed-environment.yml

@@ -0,0 +1,363 @@
+name: perturbed
+channels:
+  - defaults
+dependencies:
+  - _ipyw_jlab_nb_ext_conf=0.1.0=py37_0
+  - _libgcc_mutex=0.1=main
+  - alabaster=0.7.12=py37_0
+  - anaconda=2019.10=py37_0
+  - anaconda-client=1.7.2=py37_0
+  - anaconda-navigator=1.9.7=py37_0
+  - anaconda-project=0.8.3=py_0
+  - asn1crypto=1.0.1=py37_0
+  - astroid=2.3.1=py37_0
+  - astropy=3.2.2=py37h7b6447c_0
+  - atomicwrites=1.3.0=py37_1
+  - attrs=19.2.0=py_0
+  - babel=2.7.0=py_0
+  - backcall=0.1.0=py37_0
+  - backports=1.0=py_2
+  - backports.functools_lru_cache=1.5=py_2
+  - backports.os=0.1.1=py37_0
+  - backports.shutil_get_terminal_size=1.0.0=py37_2
+  - backports.tempfile=1.0=py_1
+  - backports.weakref=1.0.post1=py_1
+  - beautifulsoup4=4.8.0=py37_0
+  - bitarray=1.0.1=py37h7b6447c_0
+  - bkcharts=0.2=py37_0
+  - blas=1.0=mkl
+  - bleach=3.1.0=py37_0
+  - blosc=1.16.3=hd408876_0
+  - bokeh=1.3.4=py37_0
+  - boto=2.49.0=py37_0
+  - bottleneck=1.2.1=py37h035aef0_1
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2019.8.28=0
+  - cairo=1.14.12=h8948797_3
+  - certifi=2019.9.11=py37_0
+  - cffi=1.12.3=py37h2e261b9_0
+  - chardet=3.0.4=py37_1003
+  - click=7.0=py37_0
+  - cloudpickle=1.2.2=py_0
+  - clyent=1.2.2=py37_1
+  - colorama=0.4.1=py37_0
+  - conda-package-handling=1.6.0=py37h7b6447c_0
+  - conda-verify=3.4.2=py_1
+  - contextlib2=0.6.0=py_0
+  - cryptography=2.7=py37h1ba5d50_0
+  - curl=7.65.3=hbc83047_0
+  - cycler=0.10.0=py37_0
+  - cython=0.29.13=py37he6710b0_0
+  - cytoolz=0.10.0=py37h7b6447c_0
+  - dask=2.5.2=py_0
+  - dask-core=2.5.2=py_0
+  - dbus=1.13.6=h746ee38_0
+  - decorator=4.4.0=py37_1
+  - defusedxml=0.6.0=py_0
+  - distributed=2.5.2=py_0
+  - docutils=0.15.2=py37_0
+  - entrypoints=0.3=py37_0
+  - et_xmlfile=1.0.1=py37_0
+  - expat=2.2.6=he6710b0_0
+  - fastcache=1.1.0=py37h7b6447c_0
+  - filelock=3.0.12=py_0
+  - flask=1.1.1=py_0
+  - fontconfig=2.13.0=h9420a91_0
+  - freetype=2.9.1=h8a8886c_1
+  - fribidi=1.0.5=h7b6447c_0
+  - future=0.17.1=py37_0
+  - get_terminal_size=1.0.0=haa9412d_0
+  - gevent=1.4.0=py37h7b6447c_0
+  - glib=2.56.2=hd408876_0
+  - glob2=0.7=py_0
+  - gmp=6.1.2=h6c8ec71_1
+  - gmpy2=2.0.8=py37h10f8cd9_2
+  - graphite2=1.3.13=h23475e2_0
+  - greenlet=0.4.15=py37h7b6447c_0
+  - gst-plugins-base=1.14.0=hbbd80ab_1
+  - gstreamer=1.14.0=hb453b48_1
+  - h5py=2.9.0=py37h7918eee_0
+  - harfbuzz=1.8.8=hffaf4a1_0
+  - hdf5=1.10.4=hb1b8bf9_0
+  - heapdict=1.0.1=py_0
+  - html5lib=1.0.1=py37_0
+  - icu=58.2=h9c2bf20_1
+  - idna=2.8=py37_0
+  - imageio=2.6.0=py37_0
+  - imagesize=1.1.0=py37_0
+  - intel-openmp=2019.4=243
+  - ipykernel=5.1.2=py37h39e3cac_0
+  - ipython=7.8.0=py37h39e3cac_0
+  - ipython_genutils=0.2.0=py37_0
+  - ipywidgets=7.5.1=py_0
+  - isort=4.3.21=py37_0
+  - itsdangerous=1.1.0=py37_0
+  - jbig=2.1=hdba287a_0
+  - jdcal=1.4.1=py_0
+  - jedi=0.15.1=py37_0
+  - jeepney=0.4.1=py_0
+  - jinja2=2.10.3=py_0
+  - joblib=0.13.2=py37_0
+  - jpeg=9b=h024ee3a_2
+  - json5=0.8.5=py_0
+  - jsonschema=3.0.2=py37_0
+  - jupyter=1.0.0=py37_7
+  - jupyter_client=5.3.3=py37_1
+  - jupyter_console=6.0.0=py37_0
+  - jupyter_core=4.5.0=py_0
+  - jupyterlab=1.1.4=pyhf63ae98_0
+  - jupyterlab_server=1.0.6=py_0
+  - keyring=18.0.0=py37_0
+  - kiwisolver=1.1.0=py37he6710b0_0
+  - krb5=1.16.1=h173b8e3_7
+  - lazy-object-proxy=1.4.2=py37h7b6447c_0
+  - libarchive=3.3.3=h5d8350f_5
+  - libcurl=7.65.3=h20c2e04_0
+  - libedit=3.1.20181209=hc058e9b_0
+  - libffi=3.2.1=hd88cf55_4
+  - libgcc-ng=9.1.0=hdf63c60_0
+  - libgfortran-ng=7.3.0=hdf63c60_0
+  - liblief=0.9.0=h7725739_2
+  - libpng=1.6.37=hbc83047_0
+  - libsodium=1.0.16=h1bed415_0
+  - libssh2=1.8.2=h1ba5d50_0
+  - libstdcxx-ng=9.1.0=hdf63c60_0
+  - libtiff=4.0.10=h2733197_2
+  - libtool=2.4.6=h7b6447c_5
+  - libuuid=1.0.3=h1bed415_2
+  - libxcb=1.13=h1bed415_1
+  - libxml2=2.9.9=hea5a465_1
+  - libxslt=1.1.33=h7d1a2b0_0
+  - llvmlite=0.29.0=py37hd408876_0
+  - locket=0.2.0=py37_1
+  - lxml=4.4.1=py37hefd8a0e_0
+  - lz4-c=1.8.1.2=h14c3975_0
+  - lzo=2.10=h49e0be7_2
+  - markupsafe=1.1.1=py37h7b6447c_0
+  - matplotlib=3.1.1=py37h5429711_0
+  - mccabe=0.6.1=py37_1
+  - mistune=0.8.4=py37h7b6447c_0
+  - mkl=2019.4=243
+  - mkl-service=2.3.0=py37he904b0f_0
+  - mkl_fft=1.0.14=py37ha843d7b_0
+  - mkl_random=1.1.0=py37hd6b4f25_0
+  - mock=3.0.5=py37_0
+  - more-itertools=7.2.0=py37_0
+  - mpc=1.1.0=h10f8cd9_1
+  - mpfr=4.0.1=hdf1c602_3
+  - mpmath=1.1.0=py37_0
+  - msgpack-python=0.6.1=py37hfd86e86_1
+  - multipledispatch=0.6.0=py37_0
+  - navigator-updater=0.2.1=py37_0
+  - nbconvert=5.6.0=py37_1
+  - nbformat=4.4.0=py37_0
+  - ncurses=6.1=he6710b0_1
+  - networkx=2.3=py_0
+  - nltk=3.4.5=py37_0
+  - nose=1.3.7=py37_2
+  - notebook=6.0.1=py37_0
+  - numba=0.45.1=py37h962f231_0
+  - numexpr=2.7.0=py37h9e4a6bb_0
+  - numpy=1.17.2=py37haad9e8e_0
+  - numpy-base=1.17.2=py37hde5b4d6_0
+  - numpydoc=0.9.1=py_0
+  - olefile=0.46=py37_0
+  - openpyxl=3.0.0=py_0
+  - openssl=1.1.1d=h7b6447c_2
+  - packaging=19.2=py_0
+  - pandoc=2.2.3.2=0
+  - pandocfilters=1.4.2=py37_1
+  - pango=1.42.4=h049681c_0
+  - parso=0.5.1=py_0
+  - partd=1.0.0=py_0
+  - patchelf=0.9=he6710b0_3
+  - path.py=12.0.1=py_0
+  - pathlib2=2.3.5=py37_0
+  - patsy=0.5.1=py37_0
+  - pcre=8.43=he6710b0_0
+  - pep8=1.7.1=py37_0
+  - pexpect=4.7.0=py37_0
+  - pickleshare=0.7.5=py37_0
+  - pip=19.2.3=py37_0
+  - pixman=0.38.0=h7b6447c_0
+  - pkginfo=1.5.0.1=py37_0
+  - pluggy=0.13.0=py37_0
+  - ply=3.11=py37_0
+  - prometheus_client=0.7.1=py_0
+  - prompt_toolkit=2.0.10=py_0
+  - psutil=5.6.3=py37h7b6447c_0
+  - ptyprocess=0.6.0=py37_0
+  - py=1.8.0=py37_0
+  - py-lief=0.9.0=py37h7725739_2
+  - pycodestyle=2.5.0=py37_0
+  - pycosat=0.6.3=py37h14c3975_0
+  - pycparser=2.19=py37_0
+  - pycrypto=2.6.1=py37h14c3975_9
+  - pycurl=7.43.0.3=py37h1ba5d50_0
+  - pyflakes=2.1.1=py37_0
+  - pygments=2.4.2=py_0
+  - pylint=2.4.2=py37_0
+  - pyodbc=4.0.27=py37he6710b0_0
+  - pyopenssl=19.0.0=py37_0
+  - pyparsing=2.4.2=py_0
+  - pyqt=5.9.2=py37h05f1152_2
+  - pyrsistent=0.15.4=py37h7b6447c_0
+  - pysocks=1.7.1=py37_0
+  - pytables=3.5.2=py37h71ec239_1
+  - pytest=5.2.1=py37_0
+  - pytest-arraydiff=0.3=py37h39e3cac_0
+  - pytest-astropy=0.5.0=py37_0
+  - pytest-doctestplus=0.4.0=py_0
+  - pytest-openfiles=0.4.0=py_0
+  - pytest-remotedata=0.3.2=py37_0
+  - python=3.7.4=h265db76_1
+  - python-dateutil=2.8.0=py37_0
+  - python-libarchive-c=2.8=py37_13
+  - pytz=2019.3=py_0
+  - pyyaml=5.1.2=py37h7b6447c_0
+  - pyzmq=18.1.0=py37he6710b0_0
+  - qt=5.9.7=h5867ecd_1
+  - qtawesome=0.6.0=py_0
+  - qtconsole=4.5.5=py_0
+  - qtpy=1.9.0=py_0
+  - readline=7.0=h7b6447c_5
+  - requests=2.22.0=py37_0
+  - ripgrep=0.10.0=hc07d326_0
+  - rope=0.14.0=py_0
+  - ruamel_yaml=0.15.46=py37h14c3975_0
+  - scikit-learn=0.21.3=py37hd81dba3_0
+  - scipy=1.3.1=py37h7c811a0_0
+  - seaborn=0.9.0=py37_0
+  - secretstorage=3.1.1=py37_0
+  - send2trash=1.5.0=py37_0
+  - setuptools=41.4.0=py37_0
+  - simplegeneric=0.8.1=py37_2
+  - singledispatch=3.4.0.3=py37_0
+  - sip=4.19.8=py37hf484d3e_0
+  - six=1.12.0=py37_0
+  - snappy=1.1.7=hbae5bb6_3
+  - snowballstemmer=2.0.0=py_0
+  - sortedcollections=1.1.2=py37_0
+  - sortedcontainers=2.1.0=py37_0
+  - soupsieve=1.9.3=py37_0
+  - sphinx=2.2.0=py_0
+  - sphinxcontrib=1.0=py37_1
+  - sphinxcontrib-applehelp=1.0.1=py_0
+  - sphinxcontrib-devhelp=1.0.1=py_0
+  - sphinxcontrib-htmlhelp=1.0.2=py_0
+  - sphinxcontrib-jsmath=1.0.1=py_0
+  - sphinxcontrib-qthelp=1.0.2=py_0
+  - sphinxcontrib-serializinghtml=1.1.3=py_0
+  - sphinxcontrib-websupport=1.1.2=py_0
+  - spyder=3.3.6=py37_0
+  - spyder-kernels=0.5.2=py37_0
+  - sqlalchemy=1.3.9=py37h7b6447c_0
+  - sqlite=3.30.0=h7b6447c_0
+  - statsmodels=0.10.1=py37hdd07704_0
+  - sympy=1.4=py37_0
+  - tbb=2019.4=hfd86e86_0
+  - tblib=1.4.0=py_0
+  - terminado=0.8.2=py37_0
+  - testpath=0.4.2=py37_0
+  - tk=8.6.8=hbc83047_0
+  - toolz=0.10.0=py_0
+  - tornado=6.0.3=py37h7b6447c_0
+  - traitlets=4.3.3=py37_0
+  - unicodecsv=0.14.1=py37_0
+  - unixodbc=2.3.7=h14c3975_0
+  - wcwidth=0.1.7=py37_0
+  - webencodings=0.5.1=py37_1
+  - werkzeug=0.16.0=py_0
+  - wheel=0.33.6=py37_0
+  - widgetsnbextension=3.5.1=py37_0
+  - wrapt=1.11.2=py37h7b6447c_0
+  - wurlitzer=1.0.3=py37_0
+  - xlrd=1.2.0=py37_0
+  - xlsxwriter=1.2.1=py_0
+  - xlwt=1.3.0=py37_0
+  - xz=5.2.4=h14c3975_4
+  - yaml=0.1.7=had09818_2
+  - zeromq=4.3.1=he6710b0_3
+  - zict=1.0.0=py_0
+  - zipp=0.6.0=py_0
+  - zlib=1.2.11=h7b6447c_3
+  - zstd=1.3.7=h0b5b093_0
+  - pip:
+    - absl-py==0.12.0
+    - aiohttp==3.7.4.post0
+    - albumentations==0.5.2
+    - alembic==1.4.1
+    - arrow==0.17.0
+    - async-timeout==3.0.1
+    - b2sdk==1.4.0
+    - boto3==1.17.36
+    - botocore==1.20.36
+    - cachetools==4.2.1
+    - colour-demosaicing==0.1.6
+    - colour-science==0.3.16
+    - configparser==5.0.0
+    - databricks-cli==0.10.0
+    - docker==4.2.0
+    - docopt==0.6.2
+    - efficientnet-pytorch==0.6.3
+    - fsspec==0.8.7
+    - funcsigs==1.0.2
+    - gitdb==4.0.4
+    - gitpython==3.1.1
+    - google-auth==1.28.0
+    - google-auth-oauthlib==0.4.3
+    - gorilla==0.3.0
+    - grpcio==1.36.1
+    - gunicorn==20.0.4
+    - imgaug==0.4.0
+    - importlib-metadata==3.7.3
+    - jmespath==0.10.0
+    - logfury==0.1.2
+    - mako==1.1.2
+    - markdown==3.3.4
+    - mlflow==1.14.1
+    - multidict==5.1.0
+    - munch==2.5.0
+    - oauthlib==3.1.0
+    - opencv-python==4.5.1.48
+    - opencv-python-headless==4.5.1.48
+    - pandas==1.2.3
+    - pillow==8.1.2
+    - pipreqs==0.4.10
+    - plotly==4.14.3
+    - pretrainedmodels==0.7.4
+    - prettytable==2.1.0
+    - prometheus-flask-exporter==0.13.0
+    - protobuf==3.11.3
+    - pyasn1==0.4.8
+    - pyasn1-modules==0.2.8
+    - python-editor==1.0.4
+    - pytorch-lightning==1.2.5
+    - pywavelets==1.1.1
+    - querystring-parser==1.2.4
+    - rawpy==0.16.0
+    - requests-oauthlib==1.3.0
+    - retrying==1.3.3
+    - rsa==4.7.2
+    - s3transfer==0.3.6
+    - scikit-image==0.18.1
+    - segmentation-models-pytorch==0.1.3
+    - shapely==1.7.1
+    - simplejson==3.17.0
+    - smmap==3.0.2
+    - sqlparse==0.3.1
+    - tabulate==0.8.7
+    - tensorboard==2.4.1
+    - tensorboard-plugin-wit==1.8.0
+    - tifffile==2021.3.17
+    - timm==0.3.2
+    - torch==1.8.0
+    - torchmetrics==0.2.0
+    - torchvision==0.9.0
+    - tqdm==4.59.0
+    - typing-extensions==3.7.4.3
+    - urllib3==1.25.11
+    - websocket-client==0.57.0
+    - yarg==0.1.9
+    - yarl==1.6.3
+prefix: /home/nobis/anaconda3/envs/perturbed

processingpipeline/numpy_static_pipeline_show.ipynb

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

processingpipeline/pipeline.py

@@ -0,0 +1,329 @@
+"""
+Raw Image Pipeline
+"""
+__author__ = "Marco Aversa"
+
+import numpy as np
+
+from rawpy import *  # XXX: no * imports!
+from scipy import ndimage
+from scipy import fftpack
+from scipy.signal import convolve2d
+
+from skimage.filters import unsharp_mask
+from skimage.color import rgb2yuv, yuv2rgb, rgb2hsv, hsv2rgb
+from skimage.restoration import denoise_tv_chambolle, denoise_tv_bregman, denoise_nl_means, denoise_bilateral, denoise_wavelet, estimate_sigma
+
+import matplotlib.pyplot as plt
+
+from colour_demosaicing import (demosaicing_CFA_Bayer_bilinear,
+                                demosaicing_CFA_Bayer_Malvar2004,
+                                demosaicing_CFA_Bayer_Menon2007)
+
+import torch
+import numpy as np
+
+from utils.dataset import Subset
+from torch.utils.data import DataLoader
+
+from colour_demosaicing import (demosaicing_CFA_Bayer_bilinear,
+                                demosaicing_CFA_Bayer_Malvar2004,
+                                demosaicing_CFA_Bayer_Menon2007)
+
+import matplotlib.pyplot as plt
+
+
+class RawProcessingPipeline(object):
+
+    """Applies the raw-processing pipeline from pipeline.py"""
+
+    def __init__(self, camera_parameters, debayer='bilinear', sharpening='unsharp_masking', denoising='gaussian'):
+        '''
+        Args:
+            camera_parameters (tuple): (black_level, white_balance, colour_matrix)
+            debayer (str): specifies the algorithm used as debayer; choose from {'bilinear','malvar2004','menon2007'}
+            sharpening (str): specifies the algorithm used for sharpening; choose from {'sharpening_filter','unsharp_masking'}
+            denoising (str): specifies the algorithm used for denoising; choose from choose from {'gaussian_denoising','median_denoising','fft_denoising'}        
+        '''
+
+        self.camera_parameters = camera_parameters
+
+        self.debayer = debayer
+        self.sharpening = sharpening
+        self.denoising = denoising
+
+    def __call__(self, img):
+        """
+        Args:
+            img (ndarry of dtype float.32): image of size (H,W)
+        return:
+            img (tensor of dtype float): image of size (3,H,W)
+        """
+        black_level, white_balance, colour_matrix = self.camera_parameters
+        img = processing(img, black_level, white_balance, colour_matrix,
+                         debayer=self.debayer, sharpening=self.sharpening, denoising=self.denoising)
+        img = img.transpose(2, 0, 1)
+
+        return torch.Tensor(img)
+
+
+def processing(img, black_level, white_balance, colour_matrix, debayer="bilinear", sharpening="unsharp_masking",
+               sharp_radius=1.0, sharp_amount=1.0, denoising="median_filter", median_kernel_size=3,
+               gaussian_sigma=0.5, fft_fraction=0.3, weight_chambolle=0.01, weight_bregman=100,
+               sigma_bilateral=0.6, gamma=2.2, bits=16):
+    """Apply pipeline on a raw image
+
+       Args:
+           rawImg (ndarray): raw image
+           debayer (str): debayer algorithm
+           white_balance (None, ndarray): white balance array (if None it will take the default camera white balance array)
+           colour_matrix (None, ndarray): colour matrix (if None it will take the default camera colour matrix) - Size: 3x3
+           gamma (float): exponent for the non linear gamma correction.
+
+       Returns:
+           img (ndarray): post-processed image
+
+    """
+
+    # Remove Black Level
+    img = remove_blacklv(img, black_level)
+
+    # Apply demosaicing - We don't have access to these 3 functions
+    if debayer == "bilinear":
+        img = demosaicing_CFA_Bayer_bilinear(img)
+    if debayer == "malvar2004":
+        img = demosaicing_CFA_Bayer_Malvar2004(img)
+    if debayer == "menon2007":
+        img = demosaicing_CFA_Bayer_Menon2007(img)
+
+    # White Balance Correction
+
+    # Sunny images white balance array -> 2<r<2.8, g=1.0, 1.3<b<1.6
+    # Tungsten images white balance array -> 1.3<r<1.7, g=1.0, 2.2<b<2.8
+    # Shade images white balance array -> 2.4<r<3.2, g=1.0, 1.1<b<1.3
+
+    img = wb_correction(img, white_balance)
+
+    # Colour Correction
+    img = colour_correction(img, colour_matrix)
+
+    # Sharpening
+    if sharpening == "sharpening_filter":  # Fixed sharpening
+        img = sharpening_filter(img)
+    if sharpening == "unsharp_masking":  # Higher is radius and amount, higher is the sharpening
+        img = unsharp_masking(img, radius=sharp_radius, amount=sharp_amount, multichannel=True)
+
+    # Denoising
+    if denoising == "median_denoising":
+        img = median_denoising(img, size=median_kernel_size)
+    if denoising == "gaussian_denoising":
+        img = gaussian_denoising(img, sigma=gaussian_sigma)
+    if denoising == "fft_denoising":  # fft_fraction = [0.0001,0.5]
+        img = fft_denoising(img, keep_fraction=fft_fraction, row_cut=False, column_cut=True)
+
+    # We don't have access to these 3 functions
+    if denoising == "tv_chambolle":  # lower is weight, less is the denoising
+        img = denoise_tv_chambolle(img, weight=weight_chambolle, eps=0.0002, n_iter_max=200, multichannel=True)
+    if denoising == "tv_bregman":  # lower is weight, more is the denoising
+        img = denoise_tv_bregman(img, weight=weight_bregman, max_iter=100,
+                                 eps=0.001, isotropic=True, multichannel=True)
+#     if denoising == "wavelet":
+#         img = denoise_wavelet(img.copy(), sigma=None, wavelet='db1', mode='soft', wavelet_levels=None, multichannel=True,
+#                 convert2ycbcr=False, method='BayesShrink', rescale_sigma=True)
+    if denoising == "bilateral":  # higher is sigma_spatial, more is the denoising
+        img = denoise_bilateral(img, win_size=None, sigma_color=None, sigma_spatial=sigma_bilateral,
+                                bins=10000, mode='constant', cval=0, multichannel=True)
+
+    # Gamma Correction
+    img = np.clip(img, 0, 1)
+    img = adjust_gamma(img, gamma=gamma)
+
+    return img
+
+
+def get_camera_parameters(rawpyImg):
+    black_level = rawpyImg.black_level_per_channel
+    white_balance = rawpyImg.camera_whitebalance[:3]
+    colour_matrix = rawpyImg.color_matrix[:, :3].flatten().tolist()
+
+    return black_level, white_balance, colour_matrix
+
+
+def remove_blacklv(rawImg, black_level):
+    rawImg[0::2, 0::2] -= black_level[0]  # R
+    rawImg[0::2, 1::2] -= black_level[1]  # G
+    rawImg[1::2, 0::2] -= black_level[2]  # G
+    rawImg[1::2, 1::2] -= black_level[3]  # B
+
+    return rawImg
+
+
+def wb_correction(img, white_balance):
+    return img * white_balance
+
+
+def colour_correction(img, colour_matrix):
+    colour_matrix = np.array(colour_matrix).reshape(3, 3)
+    return np.einsum('ijk,lk->ijl', img, colour_matrix)
+
+
+def unsharp_masking(img, radius=1.0, amount=1.0,
+                    multichannel=False, preserve_range=True):
+
+    img = rgb2yuv(img)
+    img[:, :, 0] = unsharp_mask(img[:, :, 0], radius=radius, amount=amount,
+                                multichannel=multichannel, preserve_range=preserve_range)
+    img = yuv2rgb(img)
+    return img
+
+
+def sharpening_filter(image, iterations=1, kernel=np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])):
+
+    # https://towardsdatascience.com/image-processing-with-python-blurring-and-sharpening-for-beginners-3bcebec0583a
+
+    img_yuv = rgb2yuv(image)
+
+    for i in range(iterations):
+        img_yuv[:, :, 0] = convolve2d(img_yuv[:, :, 0], kernel, 'same', boundary='fill', fillvalue=0)
+
+    final_image = yuv2rgb(img_yuv)
+
+    return final_image
+
+
+def median_denoising(img, size=3):
+
+    img = rgb2yuv(img)
+    img[:, :, 0] = ndimage.median_filter(img[:, :, 0], size)
+    img = yuv2rgb(img)
+
+    return img
+
+
+def gaussian_denoising(img, sigma=0.5):
+
+    img = rgb2yuv(img)
+    img[:, :, 0] = ndimage.gaussian_filter(img[:, :, 0], sigma)
+    img = yuv2rgb(img)
+
+    return img
+
+
+def fft_denoising(img, keep_fraction=0.3, row_cut=False, column_cut=True):
+    """ keep_fraction = 0.5 --> same image as input
+        keep_fraction --> 0 --> remove all details """
+#   http://scipy-lectures.org/intro/scipy/auto_examples/solutions/plot_fft_image_denoise.html
+
+    im_fft = fftpack.fft2(img)
+
+    # Call ff a copy of the original transform. Numpy arrays have a copy
+    # method for this purpose.
+    im_fft2 = im_fft
+
+    # Set r and c to be the number of rows and columns of the array.
+    r, c, _ = im_fft2.shape
+
+    # Set to zero all rows with indices between r*keep_fraction and r*(1-keep_fraction):
+    if row_cut == True:
+        im_fft2[int(r * keep_fraction):int(r * (1 - keep_fraction))] = 0
+
+    # Similarly with the columns:
+    if column_cut == True:
+        im_fft2[:, int(c * keep_fraction):int(c * (1 - keep_fraction))] = 0
+
+    # Reconstruct the denoised image from the filtered spectrum, keep only the
+    # real part for display.
+    im_new = fftpack.ifft2(im_fft2).real
+
+    return im_new
+
+
+def adjust_gamma(img, gamma=1.0):
+    invGamma = 1.0 / gamma
+    img = (img ** invGamma)
+    return img
+
+
+def show_img(img, title="no_title", size=12, histo=True, bins=300, bits=16, x_range=-1):
+    """Plot image and its histogram
+
+       Args:
+           img (ndarray): image to plot
+           title (str): title of the plot
+           histo (bool): True - Plot histrograms per channel of the image. False - Plot the curve of histogram in a continue way
+           bins (int): number of bins of the histogram
+           size (int): figure size
+           bits (int): number of bits per pixel in the ndarray
+           x_range (list): maximum x range of the histogram (if -1 it will be take all x values)
+    """
+    shape = img.shape
+
+    fig = plt.figure(figsize=(size, size))
+
+    # show original image
+    fig.add_subplot(221)
+    if len(shape) > 2 and img.max() > 255:
+        img_to_show = (img.copy() * 255. / (2**bits - 1)).astype(int)
+    else:
+        img_to_show = img.copy().astype(int)
+    plt.imshow(img_to_show)
+    if title != "no_title":
+        plt.title(title)
+
+    fig.add_subplot(222)
+
+    if len(shape) > 2:
+        if histo == True:
+            plt.hist(img[:, :, 0].flatten(), bins=bins, label="Channel1", color="red", alpha=0.5)
+            plt.hist(img[:, :, 1].flatten(), bins=bins, label="Channel2", color="green", alpha=0.5)
+            plt.hist(img[:, :, 2].flatten(), bins=bins, label="Channel3", color="blue", alpha=0.5)
+            if x_range != -1:
+                plt.xlim([x_range[0], x_range[1]])
+        else:
+            h1, b1 = np.histogram(img[:, :, 0].flatten(), bins=bins)
+            h2, b2 = np.histogram(img[:, :, 1].flatten(), bins=bins)
+            h3, b3 = np.histogram(img[:, :, 2].flatten(), bins=bins)
+            plt.plot(b1[:-1], h1, label="Channel1", color="red", alpha=0.5)
+            plt.plot(b2[:-1], h2, label="Channel2", color="green", alpha=0.5)
+            plt.plot(b3[:-1], h3, label="Channel3", color="blue", alpha=0.5)
+
+        plt.legend()
+    else:
+        if histo == True:
+            plt.hist(img.flatten(), bins=bins)
+            if x_range != -1:
+                plt.xlim([x_range[0], x_range[1]])
+        else:
+            h, b = np.histogram(img.flatten(), bins=bins)
+            plt.plot(b[:-1], h)
+
+    plt.xlabel("Intensities")
+    plt.ylabel("Counts")
+
+    plt.show()
+
+
+def get_statistics(dataset, train_indices, transform=None):
+    """Calculates the mean and the standard deviation of a given sub train set of dataset
+
+    Args:
+        dataset (Subset of DroneDataset): 
+        train_indices (tensor): indicies correponding to a subset of the dataset
+        transform (Compose): list of transformations compatible with Compose to be applied before calculations
+    return:
+        mean (tensor of dtype float): size (C,1,1)
+        std (tensor of dtype float): size (C,1,1)
+    """
+
+    trainset = Subset(dataset, indices=train_indices, transform=transform)
+    dataloader = DataLoader(trainset, batch_size=len(trainset), shuffle=False)
+    dataiter = iter(dataloader)
+
+    images, labels = dataiter.next()
+
+    if len(images.shape) == 3:
+        mean, std = torch.mean(images, axis=(0, 1, 2)), torch.std(images, axis=(0, 1, 2))
+        return mean, std
+    else:
+        mean, std = torch.mean(images, axis=(0, 2, 3))[:, None, None], torch.std(images, axis=(0, 2, 3))[:, None, None]
+        return mean, std

processingpipeline/torch_pipeline.py

@@ -0,0 +1,313 @@
+import os
+from numpy.lib.function_base import interp
+import torch
+import torch.nn as nn
+if not os.path.exists('README.md'):
+    os.chdir('..')
+
+from processingpipeline.pipeline import processing as default_processing
+from utils.base import np2torch, torch2np
+
+import segmentation_models_pytorch as smp
+
+from utils.debug import debug
+
+K_G = torch.Tensor([[0, 1, 0],
+                    [1, 4, 1],
+                    [0, 1, 0]]) / 4
+
+K_RB = torch.Tensor([[1, 2, 1],
+                     [2, 4, 2],
+                     [1, 2, 1]]) / 4
+
+M_RGB_2_YUV = torch.Tensor([[0.299, 0.587, 0.114],
+                            [-0.14714119, -0.28886916, 0.43601035],
+                            [0.61497538, -0.51496512, -0.10001026]])
+M_YUV_2_RGB = torch.Tensor([[1.0000000000e+00, -4.1827794561e-09, 1.1398830414e+00],
+                            [1.0000000000e+00, -3.9464232326e-01, -5.8062183857e-01],
+                            [1.0000000000e+00, 2.0320618153e+00, -1.2232658220e-09]])
+
+K_BLUR = torch.Tensor([[6.9625e-08, 2.8089e-05, 2.0755e-04, 2.8089e-05, 6.9625e-08],
+                       [2.8089e-05, 1.1332e-02, 8.3731e-02, 1.1332e-02, 2.8089e-05],
+                       [2.0755e-04, 8.3731e-02, 6.1869e-01, 8.3731e-02, 2.0755e-04],
+                       [2.8089e-05, 1.1332e-02, 8.3731e-02, 1.1332e-02, 2.8089e-05],
+                       [6.9625e-08, 2.8089e-05, 2.0755e-04, 2.8089e-05, 6.9625e-08]])
+K_SHARP = torch.Tensor([[0, -1, 0],
+                        [-1, 5, -1],
+                        [0, -1, 0]])
+DEFAULT_CAMERA_PARAMS = (
+    [0., 0., 0., 0.],
+    [1., 1., 1.],
+    [1., 0., 0., 0., 1., 0., 0., 0., 1.],
+)
+
+
+class RawToRGB(nn.Module):
+    def __init__(self, reduce_size=True, out_channels=3, track_stages=False, normalize_mosaic=None):
+        super().__init__()
+        self.stages = None
+        self.buffer = None
+        self.reduce_size = reduce_size
+        self.out_channels = out_channels
+        self.track_stages = track_stages
+        self.normalize_mosaic = normalize_mosaic
+
+    def forward(self, raw):
+        self.stages = {}
+        self.buffer = {}
+
+        rgb = raw2rgb(raw, reduce_size=self.reduce_size, out_channels=self.out_channels)
+        self.stages['demosaic'] = rgb
+        if self.normalize_mosaic:
+            rgb = self.normalize_mosaic(rgb)
+
+        if self.track_stages and raw.requires_grad:
+            for stage in self.stages.values():
+                stage.retain_grad()
+
+        self.buffer['processed_rgb'] = rgb
+
+        return rgb
+
+
+class NNProcessing(nn.Module):
+    def __init__(self, track_stages=False, normalize_mosaic=None, batch_norm_output=True):
+        super().__init__()
+        self.stages = None
+        self.buffer = None
+        self.track_stages = track_stages
+        self.model = smp.UnetPlusPlus(
+            encoder_name='resnet34',
+            encoder_depth=3,
+            decoder_channels=[256, 128, 64],
+            in_channels=3,
+            classes=3,
+        )
+        self.batch_norm = None if not batch_norm_output else nn.BatchNorm2d(3)
+        self.normalize_mosaic = normalize_mosaic
+
+    def forward(self, raw):
+        self.stages = {}
+        self.buffer = {}
+        # self.stages['raw'] = raw
+        rgb = raw2rgb(raw)
+        if self.normalize_mosaic:
+            rgb = self.normalize_mosaic(rgb)
+        self.stages['demosaic'] = rgb
+        rgb = self.model(rgb)
+        if self.batch_norm is not None:
+            rgb = self.batch_norm(rgb)
+        self.stages['rgb'] = rgb
+
+        if self.track_stages and raw.requires_grad:
+            for stage in self.stages.values():
+                stage.retain_grad()
+
+        self.buffer['processed_rgb'] = rgb
+
+        return rgb
+
+
+class ParametrizedProcessing(nn.Module):
+    def __init__(self, camera_parameters, track_stages=False, batch_norm_output=True, noise_layer=False):
+        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.gamma_correct = nn.Parameter(torch.Tensor([2.2]))
+
+        self.debayer = Debayer()
+
+        self.sharpening_filter = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False)
+        self.sharpening_filter.weight.data[0][0] = K_SHARP.clone()
+
+        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
+
+        # if noise_layer:
+        #     for param in self.parameters():
+        #         param.requires_grad = False
+
+        self.additive_layer = nn.Parameter(0.001 * torch.randn((1, 3, 256, 256))
+                                           ) if noise_layer else None  # XXX: can this be 0?
+
+    def forward(self, raw):
+        assert raw.ndim == 3, f"needs dims (B, H, W), got {raw.shape}"
+
+        self.stages = {}
+        self.buffer = {}
+
+        # self.stages['raw'] = raw
+
+        rgb = raw2rgb(raw, black_level=self.black_level, reduce_size=False)
+        rgb = rgb.contiguous()
+        self.stages['demosaic'] = rgb
+
+        rgb = self.debayer(rgb)
+        # self.stages['debayer'] = rgb
+
+        rgb = torch.einsum('bchw,kc->bkhw', rgb, self.colour_correction).contiguous()
+        self.stages['color_correct'] = rgb
+
+        yuv = torch.einsum('bchw,kc->bkhw', rgb, self.M_RGB_2_YUV).contiguous()
+        yuv[:, [0], ...] = self.sharpening_filter(yuv[:, [0], ...])
+
+        if self.track_stages:    # keep stage in computational graph for grad information
+            rgb = torch.einsum('bchw,kc->bkhw', yuv.clone(), self.M_YUV_2_RGB).contiguous()
+            self.stages['sharpening'] = rgb
+            yuv = torch.einsum('bchw,kc->bkhw', rgb, self.M_RGB_2_YUV).contiguous()
+
+        yuv[:, [0], ...] = self.gaussian_blur(yuv[:, [0], ...])
+        rgb = torch.einsum('bchw,kc->bkhw', yuv, self.M_YUV_2_RGB).contiguous()
+        self.stages['gaussian'] = rgb
+
+        rgb = torch.clip(rgb, 1e-5, 1)
+        self.stages['clipped'] = rgb
+
+        rgb = torch.exp((1 / self.gamma_correct) * torch.log(rgb))
+        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
+
+        if self.batch_norm is not None:
+            rgb = self.batch_norm(rgb)
+
+        if self.track_stages and raw.requires_grad:
+            for stage in self.stages.values():
+                stage.retain_grad()
+
+        self.buffer['processed_rgb'] = rgb
+
+        return rgb
+
+
+class Debayer(nn.Conv2d):
+    def __init__(self):
+        super().__init__(3, 3, kernel_size=3, padding=1, padding_mode='reflect', bias=False)    # default_pipeline uses 'replicate'
+        self.weight.data.fill_(0)
+        self.weight.data[0, 0] = K_RB.clone()
+        self.weight.data[1, 1] = K_G.clone()
+        self.weight.data[2, 2] = K_RB.clone()
+
+
+def raw2rgb(raw, black_level=None, reduce_size=True, out_channels=3):
+    """transform raw image with 1 channel to rgb with 3 channels 
+    Args:
+        raw (Tensor): raw Tensor of shape (B, H, W)
+        black_level (iterable, optional): RGGB black level values to subtract
+        reduce_size (bool, optional): if False, the output image will have the same height and width 
+            as the raw input, i.e. (B, C, H, W), empty values are filled with zeros.
+            if True, the output dimensions are reduced by half (B, C, H//2, W//2), 
+            the two green channels are averaged.
+        out_channels (int, optional): number of output channels. One of {3, 4}.
+    """
+    assert out_channels in [3, 4]
+    if black_level is None:
+        black_level = [0, 0, 0, 0]
+    Bch, H, W = raw.shape
+    R = raw[:, 0::2, 0::2] - black_level[0]     # R
+    G1 = raw[:, 0::2, 1::2] - black_level[1]    # G
+    G2 = raw[:, 1::2, 0::2] - black_level[2]    # G
+    B = raw[:, 1::2, 1::2] - black_level[3]     # B
+    if reduce_size:
+        rgb = torch.zeros((Bch, out_channels, H // 2, W // 2), device=raw.device)
+        if out_channels == 3:
+            rgb[:, 0, :, :] = R
+            rgb[:, 1, :, :] = (G1 + G2) / 2
+            rgb[:, 2, :, :] = B
+        elif out_channels == 4:
+            rgb[:, 0, :, :] = R
+            rgb[:, 1, :, :] = G1
+            rgb[:, 2, :, :] = G2
+            rgb[:, 3, :, :] = B
+    else:
+        rgb = torch.zeros((Bch, out_channels, H, W), device=raw.device)
+        if out_channels == 3:
+            rgb[:, 0, 0::2, 0::2] = R
+            rgb[:, 1, 0::2, 1::2] = G1
+            rgb[:, 1, 1::2, 0::2] = G2
+            rgb[:, 2, 1::2, 1::2] = B
+        elif out_channels == 4:
+            rgb[:, 0, 0::2, 0::2] = R
+            rgb[:, 1, 0::2, 1::2] = G1
+            rgb[:, 2, 1::2, 0::2] = G2
+            rgb[:, 3, 1::2, 1::2] = B
+    return rgb
+
+
+# pipeline validation
+if __name__ == "__main__":
+
+    import torch
+    import numpy as np
+
+    if not os.path.exists('README.md'):
+        os.chdir('..')
+
+    import matplotlib.pyplot as plt
+    from utils.dataset import get_dataset
+    from utils.base import np2torch, torch2np
+
+    from utils.debug import debug
+    from processingpipeline.pipeline import processing as default_processing
+
+    raw_dataset = get_dataset('DS')
+    loader = torch.utils.data.DataLoader(raw_dataset, batch_size=1)
+    batch_raw, batch_mask = next(iter(loader))
+
+    # torch proc
+    camera_parameters = raw_dataset.camera_parameters
+    black_level = camera_parameters[0]
+
+    proc = ParametrizedProcessing(camera_parameters)
+
+    batch_rgb = proc(batch_raw)
+    rgb = batch_rgb[0]
+
+    # numpy proc
+    raw_img = batch_raw[0]
+    numpy_raw = torch2np(raw_img)
+
+    default_rgb = default_processing(numpy_raw, *camera_parameters,
+                                     sharpening='sharpening_filter', denoising='gaussian_denoising')
+
+    rgb_valid = np2torch(default_rgb)
+
+    print("pipeline norm difference:", (rgb - rgb_valid).norm().item())
+
+    rgb_mosaic = raw2rgb(batch_raw, reduce_size=False).squeeze()
+    rgb_reduced = raw2rgb(batch_raw, reduce_size=True).squeeze()
+
+    plt.figure(figsize=(16, 8))
+    plt.subplot(151)
+    plt.title('Raw')
+    plt.imshow(torch2np(raw_img))
+    plt.subplot(152)
+    plt.title('RGB Mosaic')
+    plt.imshow(torch2np(rgb_mosaic))
+    plt.subplot(153)
+    plt.title('RGB Reduced')
+    plt.imshow(torch2np(rgb_reduced))
+    plt.subplot(154)
+    plt.title('Torch Pipeline')
+    plt.imshow(torch2np(rgb))
+    plt.subplot(155)
+    plt.title('Default Pipeline')
+    plt.imshow(torch2np(rgb_valid))
+    plt.show()
+
+    # assert rgb.allclose(rgb_valid)

sanity_checks_and_statistics.ipynb

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

show_classification_results.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62c5dbc4bb22ecd26bc691c1f574b6bcf07b7cd48f62668506955df3513afe55
+size 10556940

show_results.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+
+datasets='Microscopy Drone'
+augmentations='weak strong none'
+
+for augment in $augmentations
+    do
+    for data in $datasets
+        do
+        
+        python show_results.py \
+        --dataset $data \
+        --augmentation $augment \
+
+    done
+done

train.py

@@ -0,0 +1,420 @@
+import os
+import sys
+import copy
+import argparse
+
+import torch
+import torch.nn as nn
+
+import mlflow.pytorch
+from torch.utils.data import DataLoader
+from torchvision.models import resnet18
+import torchvision.transforms as T
+from pytorch_lightning.metrics.functional import accuracy
+import pytorch_lightning as pl
+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.augmentation import get_augmentation
+from utils.dataset import Subset, get_dataset, k_fold
+
+from processingpipeline.pipeline import RawProcessingPipeline
+from processingpipeline.torch_pipeline import raw2rgb, RawToRGB, ParametrizedProcessing, NNProcessing
+
+from models.classifier import log_tensor, resnet_model, LitModel, TrackImagesCallback
+
+import segmentation_models_pytorch as smp
+
+from utils.pytorch_ssim import SSIM
+
+# args to set up task
+parser = argparse.ArgumentParser(description="classification_task")
+parser.add_argument("--tracking_uri", type=str,
+                    default="http://deplo-mlflo-1ssxo94f973sj-890390d809901dbf.elb.eu-central-1.amazonaws.com", help='URI of the mlflow server on AWS')
+parser.add_argument("--processor_uri", type=str, default=None,
+                    help='URI of the processing model (e.g. s3://mlflow-artifacts-821771080529/1/5fa754c566e3466690b1d309a476340f/artifacts/processing-model)')
+parser.add_argument("--classifier_uri", type=str, default=None,
+                    help='URI of the net (e.g. s3://mlflow-artifacts-821771080529/1/5fa754c566e3466690b1d309a476340f/artifacts/prediction-model)')
+parser.add_argument("--state_dict_uri", type=str,
+                    default=None, help='URI of the indices you want to load (e.g. s3://mlflow-artifacts-601883093460/7/4326da05aca54107be8c554de0674a14/artifacts/training')
+
+parser.add_argument("--experiment_name", type=str,
+                    default='classification learnable pipeline', help='Specify the experiment you are running, e.g. end2end segmentation')
+parser.add_argument("--run_name", type=str,
+                    default='test run', help='Specify the name of your run')
+
+parser.add_argument("--log_model", type=str2bool, default=True, help='Enables model logging')
+parser.add_argument("--save_locally", action='store_true',
+                    help='Model will be saved locally if action is taken')   # TODO: bypass mlflow
+
+parser.add_argument("--track_processing", action='store_true',
+                    help='Save images after each trasformation of the pipeline for the test set')
+parser.add_argument("--track_processing_gradients", action='store_true',
+                    help='Save images of gradients after each trasformation of the pipeline for the test set')
+parser.add_argument("--track_save_tensors", action='store_true',
+                    help='Save the torch tensors after each trasformation of the pipeline for the test set')
+parser.add_argument("--track_predictions", action='store_true',
+                    help='Save images after each trasformation of the pipeline for the test set + input gradient')
+parser.add_argument("--track_n_images", default=5,
+                    help='Track the n first elements of dataset. Only used for args.track_processing=True')
+parser.add_argument("--track_every_epoch", action='store_true', help='Track images every epoch or once after training')
+
+# args to create dataset
+parser.add_argument("--seed", type=int, default=1, help='Global seed')
+parser.add_argument("--dataset", type=str, default='Microscopy',
+                    choices=["Drone", "DroneSegmentation", "Microscopy"], help='Select dataset')
+
+parser.add_argument("--n_splits", type=int, default=1, help='Number of splits used for training')
+parser.add_argument("--train_size", type=float, default=0.8, help='Fraction of training points in dataset')
+
+# args for training
+parser.add_argument("--lr", type=float, default=1e-5, help="learning rate used for training")
+parser.add_argument("--epochs", type=int, default=3, help="numper of epochs")
+parser.add_argument("--batch_size", type=int, default=32, help="Training batch size")
+parser.add_argument("--augmentation", type=str, default='none',
+                    choices=["none", "weak", "strong"], help="Applies augmentation to training")
+parser.add_argument("--augmentation_on_valid_epoch", action='store_true',
+                    help='Track images every epoch or once after training')  # TODO: implement, actually should be disabled by default for 'val' and 'test
+parser.add_argument("--check_val_every_n_epoch", type=int, default=1)
+
+# args to specify the processing
+parser.add_argument("--processing_mode", type=str, default="parametrized",
+                    choices=["parametrized", "static", "neural_network", "none"],
+                    help="Which type of raw to rgb processing should be used")
+
+# args to specify model
+parser.add_argument("--classifier_network", type=str, default='ResNet18',
+                    help='Type of pretrained network')  # TODO: implement different choices
+parser.add_argument("--classifier_pretrained", action='store_true',
+                    help='Whether to use a pre-trained model or not')
+parser.add_argument("--smp_encoder", type=str, default='resnet34', help='segmentation model encoder')
+
+parser.add_argument("--freeze_processor", action='store_true', help="Freeze raw to rgb processing model weights")
+parser.add_argument("--freeze_classifier", action='store_true', help="Freeze classification model weights")
+
+# args to specify static pipeline transformations
+parser.add_argument("--sp_debayer", type=str, default='bilinear',
+                    choices=['bilinear', 'malvar2004', 'menon2007'], help="Specify algorithm used as debayer")
+parser.add_argument("--sp_sharpening", type=str, default='sharpening_filter',
+                    choices=['sharpening_filter', 'unsharp_masking'], help="Specify algorithm used for sharpening")
+parser.add_argument("--sp_denoising", type=str, default='gaussian_denoising',
+                    choices=['gaussian_denoising', 'median_denoising', 'fft_denoising'], help="Specify algorithm used for denoising")
+
+# args to choose training mode
+parser.add_argument("--adv_training", action='store_true', help="Enable adversarial training")
+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("--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()
+
+
+def run_train(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
+
+    # dataset
+
+    dataset = get_dataset(args.dataset)
+
+    print(f'dataset: {type(dataset).__name__}[{len(dataset)}]')
+    print(f'task: {dataset.task}')
+    print(f'mode: {training_mode} training')
+    print(f'# cross-validation subsets: {args.n_splits}')
+    pl.seed_everything(args.seed)
+    idxs_kfold = k_fold(dataset, n_splits=args.n_splits, seed=args.seed, train_size=args.train_size)
+
+    with mlflow.start_run(run_name=args.run_name) as parent_run:
+
+        for k_iter, idxs in enumerate(idxs_kfold):
+
+            print(f"K_fold subset: {k_iter+1}/{args.n_splits}")
+
+            if args.processing_mode == 'static':
+                if args.dataset == "Drone" or args.dataset == "DroneSegmentation":
+                    mean = torch.tensor([0.35, 0.36, 0.35])
+                    std = torch.tensor([0.12, 0.11, 0.12])
+                elif args.dataset == "Microscopy":
+                    mean = torch.tensor([0.91, 0.84, 0.94])
+                    std = torch.tensor([0.08, 0.12, 0.05])
+
+                dataset.transform = T.Compose([RawProcessingPipeline(
+                    camera_parameters=dataset.camera_parameters,
+                    debayer=args.sp_debayer,
+                    sharpening=args.sp_sharpening,
+                    denoising=args.sp_denoising,
+                ), T.Normalize(mean, std)])
+                # XXX: Not clean
+
+                processor = nn.Identity()
+
+            if args.processor_uri is not None and args.processing_mode != 'none':
+                print('Fetching processor: ', end='')
+                model = fetch_from_mlflow(args.processor_uri, use_cache=args.cache_downloaded_models)
+                processor = model.processor
+                for param in processor.parameters():
+                    param.requires_grad = True
+                model.processor = None
+                del model
+            else:
+                print(f'processing_mode: {args.processing_mode}')
+                normalize_mosaic = None   # normalize after raw has been passed to raw2rgb
+                if args.dataset == "Microscopy":
+                    mosaic_mean = [0.5663, 0.1401, 0.0731]
+                    mosaic_std = [0.097, 0.0423, 0.008]
+                    normalize_mosaic = T.Normalize(mosaic_mean, mosaic_std)
+
+                track_stages = args.track_processing or args.track_processing_gradients
+                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?
+                    )
+                elif args.processing_mode == 'neural_network':
+                    processor = NNProcessing(track_stages=track_stages,
+                                             normalize_mosaic=normalize_mosaic, batch_norm_output=True)
+                elif args.processing_mode == 'none':
+                    processor = RawToRGB(reduce_size=True, out_channels=3, track_stages=track_stages,
+                                         normalize_mosaic=normalize_mosaic)
+
+            if args.classifier_uri:  # fetch classifier
+                print('Fetching classifier: ', end='')
+                model = fetch_from_mlflow(args.classifier_uri, use_cache=args.cache_downloaded_models)
+                classifier = model.classifier
+                model.classifier = None
+                del model
+            else:
+                if dataset.task == 'classification':
+                    classifier = resnet_model(
+                        model=resnet18,
+                        pretrained=args.classifier_pretrained,
+                        in_channels=3,
+                        fc_out_features=len(dataset.classes)
+                    )
+                else:
+                    # XXX: add other network choices to args.smp_network (FPN) and args.network
+                    classifier = smp.UnetPlusPlus(
+                        encoder_name=args.smp_encoder,
+                        encoder_depth=5,
+                        encoder_weights='imagenet',
+                        in_channels=3,
+                        classes=1,
+                        activation=None,
+                    )
+
+            if args.freeze_processor and len(list(iter(processor.parameters()))) == 0:
+                print('Note: freezing processor without parameters.')
+            assert not (args.freeze_processor and args.freeze_classifier), 'Likely no parameters to train.'
+
+            if dataset.task == 'classification':
+                loss = nn.CrossEntropyLoss()
+                metrics = [accuracy]
+            else:
+                # loss = utils.base.smp_get_loss(args.smp_loss)    # XXX: add other losses to args.smp_loss
+                loss = smp.losses.DiceLoss(mode='binary', from_logits=True)
+                metrics = [smp.utils.metrics.IoU()]
+
+            loss_aux = None
+
+            if args.adv_training:
+
+                assert args.processing_mode == 'parametrized', f"Processing mode ({args.processing_mode}) should be set to 'parametrized' for adversarial training"
+                assert args.freeze_classifier, "Classifier should be frozen for adversarial training"
+                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.eval()
+                processor_default.to(DEVICE)
+                # debug(processor_default)
+
+                def l2_regularization(x, y):
+                    return (x - y).norm()
+
+                if args.adv_aux_loss == 'l2':
+                    regularization = l2_regularization
+                elif args.adv_aux_loss == 'ssim':
+                    regularization = SSIM(window_size=11)
+                else:
+                    NotImplementedError(args.adv_aux_loss)
+
+                class AuxLoss(nn.Module):
+                    def __init__(self, loss_aux, weight=1):
+                        super().__init__()
+                        self.loss_aux = loss_aux
+                        self.weight = weight
+
+                    def forward(self, x):
+                        x_reference = processor_default(x)
+                        x_processed = processor.buffer['processed_rgb']
+                        return self.weight * self.loss_aux(x_reference, x_processed)
+
+                class WeightedLoss(nn.Module):
+                    def __init__(self, loss, weight=1):
+                        super().__init__()
+                        self.loss = loss
+                        self.weight = weight
+
+                    def forward(self, x, y):
+                        return self.weight * self.loss(x, y)
+
+                    def __repr__(self):
+                        return f'{self.weight} * {get_name(self.loss)}'
+
+                loss = WeightedLoss(loss=nn.CrossEntropyLoss(), weight=-1)
+                # loss = WeightedLoss(loss=nn.CrossEntropyLoss(), weight=0)
+                loss_aux = AuxLoss(
+                    loss_aux=regularization,
+                    weight=args.adv_aux_weight,
+                )
+
+            augmentation = get_augmentation(args.augmentation)
+
+            model = LitModel(
+                classifier=classifier,
+                processor=processor,
+                loss=loss,
+                loss_aux=loss_aux,
+                adv_training=args.adv_training,
+                metrics=metrics,
+                augmentation=augmentation,
+                is_segmentation_task=dataset.task == 'segmentation',
+                freeze_classifier=args.freeze_classifier,
+                freeze_processor=args.freeze_processor,
+            )
+
+            # get train_set_dict
+            if args.state_dict_uri:
+                state_dict = mlflow.pytorch.load_state_dict(args.state_dict_uri)
+                train_indices = state_dict['train_indices']
+                valid_indices = state_dict['valid_indices']
+            else:
+                train_indices = idxs[0]
+                valid_indices = idxs[1]
+                state_dict = vars(args).copy()
+
+            track_indices = list(range(args.track_n_images))
+
+            if dataset.task == 'classification':
+                state_dict['classes'] = dataset.classes
+            state_dict['device'] = DEVICE
+            state_dict['train_indices'] = train_indices
+            state_dict['valid_indices'] = valid_indices
+            state_dict['elements in train set'] = len(train_indices)
+            state_dict['elements in test set'] = len(valid_indices)
+
+            if args.test_run:
+                train_indices = train_indices[:args.batch_size]
+                valid_indices = valid_indices[:args.batch_size]
+
+            train_set = Subset(dataset, indices=train_indices)
+            valid_set = Subset(dataset, indices=valid_indices)
+            track_set = Subset(dataset, indices=track_indices)
+
+            train_loader = DataLoader(train_set, batch_size=args.batch_size, num_workers=16, shuffle=True)
+            valid_loader = DataLoader(valid_set, batch_size=args.batch_size, num_workers=16, shuffle=False)
+            track_loader = DataLoader(track_set, batch_size=args.batch_size, num_workers=16, shuffle=False)
+
+            with mlflow.start_run(run_name=f"{args.run_name}_{k_iter}", nested=True) as child_run:
+
+                #mlflow.pytorch.autolog(silent=True)
+
+                if k_iter == 0:
+                    display_mlflow_run_info(child_run)
+
+                mlflow.pytorch.log_state_dict(state_dict, artifact_path=None)
+
+                hparams = {
+                    'dataset': args.dataset,
+                    'processing_mode': args.processing_mode,
+                    'training_mode': training_mode,
+                }
+                if training_mode == 'adversarial':
+                    hparams['adv_aux_weight'] = args.adv_aux_weight
+                    hparams['adv_aux_loss'] = args.adv_aux_loss
+
+                mlflow.log_params(hparams)
+
+                with open('results/state_dict.txt', 'w') as f:
+                    f.write('python ' + ' '.join(sys.argv) + '\n')
+                    f.write('\n'.join([f'{k}={v}' for k, v in state_dict.items()]))
+                mlflow.log_artifact('results/state_dict.txt', artifact_path=None)
+
+                mlf_logger = pl.loggers.MLFlowLogger(experiment_name=args.experiment_name,
+                                                     tracking_uri=args.tracking_uri,)
+                mlf_logger._run_id = child_run.info.run_id
+
+                callbacks = []
+                if args.track_processing:
+                    callbacks += [TrackImagesCallback(track_loader,
+                                                      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 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, 
+                #    else:
+                #        checkpoint_callback = ModelCheckpoint(monitor="val_iou_score")
+                #callbacks += [checkpoint_callback]
+
+                trainer = pl.Trainer(
+                    gpus=1 if DEVICE == 'cuda' else 0,
+                    min_epochs=args.epochs,
+                    max_epochs=args.epochs,
+                    logger=mlf_logger,
+                    callbacks=callbacks,
+                    check_val_every_n_epoch=args.check_val_every_n_epoch,
+                    #checkpoint_callback=True,
+                )
+
+                if args.log_model:
+                    mlflow.pytorch.autolog(log_every_n_epoch=10)
+                    print(f'model_uri="{mlflow.get_artifact_uri()}/model"')
+
+                t = trainer.fit(
+                    model,
+                    train_dataloader=train_loader,
+                    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()}")
+    return model
+
+
+if __name__ == '__main__':
+    model = run_train(args)

train.sh

@@ -0,0 +1,81 @@
+#!/bin/bash
+
+# # Parametrized Training
+# 100 epochs, frozen_processor: http://deplo-mlflo-1ssxo94f973sj-890390d809901dbf.elb.eu-central-1.amazonaws.com/#/experiments/49/runs/2803f44514e34a0f87d591520706e876
+# model_uri="s3://mlflow-artifacts-601883093460/49/2803f44514e34a0f87d591520706e876/artifacts/model"
+
+# used for training current model to 100% train and 80% val accuracy
+# python train.py \
+# --experiment_name parametrized \
+# --classifier_uri "${model_uri}" \
+# --run_name par_full_kurt \
+# --dataset Microscopy \
+# --lr 1e-5 \
+# --epochs 50 \
+# --freeze_classifier \
+
+# --freeze_processor \
+
+# # Adversarial Training
+
+# python train.py \
+# --experiment_name adversarial \
+# --run_name adv_frozen_processor \
+# --classifier_uri "${model_uri}" \
+# --dataset Microscopy \
+# --adv_training \
+# --lr 1e-3 \
+# --epochs 7 \
+# --freeze_classifier \
+# --track_processing \
+# --track_every_epoch \
+# --log_model=False \
+# --adv_aux_weight=0.1 \
+# --adv_aux_loss "l2" \
+
+# --adv_aux_weight=2e-5 \
+# --adv_aux_weight=2e-5 \
+# --adv_aux_weight=1.9e-5 \
+
+# Cross pipeline training (Segmentation/Classification)
+
+# Static Pipeline Script        
+
+# datasets="Microscopy Drone DroneSegmentation"
+datasets="DroneSegmentation"
+augmentations="weak strong none"
+
+demosaicings="bilinear malvar2004 menon2007"
+sharpenings="sharpening_filter unsharp_masking"
+denoisings="median_denoising gaussian_denoising"
+
+for augment in $augmentations
+    do
+    for data in $datasets
+        do
+        for demosaicing in $demosaicings 
+            do
+            for sharpening in $sharpenings
+                do
+                for denoising in $denoisings
+                    do
+
+                    python train.py \
+                    --experiment_name ABtesting \
+                    --run_name "$data"_"$demosaicing"_"$sharpening"_"$denoising"_"$augment" \
+                    --dataset "$data" \
+                    --batch_size 4 \
+                    --lr 1e-5 \
+                    --epochs 100 \
+                    --sp_debayer "$demosaicing" \
+                    --sp_sharpening "$sharpening" \
+                    --sp_denoising "$denoising" \
+                    --processing_mode "static" \
+                    --augmentation "$augment" \
+                    --n_split 5 \
+                    
+                done
+            done
+        done
+    done
+done

utils/Cperturb.py

@@ -0,0 +1,475 @@
+'''
+Code extracted from the paper:
+
+@articlehendrycks2019robustness,
+  title=Benchmarking Neural Network Robustness to Common Corruptions and Perturbations,
+  author=Dan Hendrycks and Thomas Dietterich,
+  journal=Proceedings of the International Conference on Learning Representations,
+  year=2019
+}
+
+The code is modified to fit with our model
+'''
+
+import os
+from PIL import Image
+import os.path
+import time
+import torch
+import torchvision.datasets as dset
+import torchvision.transforms as trn
+import torch.utils.data as data
+import numpy as np
+
+from PIL import Image
+
+
+# /////////////// Distortion Helpers ///////////////
+
+import skimage as sk
+from skimage.filters import gaussian
+from io import BytesIO
+from wand.image import Image as WandImage
+from wand.api import library as wandlibrary
+import wand.color as WandColor
+import ctypes
+from PIL import Image as PILImage
+import cv2
+from scipy.ndimage import zoom as scizoom
+from scipy.ndimage.interpolation import map_coordinates
+import warnings
+
+warnings.simplefilter("ignore", UserWarning)
+
+
+def disk(radius, alias_blur=0.1, dtype=np.float32):
+    if radius <= 8:
+        L = np.arange(-8, 8 + 1)
+        ksize = (3, 3)
+    else:
+        L = np.arange(-radius, radius + 1)
+        ksize = (5, 5)
+    X, Y = np.meshgrid(L, L)
+    aliased_disk = np.array((X ** 2 + Y ** 2) <= radius ** 2, dtype=dtype)
+    aliased_disk /= np.sum(aliased_disk)
+
+    # supersample disk to antialias
+    return cv2.GaussianBlur(aliased_disk, ksize=ksize, sigmaX=alias_blur)
+
+
+# Tell Python about the C method
+wandlibrary.MagickMotionBlurImage.argtypes = (ctypes.c_void_p,  # wand
+                                              ctypes.c_double,  # radius
+                                              ctypes.c_double,  # sigma
+                                              ctypes.c_double)  # angle
+
+
+# Extend wand.image.Image class to include method signature
+class MotionImage(WandImage):
+    def motion_blur(self, radius=0.0, sigma=0.0, angle=0.0):
+        wandlibrary.MagickMotionBlurImage(self.wand, radius, sigma, angle)
+
+
+# modification of https://github.com/FLHerne/mapgen/blob/master/diamondsquare.py
+def plasma_fractal(mapsize=32, wibbledecay=3):
+    """
+    Generate a heightmap using diamond-square algorithm.
+    Return square 2d array, side length 'mapsize', of floats in range 0-255.
+    'mapsize' must be a power of two.
+    """
+    assert (mapsize & (mapsize - 1) == 0)
+    maparray = np.empty((mapsize, mapsize), dtype=np.float_)
+    maparray[0, 0] = 0
+    stepsize = mapsize
+    wibble = 100
+
+    def wibbledmean(array):
+        return array / 4 + wibble * np.random.uniform(-wibble, wibble, array.shape)
+
+    def fillsquares():
+        """For each square of points stepsize apart,
+           calculate middle value as mean of points + wibble"""
+        cornerref = maparray[0:mapsize:stepsize, 0:mapsize:stepsize]
+        squareaccum = cornerref + np.roll(cornerref, shift=-1, axis=0)
+        squareaccum += np.roll(squareaccum, shift=-1, axis=1)
+        maparray[stepsize // 2:mapsize:stepsize,
+        stepsize // 2:mapsize:stepsize] = wibbledmean(squareaccum)
+
+    def filldiamonds():
+        """For each diamond of points stepsize apart,
+           calculate middle value as mean of points + wibble"""
+        mapsize = maparray.shape[0]
+        drgrid = maparray[stepsize // 2:mapsize:stepsize, stepsize // 2:mapsize:stepsize]
+        ulgrid = maparray[0:mapsize:stepsize, 0:mapsize:stepsize]
+        ldrsum = drgrid + np.roll(drgrid, 1, axis=0)
+        lulsum = ulgrid + np.roll(ulgrid, -1, axis=1)
+        ltsum = ldrsum + lulsum
+        maparray[0:mapsize:stepsize, stepsize // 2:mapsize:stepsize] = wibbledmean(ltsum)
+        tdrsum = drgrid + np.roll(drgrid, 1, axis=1)
+        tulsum = ulgrid + np.roll(ulgrid, -1, axis=0)
+        ttsum = tdrsum + tulsum
+        maparray[stepsize // 2:mapsize:stepsize, 0:mapsize:stepsize] = wibbledmean(ttsum)
+
+    while stepsize >= 2:
+        fillsquares()
+        filldiamonds()
+        stepsize //= 2
+        wibble /= wibbledecay
+
+    maparray -= maparray.min()
+    return maparray / maparray.max()
+
+
+def clipped_zoom(img, zoom_factor):
+    h = img.shape[0]
+    # ceil crop height(= crop width)
+    ch = int(np.ceil(h / zoom_factor))
+
+    top = (h - ch) // 2
+    img = scizoom(img[top:top + ch, top:top + ch], (zoom_factor, zoom_factor, 1), order=1)
+    # trim off any extra pixels
+    trim_top = (img.shape[0] - h) // 2
+
+    return img[trim_top:trim_top + h, trim_top:trim_top + h]
+
+
+# /////////////// End Distortion Helpers ///////////////
+
+
+# /////////////// Distortions ///////////////
+
+class Distortions:
+    def __init__(self, severity=1, transform='identity'):
+        self.severity = severity
+        self.transform = transform
+
+    def __call__(self, img):
+        assert torch.is_tensor(img), 'Input data need to be a torch.tensor'
+        assert len(img.shape) == 3, 'Input image should be RGB'
+        img = self.torch2np(img)
+        t = getattr(self, self.transform)
+        img = t(img, self.severity)
+        return self.np2torch(img).float()
+
+    def np2torch(self,x):
+        return torch.tensor(x).permute(2,0,1)
+
+    def torch2np(self,x):
+        return np.array(x.permute(1,2,0))
+    
+    def identity(self,x, severity=1):
+        return x
+
+    def gaussian_noise(self, x, severity=1):
+        c = [0.04, 0.06, .08, .09, .10][severity - 1]
+        return np.clip(x + np.random.normal(size=x.shape, scale=c), 0, 1)
+
+
+    def shot_noise(self, x, severity=1):
+        c = [500, 250, 100, 75, 50][severity - 1]
+        return np.clip(np.random.poisson(x * c) / c, 0, 1)
+
+
+    def impulse_noise(self, x, severity=1):
+        c = [.01, .02, .03, .05, .07][severity - 1]
+
+        x = sk.util.random_noise(x, mode='s&p', amount=c)
+        return np.clip(x, 0, 1)
+
+
+    def speckle_noise(self, x, severity=1):
+        c = [.06, .1, .12, .16, .2][severity - 1]
+        return np.clip(x + x * np.random.normal(size=x.shape, scale=c), 0, 1)
+
+
+    def gaussian_blur(self, x, severity=1):
+        c = [.4, .6, 0.7, .8, 1][severity - 1]
+
+        x = gaussian(x, sigma=c, multichannel=True)
+        return np.clip(x, 0, 1)
+
+
+    def glass_blur(self, x, severity=1):
+        # sigma, max_delta, iterations
+        c = [(0.05,1,1), (0.25,1,1), (0.4,1,1), (0.25,1,2), (0.4,1,2)][severity - 1]
+
+        x = gaussian(x, sigma=c[0], multichannel=True)
+
+        # locally shuffle pixels
+        for i in range(c[2]):
+            for h in range(32 - c[1], c[1], -1):
+                for w in range(32 - c[1], c[1], -1):
+                    dx, dy = np.random.randint(-c[1], c[1], size=(2,))
+                    h_prime, w_prime = h + dy, w + dx
+                    # swap
+                    x[h, w], x[h_prime, w_prime] = x[h_prime, w_prime], x[h, w]
+
+        return np.clip(gaussian(x, sigma=c[0], multichannel=True), 0, 1)
+
+
+    def defocus_blur(self, x, severity=1):
+        c = [(0.3, 0.4), (0.4, 0.5), (0.5, 0.6), (1, 0.2), (1.5, 0.1)][severity - 1]
+        kernel = disk(radius=c[0], alias_blur=c[1])
+
+        channels = []
+        for d in range(3):
+            channels.append(cv2.filter2D(x[:, :, d], -1, kernel))
+        channels = np.array(channels).transpose((1, 2, 0))  # 3x32x32 -> 32x32x3
+
+        return np.clip(channels, 0, 1)
+
+
+    def motion_blur(self, x, severity=1):
+        c = [(6,1), (6,1.5), (6,2), (8,2), (9,2.5)][severity - 1]
+
+        output = BytesIO()
+        x.save(output, format='PNG')
+        x = MotionImage(blob=output.getvalue())
+
+        x.motion_blur(radius=c[0], sigma=c[1], angle=np.random.uniform(-45, 45))
+
+        x = cv2.imdecode(np.fromstring(x.make_blob(), np.uint8),
+                        cv2.IMREAD_UNCHANGED)
+
+        if x.shape != (32, 32):
+            return np.clip(x[..., [2, 1, 0]], 0, 1)  # BGR to RGB
+        else:  # greyscale to RGB
+            return np.clip(np.array([x, x, x]).transpose((1, 2, 0)), 0, 1)
+
+
+    def zoom_blur(self, x, severity=1):
+        c = [np.arange(1, 1.06, 0.01), np.arange(1, 1.11, 0.01), np.arange(1, 1.16, 0.01),
+            np.arange(1, 1.21, 0.01), np.arange(1, 1.26, 0.01)][severity - 1]
+        out = np.zeros_like(x)
+        for zoom_factor in c:
+            out += clipped_zoom(x, zoom_factor)
+
+        x = (x + out) / (len(c) + 1)
+        return np.clip(x, 0, 1)
+
+
+    def fog(self, x, severity=1):
+        c = [(.2,3), (.5,3), (0.75,2.5), (1,2), (1.5,1.75)][severity - 1]
+        max_val = x.max()
+        x += c[0] * plasma_fractal(wibbledecay=c[1])[:32, :32][..., np.newaxis]
+        return np.clip(x * max_val / (max_val + c[0]), 0, 1)
+
+
+    def frost(self, x, severity=1):
+        c = [(1, 0.2), (1, 0.3), (0.9, 0.4), (0.85, 0.4), (0.75, 0.45)][severity - 1]
+        idx = np.random.randint(5)
+        filename = ['./frost1.png', './frost2.png', './frost3.png', './frost4.jpg', './frost5.jpg', './frost6.jpg'][idx]
+        frost = cv2.imread(filename)
+        frost = cv2.resize(frost, (0, 0), fx=0.2, fy=0.2)
+        # randomly crop and convert to rgb
+        x_start, y_start = np.random.randint(0, frost.shape[0] - 32), np.random.randint(0, frost.shape[1] - 32)
+        frost = frost[x_start:x_start + 32, y_start:y_start + 32][..., [2, 1, 0]]
+
+        return np.clip(c[0] * np.array(x) + c[1] * frost, 0, 1)
+
+
+    def snow(self, x, severity=1):
+        c = [(0.1,0.2,1,0.6,8,3,0.95),
+            (0.1,0.2,1,0.5,10,4,0.9),
+            (0.15,0.3,1.75,0.55,10,4,0.9),
+            (0.25,0.3,2.25,0.6,12,6,0.85),
+            (0.3,0.3,1.25,0.65,14,12,0.8)][severity - 1]
+
+        snow_layer = np.random.normal(size=x.shape[:2], loc=c[0], scale=c[1])  # [:2] for monochrome
+
+        snow_layer = clipped_zoom(snow_layer[..., np.newaxis], c[2])
+        snow_layer[snow_layer < c[3]] = 0
+
+        snow_layer = PILImage.fromarray((np.clip(snow_layer.squeeze(), 0, 1) * 255).astype(np.uint8), mode='L')
+        output = BytesIO()
+        snow_layer.save(output, format='PNG')
+        snow_layer = MotionImage(blob=output.getvalue())
+
+        snow_layer.motion_blur(radius=c[4], sigma=c[5], angle=np.random.uniform(-135, -45))
+
+        snow_layer = cv2.imdecode(np.fromstring(snow_layer.make_blob(), np.uint8),
+                                cv2.IMREAD_UNCHANGED) / (2**16-1)
+        snow_layer = snow_layer[..., np.newaxis]
+
+        x = c[6] * x + (1 - c[6]) * np.maximum(x, cv2.cvtColor(x, cv2.COLOR_RGB2GRAY).reshape(32, 32, 1) * 1.5 + 0.5)
+        return np.clip(x + snow_layer + np.rot90(snow_layer, k=2), 0, 1)
+
+
+    def spatter(self, x, severity=1):
+        c = [(0.62,0.1,0.7,0.7,0.5,0),
+            (0.65,0.1,0.8,0.7,0.5,0),
+            (0.65,0.3,1,0.69,0.5,0),
+            (0.65,0.1,0.7,0.69,0.6,1),
+            (0.65,0.1,0.5,0.68,0.6,1)][severity - 1]
+
+        liquid_layer = np.random.normal(size=x.shape[:2], loc=c[0], scale=c[1])
+
+        liquid_layer = gaussian(liquid_layer, sigma=c[2])
+        liquid_layer[liquid_layer < c[3]] = 0
+        if c[5] == 0:
+            liquid_layer = (liquid_layer * (2**16-1)).astype(np.uint8)
+            dist = (2**16-1) - cv2.Canny(liquid_layer, 50, 150)
+            dist = cv2.distanceTransform(dist, cv2.DIST_L2, 5)
+            _, dist = cv2.threshold(dist, 20, 20, cv2.THRESH_TRUNC)
+            dist = cv2.blur(dist, (3, 3)).astype(np.uint8)
+            dist = cv2.equalizeHist(dist)
+            #     ker = np.array([[-1,-2,-3],[-2,0,0],[-3,0,1]], dtype=np.float32)
+            #     ker -= np.mean(ker)
+            ker = np.array([[-2, -1, 0], [-1, 1, 1], [0, 1, 2]])
+            dist = cv2.filter2D(dist, cv2.CV_8U, ker)
+            dist = cv2.blur(dist, (3, 3)).astype(np.float32)
+
+            m = cv2.cvtColor(liquid_layer * dist, cv2.COLOR_GRAY2BGRA)
+            m /= np.max(m, axis=(0, 1))
+            m *= c[4]
+
+            # water is pale turqouise
+            color = np.concatenate((175 / 255. * np.ones_like(m[..., :1]),
+                                    238 / 255. * np.ones_like(m[..., :1]),
+                                    238 / 255. * np.ones_like(m[..., :1])), axis=2)
+
+            color = cv2.cvtColor(color, cv2.COLOR_BGR2BGRA)
+            x = cv2.cvtColor(x, cv2.COLOR_BGR2BGRA)
+
+            return cv2.cvtColor(np.clip(x + m * color, 0, 1), cv2.COLOR_BGRA2BGR) * (2**16-1)
+        else:
+            m = np.where(liquid_layer > c[3], 1, 0)
+            m = gaussian(m.astype(np.float32), sigma=c[4])
+            m[m < 0.8] = 0
+            #         m = np.abs(m) ** (1/c[4])
+
+            # mud brown
+            color = np.concatenate((63 / 255. * np.ones_like(x[..., :1]),
+                                    42 / 255. * np.ones_like(x[..., :1]),
+                                    20 / 255. * np.ones_like(x[..., :1])), axis=2)
+
+            color *= m[..., np.newaxis]
+            x *= (1 - m[..., np.newaxis])
+
+            return np.clip(x + color, 0, 1)
+
+
+    def contrast(self, x, severity=1):
+        c = [.75, .5, .4, .3, 0.15][severity - 1]
+        means = np.mean(x, axis=(0, 1), keepdims=True)
+        return np.clip((x - means) * c + means, 0, 1)
+
+
+    def brightness(self, x, severity=1):
+        c = [.05, .1, .15, .2, .3][severity - 1]
+
+        x = sk.color.rgb2hsv(x)
+        x[:, :, 2] = np.clip(x[:, :, 2] + c, 0, 1)
+        x = sk.color.hsv2rgb(x)
+
+        return np.clip(x, 0, 1)
+
+
+    def saturate(self, x, severity=1):
+        c = [(0.3, 0), (0.1, 0), (1.5, 0), (2, 0.1), (2.5, 0.2)][severity - 1]
+
+        x = sk.color.rgb2hsv(x)
+        x[:, :, 1] = np.clip(x[:, :, 1] * c[0] + c[1], 0, 1)
+        x = sk.color.hsv2rgb(x)
+
+        return np.clip(x, 0, 1)
+
+
+    def jpeg_compression(self, x, severity=1):
+        c = [80, 65, 58, 50, 40][severity - 1]
+
+        output = BytesIO()
+        x.save(output, 'JPEG', quality=c)
+        x = PILImage.open(output)
+
+        return x
+
+
+    def pixelate(self, x, severity=1):
+        c = [0.95, 0.9, 0.85, 0.75, 0.65][severity - 1]
+
+        x = x.resize((int(32 * c), int(32 * c)), PILImage.BOX)
+        x = x.resize((32, 32), PILImage.BOX)
+
+        return x
+
+
+    # mod of https://gist.github.com/erniejunior/601cdf56d2b424757de5
+    def elastic_transform(self, image, severity=1):
+        IMSIZE = 32
+        c = [(IMSIZE*0, IMSIZE*0, IMSIZE*0.08),
+            (IMSIZE*0.05, IMSIZE*0.2, IMSIZE*0.07),
+            (IMSIZE*0.08, IMSIZE*0.06, IMSIZE*0.06),
+            (IMSIZE*0.1, IMSIZE*0.04, IMSIZE*0.05),
+            (IMSIZE*0.1, IMSIZE*0.03, IMSIZE*0.03)][severity - 1]
+
+        shape = image.shape
+        shape_size = shape[:2]
+
+        # random affine
+        center_square = np.float32(shape_size) // 2
+        square_size = min(shape_size) // 3
+        pts1 = np.float32([center_square + square_size,
+                        [center_square[0] + square_size, center_square[1] - square_size],
+                        center_square - square_size])
+        pts2 = pts1 + np.random.uniform(-c[2], c[2], size=pts1.shape).astype(np.float32)
+        M = cv2.getAffineTransform(pts1, pts2)
+        image = cv2.warpAffine(image, M, shape_size[::-1], borderMode=cv2.BORDER_REFLECT_101)
+
+        dx = (gaussian(np.random.uniform(-1, 1, size=shape[:2]),
+                    c[1], mode='reflect', truncate=3) * c[0]).astype(np.float32)
+        dy = (gaussian(np.random.uniform(-1, 1, size=shape[:2]),
+                    c[1], mode='reflect', truncate=3) * c[0]).astype(np.float32)
+        dx, dy = dx[..., np.newaxis], dy[..., np.newaxis]
+
+        x, y, z = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]), np.arange(shape[2]))
+        indices = np.reshape(y + dy, (-1, 1)), np.reshape(x + dx, (-1, 1)), np.reshape(z, (-1, 1))
+        return np.clip(map_coordinates(image, indices, order=1, mode='reflect').reshape(shape), 0, 1)
+
+if __name__=='__main__':
+    import os
+
+    import numpy as np
+    import matplotlib.pyplot as plt
+    import tifffile as tiff
+    import torch
+
+    os.system('cd ..')
+
+    img = tiff.imread('/home/marco/perturbed-minds/perturbed-minds/data/microscopy/images/rgb_scale100/Ma190c_lame1_zone1_composite_Mcropped_1.tiff')
+    img = np.array(img)/(2**16-1)
+    img = torch.tensor(img).permute(2,0,1)
+
+    def identity(x, sev):
+            return x
+
+    if not os.path.exists('results/Cimages'):
+        os.makedirs('results/Cimages')
+
+    transformations = ['gaussian_noise', 'shot_noise', 'impulse_noise', 'speckle_noise', 
+                    'gaussian_blur', 'zoom_blur', 'contrast', 'brightness', 'saturate', 'elastic_transform']
+
+# glass_blur, defocus_blur, motion_blur, fog, frost, snow, spatter, jpeg_compression, pixelate,
+
+    plt.figure()
+    plt.imshow(img.permute(1,2,0))
+    plt.title('identity')
+    plt.show()
+    plt.savefig(f'results/Cimages/1_identity.png')
+
+
+    for i,t in enumerate(transformations):
+
+        fig = plt.figure(figsize=(25,5))
+        columns = 5
+        rows = 1        
+        
+        for sev in range(1,6):
+                dist = Distortions(severity=sev, transform=t)     
+                fig.add_subplot(rows, columns, sev)
+                plt.imshow(dist(img).permute(1,2,0))
+                plt.title(f'{t} {sev}')
+                plt.xticks([], [])
+                plt.yticks([], [])
+        plt.show()
+        plt.savefig(f'results/Cimages/{i+2}_{t}.png')

utils/augmentation.py

@@ -0,0 +1,132 @@
+import random
+import numpy as np
+
+import torch
+import torchvision.transforms as T
+
+
+class RandomRotate90():  # Note: not the same as T.RandomRotation(90)
+    def __call__(self, x):
+        x = x.rot90(random.randint(0, 3), dims=(-1, -2))
+        return x
+
+    def __repr__(self):
+        return self.__class__.__name__
+
+
+class AddGaussianNoise():
+    def __init__(self, std=0.01):
+        self.std = std
+
+    def __call__(self, x):
+        # noise = torch.randn_like(x) * self.std
+        # out = x + noise
+        # debug(x)
+        # debug(noise)
+        # debug(out)
+        return x + torch.randn_like(x) * self.std
+
+    def __repr__(self):
+        return self.__class__.__name__ + f'(std={self.std})'
+
+
+def set_global_seed(seed):
+    torch.random.manual_seed(seed)
+    np.random.seed(seed % (2**32 - 1))
+    random.seed(seed)
+
+
+class ComposeState(T.Compose):
+    def __init__(self, transforms):
+        self.transforms = []
+        self.mask_transforms = []
+
+        for t in transforms:
+            apply_for_mask = True
+            if isinstance(t, tuple):
+                t, apply_for_mask = t
+            self.transforms.append(t)
+            if apply_for_mask:
+                self.mask_transforms.append(t)
+
+        self.seed = None
+
+    # @debug
+    def __call__(self, x, retain_state=False, mask_transform=False):
+        if self.seed is not None:   # retain previous state
+            set_global_seed(self.seed)
+        if retain_state:    # save state for next call
+            self.seed = self.seed or torch.seed()
+            set_global_seed(self.seed)
+        else:
+            self.seed = None    # reset / ignore state
+
+        transforms = self.transforms if not mask_transform else self.mask_transforms
+        for t in transforms:
+            x = t(x)
+        return x
+
+
+augmentation_weak = ComposeState([
+    T.RandomHorizontalFlip(),
+    T.RandomVerticalFlip(),
+    RandomRotate90(),
+])
+
+
+augmentation_strong = ComposeState([
+    T.RandomHorizontalFlip(p=0.5),
+    T.RandomVerticalFlip(p=0.5),
+    T.RandomApply([T.RandomRotation(90)], p=0.5),
+    # (transform, apply_to_mask=True)
+    (T.RandomApply([AddGaussianNoise(std=0.0005)], p=0.5), False),
+    (T.RandomAdjustSharpness(0.5, p=0.5), False),
+])
+
+
+def get_augmentation(type):
+    if type == 'none':
+        return None
+    if type == 'weak':
+        return augmentation_weak
+    if type == 'strong':
+        return augmentation_strong
+
+
+if __name__ == '__main__':
+    import os
+    if not os.path.exists('README.md'):
+        os.chdir('..')
+
+    # from utils.debug import debug
+    from utils.dataset import get_dataset
+    import matplotlib.pyplot as plt
+
+    dataset = get_dataset('DS')  # drone segmentation
+    img, mask = dataset[10]
+    mask = (mask + 0.2) / 1.2
+
+    plt.figure(figsize=(14, 8))
+    plt.subplot(121)
+    plt.imshow(img)
+    plt.subplot(122)
+    plt.imshow(mask)
+    plt.suptitle('no augmentation')
+    plt.show()
+
+    from utils.base import np2torch, torch2np
+    img, mask = np2torch(img), np2torch(mask)
+
+    # from utils.augmentation import get_augmentation
+    augmentation = get_augmentation('strong')
+
+    set_global_seed(1)
+
+    for i in range(1, 4):
+        plt.figure(figsize=(14, 8))
+        plt.subplot(121)
+        plt.imshow(torch2np(augmentation(img.unsqueeze(0), retain_state=True)).squeeze())
+        plt.subplot(122)
+        plt.imshow(torch2np(augmentation(mask.unsqueeze(0), mask_transform=True)).squeeze())
+        plt.suptitle(f'augmentation test {i}')
+        plt.show()

utils/base.py

@@ -0,0 +1,330 @@
+"""
+Utilities for other scripts
+"""
+
+import os
+import shutil
+
+import random
+
+import torch
+import mlflow
+from mlflow.tracking import MlflowClient
+import numpy as np
+
+from IPython.display import display, Markdown
+
+from b2sdk.v1 import *
+
+import argparse
+
+
+class SmartFormatter(argparse.HelpFormatter):
+        
+    def _split_lines(self, text, width):
+        if text.startswith('R|'):
+            return text[2:].splitlines()  
+        # this is the RawTextHelpFormatter._split_lines
+        return argparse.HelpFormatter._split_lines(self, text, width)
+
+
+def str2bool(string):
+    return string == 'True'
+
+
+def np2torch(nparray):
+    """Convert numpy array to torch tensor
+       For array with more than 3 channels, it is better to use an input array in the format BxHxWxC
+
+       Args:
+           numpy array (ndarray) BxHxWxC
+       Returns:
+           torch tensor (tensor) BxCxHxW"""
+
+    tensor = torch.Tensor(nparray)
+
+    if tensor.ndim == 2:
+        return tensor
+    if tensor.ndim == 3:
+        height, width, channels = tensor.shape
+        if channels <= 3:  # Single image with more channels (HxWxC)
+            return tensor.permute(2, 0, 1)
+
+    if tensor.ndim == 4:  # More images with more channels (BxHxWxC)
+        return tensor.permute(0, 3, 1, 2)
+
+    return tensor
+
+
+def torch2np(torchtensor):
+    """Convert torch tensor to numpy array 
+       For tensor with more than 3 channels or batch, it is better to use an input tensor in the format BxCxHxW
+
+       Args:
+           torch tensor (tensor) BxCxHxW
+       Returns:
+           numpy array (ndarray) BxHxWxC"""
+
+    ndarray = torchtensor.detach().cpu().numpy().astype(np.float32)
+
+    if ndarray.ndim == 3:  # Single image with more channels (CxHxW)
+        channels, height, width = ndarray.shape
+        if channels <= 3:
+            return ndarray.transpose(1, 2, 0)
+
+    if ndarray.ndim == 4:  # More images with more channels (BxCxHxW)
+        return ndarray.transpose(0, 2, 3, 1)
+
+    return ndarray
+
+
+def set_random_seed(seed):
+    np.random.seed(seed)  # cpu vars
+    torch.manual_seed(seed)  # cpu  vars
+    random.seed(seed)  # Python
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)  # gpu vars
+        torch.backends.cudnn.deterministic = True  # needed
+        torch.backends.cudnn.benchmark = False
+
+
+def normalize(img):
+    """Normalize images
+
+       Args:
+            imgs (ndarray): image to normalize --> size: (Height,Width,Channels)
+       Returns:
+            normalized (ndarray): normalized image
+            mu (ndarray): mean
+            sigma (ndarray): standard deviation
+    """
+
+    img = img.astype(float)
+
+    if len(img.shape) == 2:
+        img = img[:, :, np.newaxis]
+
+    height, width, channels = img.shape
+
+    mu, sigma = np.empty(channels), np.empty(channels)
+
+    for ch in range(channels):
+        temp_mu = img[:, :, ch].mean()
+        temp_sigma = img[:, :, ch].std()
+
+        img[:, :, ch] = (img[:, :, ch] - temp_mu) / (temp_sigma + 1e-4)
+
+        mu[ch] = temp_mu
+        sigma[ch] = temp_sigma
+
+    return img, mu, sigma
+
+
+def b2_list_files(folder=''):
+    bucket = get_b2_bucket()
+    for file_info, _ in bucket.ls(folder, show_versions=False):
+        print(file_info.file_name)
+
+
+def get_b2_bucket():
+    bucket_name = 'perturbed-minds'
+    application_key_id = '003d6b042de536a0000000004'
+    application_key = 'K003E5Cr+BAYlvSHfg2ynLtvS5aNq78'
+    info = InMemoryAccountInfo()
+    b2_api = B2Api(info)
+    b2_api.authorize_account('production', application_key_id, application_key)
+    bucket = b2_api.get_bucket_by_name(bucket_name)
+    return bucket
+
+
+def b2_download_folder(b2_dir, local_dir, force_download=False, mirror_folder=True):
+    """Downloads a folder from the b2 bucket and optionally cleans
+    up files that are no longer on the server
+
+    Args:
+        b2_dir (str): path to folder on the b2 server
+        local_dir (str): path to folder on the local machine
+        force_download (bool, optional): force the download, if set to `False`, 
+            files with matching names on the local machine will be skipped
+        mirror_folder (bool, optional): if set to `True`, files that are found in
+            the local directory, but are not on the server will be deleted
+    """
+    bucket = get_b2_bucket()
+
+    if not os.path.exists(local_dir):
+        os.makedirs(local_dir)
+    elif not force_download:
+        return
+
+    download_files = [file_info.file_name.split(b2_dir + '/')[-1]
+                      for file_info, _ in bucket.ls(b2_dir, show_versions=False)]
+
+    for file_name in download_files:
+        if file_name.endswith('/.bzEmpty'):  # subdirectory, download recursively
+            subdir = file_name.replace('/.bzEmpty', '')
+            if len(subdir) > 0:
+                b2_subdir = os.path.join(b2_dir, subdir)
+                local_subdir = os.path.join(local_dir, subdir)
+                if b2_subdir != b2_dir:
+                    b2_download_folder(b2_subdir, local_subdir, force_download=force_download,
+                                       mirror_folder=mirror_folder)
+        else:   # file
+            b2_file = os.path.join(b2_dir, file_name)
+            local_file = os.path.join(local_dir, file_name)
+            if not os.path.exists(local_file) or force_download:
+                print(f"downloading b2://{b2_file} -> {local_file}")
+                bucket.download_file_by_name(b2_file, DownloadDestLocalFile(local_file))
+
+    if mirror_folder:   # remove all files that are not on the b2 server anymore
+        for i, file in enumerate(download_files):
+            if file.endswith('/.bzEmpty'):  # subdirectory, download recursively
+                download_files[i] = file.replace('/.bzEmpty', '')
+        for file_name in os.listdir(local_dir):
+            if file_name not in download_files:
+                local_file = os.path.join(local_dir, file_name)
+                print(f"deleting {local_file}")
+                if os.path.isdir(local_file):
+                    shutil.rmtree(local_file)
+                else:
+                    os.remove(local_file)
+
+
+def get_name(obj):
+    return obj.__name__ if hasattr(obj, '__name__') else type(obj).__name__
+
+
+def get_mlflow_model_by_name(experiment_name, run_name, 
+                    tracking_uri = "http://deplo-mlflo-1ssxo94f973sj-890390d809901dbf.elb.eu-central-1.amazonaws.com",
+                    download_model = True):
+    
+    # 0. mlflow basics
+    mlflow.set_tracking_uri(tracking_uri)
+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
+
+    # # 1. use get_experiment_by_name to get experiment objec
+    experiment = mlflow.get_experiment_by_name(experiment_name)
+
+    # # 2. use search_runs with experiment_id for string search query
+    if os.path.isfile('cache/runs_names.pkl'):
+        runs = pd.read_pickle('cache/runs_names.pkl')
+        if runs['tags.mlflow.runName'][runs['tags.mlflow.runName'] == run_name].empty:
+            runs = fetch_runs_list_mlflow(experiment) #returns a pandas data frame where each row is a run (if several exist under that name) 
+    else: 
+        runs = fetch_runs_list_mlflow(experiment) #returns a pandas data frame where each row is a run (if several exist under that name) 
+
+    # 3. get the selected run between all runs inside the selected experiment
+    run = runs.loc[runs['tags.mlflow.runName'] == run_name]
+
+    # 4. check if there is only a run with that name
+    assert len(run) == 1, "More runs with this name"
+    index_run = run.index[0]
+    artifact_uri = run.loc[index_run, 'artifact_uri']
+
+    # 5. load state_dict of your run
+    state_dict = mlflow.pytorch.load_state_dict(artifact_uri)
+
+    # 6. load model of your run
+    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+    # model = mlflow.pytorch.load_model(os.path.join(
+    #         artifact_uri, "model"), map_location=torch.device(DEVICE))
+    model = fetch_from_mlflow(os.path.join(
+                artifact_uri, "model"), use_cache=True, download_model=download_model)
+
+    return state_dict, model
+
+def data_loader_mean_and_std(data_loader, transform=None):
+    means = []
+    stds = []
+    for x, y in data_loader:
+        if transform is not None:
+            x = transform(x)
+        means.append(x.mean(dim=(0, 2, 3)).unsqueeze(0))
+        stds.append(x.std(dim=(0, 2, 3)).unsqueeze(0))
+    return torch.cat(means).mean(dim=0), torch.cat(stds).mean(dim=0)
+
+def fetch_runs_list_mlflow(experiment):        
+    runs = mlflow.search_runs(experiment.experiment_id)
+    runs.to_pickle('cache/runs_names.pkl')  # where to save it, usually as a .pkl
+    return runs
+
+def fetch_from_mlflow(uri, use_cache=True, download_model=True):
+    cache_loc = os.path.join('cache', uri.split('//')[1]) + '.pt'
+    if use_cache and os.path.exists(cache_loc):
+        print(f'loading cached model from {cache_loc} ...')
+        return torch.load(cache_loc)
+    else:
+        print(f'fetching model from {uri} ...')
+        model = mlflow.pytorch.load_model(uri)
+        os.makedirs(os.path.dirname(cache_loc), exist_ok=True)
+        if download_model:
+            torch.save(model, cache_loc, pickle_module=mlflow.pytorch.pickle_module)
+        return model
+
+
+def display_mlflow_run_info(run):
+    uri = mlflow.get_tracking_uri()
+    experiment_id = run.info.experiment_id
+    experiment_name = mlflow.get_experiment(experiment_id).name
+    run_id = run.info.run_id
+    run_name = run.data.tags['mlflow.runName']
+    experiment_url = f'{uri}/#/experiments/{experiment_id}'
+    run_url = f'{experiment_url}/runs/{run_id}'
+
+    print(f'view results at {run_url}')
+    display(Markdown(
+        f"[<a href='{experiment_url}'>experiment {experiment_id} '{experiment_name}'</a>]"
+        f" > "
+        f"[<a href='{run_url}'>run '{run_name}' {run_id}</a>]"
+    ))
+    print('')
+
+
+def get_train_test_indices_drone(df, frac, seed=None):
+    """ Split indices of a DataFrame with binary and balanced labels into balanced subindices
+
+   Args:
+        df (pd.DataFrame): {0,1}-labeled data
+        frac (float): fraction of indicies in first subset
+        random_seed (int): random seed used as random state in np.random and as argument for random.seed()
+   Returns:
+       train_indices (torch.tensor): balanced subset of indices corresponding to rows in the DataFrame
+       test_indices (torch.tensor): balanced subset of indices corresponding to rows in the DataFrame
+    """
+
+    split_idx = int(len(df) * frac / 2)
+    df_with = df[df['label'] == 1]
+    df_without = df[df['label'] == 0]
+
+    np.random.seed(seed)
+    df_with_train = df_with.sample(n=split_idx, random_state=seed)
+    df_with_test = df_with.drop(df_with_train.index)
+
+    df_without_train = df_without.sample(n=split_idx, random_state=seed)
+    df_without_test = df_without.drop(df_without_train.index)
+
+    train_indices = list(df_without_train.index) + list(df_with_train.index)
+    test_indices = list(df_without_test.index) + list(df_with_test.index)
+
+    """"
+    print('fraction of 1-label in train set: {}'.format(len(df_with_train)/(len(df_with_train) + len(df_without_train))))
+    print('fraction of 1-label in test set: {}'.format(len(df_with_test)/(len(df_with_test) + len(df_with_test))))
+    """
+
+    return train_indices, test_indices
+
+
+def smp_get_loss(loss):
+    if loss == "Dice":
+        return smp.losses.DiceLoss(mode='binary', from_logits=True)
+    if loss == "BCE":
+        return nn.BCELoss()
+    elif loss == "BCEWithLogits":
+        return smp.losses.BCEWithLogitsLoss()
+    elif loss == "DicyBCE":
+        from pytorch_toolbelt import losses as ptbl
+        return ptbl.JointLoss(ptbl.DiceLoss(mode='binary', from_logits=False),
+                              nn.BCELoss(),
+                              first_weight=args.dice_weight,
+                              second_weight=args.bce_weight)

utils/dataset.py

@@ -0,0 +1,622 @@
+import os
+import shutil
+import rawpy
+import random
+from PIL import Image
+import tifffile as tiff
+import zipfile
+
+import numpy as np
+import pandas as pd
+
+from torch.utils.data import Dataset, DataLoader, TensorDataset
+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.base import np2torch, torch2np, b2_download_folder
+
+IMAGE_FILE_TYPES = ['dng', 'png', 'tif', 'tiff']
+
+
+def get_dataset(name, I_ratio=1.0):
+    # DroneDataset
+    if name in ('DC', 'Drone', 'DroneClassification', 'DroneDatasetClassificationTiled'):
+        return DroneDatasetClassificationTiled(I_ratio=I_ratio)
+    if name in ('DS', 'DroneSegmentation', 'DroneDatasetSegmentationTiled'):
+        return DroneDatasetSegmentationTiled(I_ratio=I_ratio)
+
+    # MicroscopyDataset
+    if name in ('M', 'Microscopy', 'MicroscopyDataset'):
+        return MicroscopyDataset(I_ratio=I_ratio)
+
+    # for testing
+    if name in ('DSF', 'DroneDatasetSegmentationFull'):
+        return DroneDatasetSegmentationFull(I_ratio=I_ratio)
+    if name in ('MRGB', 'MicroscopyRGB', 'MicroscopyDatasetRGB'):
+        return MicroscopyDatasetRGB(I_ratio=I_ratio)
+
+    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.
+    Files are expected to be of the same filetype.
+
+    Args:
+        img_dir (str): path to image folder
+        mask_dir (str): path to mask folder
+        transform (callable, optional): transformation to apply to image and mask
+        bits (int, optional): normalize image by dividing by 2^bits - 1
+    """
+
+    task = 'classification'
+
+    def __init__(self, img_dir, labels, transform=None, bits=1):
+
+        self.img_dir = img_dir
+        self.labels = labels
+
+        self.images = list_images_in_dir(img_dir)
+
+        assert len(self.images) == len(self.labels)
+
+        self.transform = transform
+        self.bits = bits
+
+    def __repr__(self):
+        rep = f"{type(self).__name__}: ImageFolderDataset[{len(self.images)}]"
+        for n, (img, label) in enumerate(zip(self.images, self.labels)):
+            rep += f'\nimage: {img}\tlabel: {label}'
+            if n > 10:
+                rep += '\n...'
+                break
+        return rep
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+
+        label = self.labels[idx]
+
+        img = load_image(self.images[idx])
+        img = img / (2**self.bits - 1)
+        if self.transform is not None:
+            img = self.transform(img)
+
+        if len(img.shape) == 2:
+            assert img.shape == (256, 256), f"Invalid size for {self.images[idx]}"
+        else:
+            assert img.shape == (3, 256, 256), f"Invalid size for {self.images[idx]}"
+
+        return img, label
+
+
+class ImageFolderDatasetSegmentation(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.
+    Files are expected to be of the same filetype.
+
+    Args:
+        img_dir (str): path to image folder
+        mask_dir (str): path to mask folder
+        transform (callable, optional): transformation to apply to image and mask
+        bits (int, optional): normalize image by dividing by 2^bits - 1
+    """
+
+    task = 'segmentation'
+
+    def __init__(self, img_dir, mask_dir, transform=None, bits=1):
+
+        self.img_dir = img_dir
+        self.mask_dir = mask_dir
+
+        self.images = list_images_in_dir(img_dir)
+        self.masks = list_images_in_dir(mask_dir)
+
+        check_image_folder_consistency(self.images, self.masks)
+
+        self.transform = transform
+        self.bits = bits
+
+    def __repr__(self):
+        rep = f"{type(self).__name__}: ImageFolderDatasetSegmentation[{len(self.images)}]"
+        for n, (img, mask) in enumerate(zip(self.images, self.masks)):
+            rep += f'\nimage: {img}\tmask: {mask}'
+            if n > 10:
+                rep += '\n...'
+                break
+        return rep
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+
+        img = load_image(self.images[idx])
+        mask = load_image(self.masks[idx])
+
+        img = img / (2**self.bits - 1)
+        mask = (mask > 0).astype(np.float32)
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        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
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __repr__(self):
+        return f"Subset [{len(self.dataset)}] of " + repr(self.dataset)
+
+    def __getitem__(self, idx):
+        x, y = self.dataset[idx]
+        if self.transform is not None:
+            x = self.transform(x)
+        return x, y
+
+class Subset(Dataset):
+    """Define a subset of a dataset by only selecting given indices.
+
+    Args:
+        dataset (Dataset): full dataset
+        indices (list): subset indices
+    """
+
+    def __init__(self, dataset, indices=None, transform=None):
+        self.dataset = dataset
+        self.indices = indices if indices is not None else range(len(dataset))
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.indices)
+
+    def __repr__(self):
+        return f"Subset [{len(self)}] of " + repr(self.dataset)
+
+    def __getitem__(self, idx):
+        x, y = self.dataset[self.indices[idx]]
+        if self.transform is not None:
+            x = self.transform(x)
+        return x, y
+
+
+class DroneDatasetSegmentationFull(ImageFolderDatasetSegmentation):
+    """Dataset consisting of full-sized numpy images and masks. Images are normalized to range [0, 1].
+    """
+
+    black_level = [0.0625, 0.0626, 0.0625, 0.0626]
+    white_balance = [2.86653646, 1., 1.73079425]
+    colour_matrix = [1.50768983, -0.33571374, -0.17197604, -0.23048614,
+                     1.70698738, -0.47650126, -0.03119153, -0.32803956, 1.35923111]
+    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] 
+
+        img_dir = f'data/drone/images_full/raw_scale{int(I_ratio*100):03d}'
+        mask_dir = 'data/drone/masks_full'
+
+        download_drone_dataset(force_download)  # XXX: zip files and add checksum? date?
+
+        super().__init__(img_dir=img_dir, mask_dir=mask_dir, transform=transform, bits=bits)
+
+
+class DroneDatasetSegmentationTiled(ImageFolderDatasetSegmentation):
+    """Dataset consisting of tiled numpy images and masks. Images are in range [0, 1]
+    Args:
+        tile_size (int, optional): size of the tiled images. Defaults to 256.
+    """
+
+    camera_parameters = DroneDatasetSegmentationFull.camera_parameters
+
+    def __init__(self, I_ratio=1.0, transform=None):
+
+        tile_size = 256
+
+        img_dir = f'data/drone/images_tiles_{tile_size}/raw_scale{int(I_ratio*100):03d}'
+        mask_dir = f'data/drone/masks_tiles_{tile_size}'
+
+        if not os.path.exists(img_dir) or not os.path.exists(mask_dir):
+            dataset_full = DroneDatasetSegmentationFull(I_ratio=I_ratio, bits=1)
+            print("tiling dataset..")
+            create_tiles_dataset(dataset_full, img_dir, mask_dir, tile_size=tile_size)
+
+        super().__init__(img_dir=img_dir, mask_dir=mask_dir, transform=transform, bits=16)
+
+
+class DroneDatasetClassificationTiled(ImageFolderDataset):
+
+    camera_parameters = DroneDatasetSegmentationFull.camera_parameters
+
+    def __init__(self, I_ratio=1.0, transform=None):
+
+        random_state = 72
+        tile_size = 256
+        thr = 0.01
+
+        img_dir = f'data/drone/classification/images_tiles_{tile_size}/raw_scale{int(I_ratio*100):03d}_thr_{thr}'
+        mask_dir = f'data/drone/classification/masks_tiles_{tile_size}_thr_{thr}'
+        df_path = f'data/drone/classification/dataset_tiles_{tile_size}_{random_state}_{thr}.csv'
+
+        if not os.path.exists(img_dir) or not os.path.exists(mask_dir):
+            dataset_full = DroneDatasetSegmentationFull(I_ratio=I_ratio, bits=1)
+            print("tiling dataset..")
+            create_tiles_dataset_binary(dataset_full, img_dir, mask_dir, random_state, thr, tile_size=tile_size)
+
+        self.classes = ['car', 'no car']
+        self.df = pd.read_csv(df_path)
+        labels = self.df['label'].to_list()
+
+        super().__init__(img_dir=img_dir, labels=labels, transform=transform, bits=16)
+
+        images, class_labels = read_label_csv(self.df)
+        self.images = [os.path.join(self.img_dir, image) for image in images]
+        self.labels = class_labels
+
+
+class MicroscopyDataset(ImageFolderDataset):
+    """MicroscopyDataset raw images
+
+    Args:
+        I_ratio (float): Original image rescaled by this factor, possible values [0.01,0.05,0.1,0.25,0.5,0.75,1.0]
+        raw (bool): Select rgb dataset or raw dataset
+        transform (callable, optional): transformation to apply to image and mask
+        bits (int, optional): normalize image by dividing by 2^bits - 1
+    """
+
+    black_level = [9.834368023181512e-06, 9.834368023181512e-06, 9.834368023181512e-06, 9.834368023181512e-06]
+    white_balance = [-0.6567, 1.9673, 3.5304]
+    colour_matrix = [-2.0338, 0.0933, 0.4157, -0.0286, 2.6464, -0.0574, -0.5516, -0.0947, 2.9308]
+
+    camera_parameters = black_level, white_balance, colour_matrix
+
+    dataset_mean = [0.91, 0.84, 0.94]
+    dataset_std = [0.08, 0.12, 0.05]
+
+    def __init__(self, I_ratio=1.0, transform=None, bits=16, force_download=False):
+
+        assert I_ratio in [0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 1.0]
+
+        download_microscopy_dataset(force_download=force_download)
+
+        self.img_dir = f'data/microscopy/images/raw_scale{int(I_ratio*100):03d}'
+        self.transform = transform
+        self.bits = bits
+
+        self.label_file = 'data/microscopy/labels/Ma190c_annotations.dat'
+
+        self.valid_classes = ['BAS', 'EBO', 'EOS', 'KSC', 'LYA', 'LYT', 'MMZ', 'MOB',
+                              'MON', 'MYB', 'MYO', 'NGB', 'NGS', 'PMB', 'PMO', 'UNC']
+
+        self.invalid_files = ['Ma190c_lame3_zone13_composite_Mcropped_2.tiff', ]
+
+        images, class_labels = read_label_file(self.label_file)
+
+        # filter classes with low appearance
+        self.valid_classes = [class_label for class_label in self.valid_classes
+                              if class_labels.count(class_label) > 4]
+
+        # remove invalid classes and invalid files from (images, class_labels)
+        images, class_labels = list(zip(*[
+            (image, class_label)
+            for image, class_label in zip(images, class_labels)
+            if class_label in self.valid_classes and image not in self.invalid_files
+        ]))
+
+        self.classes = list(sorted({*class_labels}))
+
+        # store full path
+        self.images = [os.path.join(self.img_dir, image) for image in images]
+
+        # reindex labels
+        self.labels = [self.classes.index(class_label) for class_label in class_labels]
+
+
+class MicroscopyDatasetRGB(MicroscopyDataset):
+    """MicroscopyDataset RGB images
+
+    Args:
+        I_ratio (float): Original image rescaled by this factor, possible values [0.01,0.05,0.1,0.25,0.5,0.75,1.0]
+        raw (bool): Select rgb dataset or raw dataset
+        transform (callable, optional): transformation to apply to image and mask
+        bits (int, optional): normalize image by dividing by 2^bits - 1
+    """
+    camera_parameters = None
+
+    dataset_mean = None
+    dataset_std = None
+
+    def __init__(self, I_ratio=1.0, transform=None, bits=16, force_download=False):
+        super().__init__(I_ratio=I_ratio, transform=transform, bits=bits, force_download=force_download)
+        self.images = [image.replace('raw', 'rgb') for image in self.images]  # XXX: hack
+
+
+def read_label_file(label_file_path):
+
+    images = []
+    class_labels = []
+
+    with open(label_file_path, "rb") as data:
+        for line in data:
+            file_name, class_label = line.decode("utf-8").split()
+            image = file_name + '.tiff'
+            images.append(image)
+            class_labels.append(class_label)
+
+    return images, class_labels
+
+
+def read_label_csv(df):
+
+    images = []
+    class_labels = []
+
+    for file_name, label in zip(df['file name'], df['label']):
+        image = file_name + '.tif'
+        images.append(image)
+        class_labels.append(int(label))
+    return images, class_labels
+
+
+def download_drone_dataset(force_download):
+    b2_download_folder('drone/images', 'data/drone/images_full', force_download=force_download)
+    b2_download_folder('drone/masks', 'data/drone/masks_full', force_download=force_download)
+    unzip_drone_images()
+
+
+def download_microscopy_dataset(force_download):
+    b2_download_folder('Data histopathology/WhiteCellsImages',
+                       'data/microscopy/images', force_download=force_download)
+    b2_download_folder('Data histopathology/WhiteCellsLabels',
+                       'data/microscopy/labels', force_download=force_download)
+    unzip_microscopy_images()
+
+
+def unzip_microscopy_images():
+            
+    if os.path.isfile('data/microscopy/labels/.bzEmpty'):
+        os.remove('data/microscopy/labels/.bzEmpty')
+
+    for file in os.listdir('data/microscopy/images'):
+        if file.endswith(".zip"):
+            zip = zipfile.ZipFile(os.path.join('data/microscopy/images', file))
+            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'):
+        os.remove('data/drone/masks_full/.bzEmpty')
+
+    for file in os.listdir('data/drone/images_full'):
+        if file.endswith(".zip"):
+            zip = zipfile.ZipFile(os.path.join('data/drone/images_full', file))
+            zip.extractall('data/drone/images_full')
+            os.remove(os.path.join('data/drone/images_full', file))
+
+
+def create_tiles_dataset(dataset, img_dir, mask_dir, tile_size=256):
+    for folder in [img_dir, mask_dir]:
+        if not os.path.exists(folder):
+            os.makedirs(folder)
+    for n, (img, mask) in enumerate(dataset):
+        tiled_img = split_img(img, ROIs=(tile_size, tile_size), step=(tile_size, tile_size))
+        tiled_mask = split_img(mask, ROIs=(tile_size, tile_size), step=(tile_size, tile_size))
+        tiled_img, tiled_mask = class_detection(tiled_img, tiled_mask)  # Remove images without cars in it
+        for i, (sub_img, sub_mask) in enumerate(zip(tiled_img, tiled_mask)):
+            tile_id = f"{n:02d}_{i:05d}"
+            Image.fromarray(sub_img).save(os.path.join(img_dir, tile_id + '.tif'))
+            Image.fromarray(sub_mask > 0).save(os.path.join(mask_dir, tile_id + '.png'))
+
+
+def create_tiles_dataset_binary(dataset, img_dir, mask_dir, random_state, thr, tile_size=256):
+
+    for folder in [img_dir, mask_dir]:
+        if not os.path.exists(folder):
+            os.makedirs(folder)
+
+    ids = []
+    labels = []
+
+    for n, (img, mask) in enumerate(dataset):
+        tiled_img = split_img(img, ROIs=(tile_size, tile_size), step=(tile_size, tile_size))
+        tiled_mask = split_img(mask, ROIs=(tile_size, tile_size), step=(tile_size, tile_size))
+
+        X_with, X_without, Y_with, Y_without = binary_class_detection(
+            tiled_img, tiled_mask, random_state, thr)  # creates balanced arrays with class and without class
+
+        for i, (sub_X_with, sub_Y_with) in enumerate(zip(X_with, Y_with)):
+            tile_id = f"{n:02d}_{i:05d}"
+            ids.append(tile_id)
+            labels.append(0)
+            Image.fromarray(sub_X_with).save(os.path.join(img_dir, tile_id + '.tif'))
+            Image.fromarray(sub_Y_with > 0).save(os.path.join(mask_dir, tile_id + '.png'))
+        for j, (sub_X_without, sub_Y_without) in enumerate(zip(X_without, Y_without)):
+            tile_id = f"{n:02d}_{i+1+j:05d}"
+            ids.append(tile_id)
+            labels.append(1)
+            Image.fromarray(sub_X_without).save(os.path.join(img_dir, tile_id + '.tif'))
+            Image.fromarray(sub_Y_without > 0).save(os.path.join(mask_dir, tile_id + '.png'))
+           # Image.fromarray(sub_mask).save(os.path.join(mask_dir, tile_id + '.png'))
+
+    df = pd.DataFrame({'file name': ids, 'label': labels})
+
+    df_loc = f'data/drone/classification/dataset_tiles_{tile_size}_{random_state}_{thr}.csv'
+    df.to_csv(df_loc)
+
+    return
+
+
+def class_detection(X, Y):
+    """Split dataset in images which has the class in the target
+
+       Args:
+            X (ndarray): input image
+            Y (ndarray): target with segmentation map (images with {0,1} values where it is 1 when there is the class)
+       Returns:
+           X_with_class (ndarray): input regions with the selected class 
+           Y_with_class (ndarray): target regions with the selected class 
+           X_without_class (ndarray): input regions without the selected class 
+           Y_without_class (ndarray): target regions without the selected class 
+    """
+
+    with_class = []
+    without_class = []
+    for i, img in enumerate(Y):
+        if img.mean() == 0:
+            without_class.append(i)
+        else:
+            with_class.append(i)
+
+    X_with_class = np.delete(X, without_class, 0)
+    Y_with_class = np.delete(Y, without_class, 0)
+
+    return X_with_class, Y_with_class
+
+
+def binary_class_detection(X, Y, random_seed, thr):
+    """Splits subimages in subimages with the selected class and without the selected class by calculating the mean of the submasks; subimages with 0 < submask.mean()<=thr are disregared
+
+
+
+       Args:
+            X (ndarray): input image
+            Y (ndarray): target with segmentation map (images with {0,1} values where it is 1 when there is the class)
+            thr (flaot): sub images are not considered if 0 < sub_target.mean() <= thr 
+            balanced (bool): number of returned sub images is equal for both classes if true 
+            random_seed (None or int): selection of sub images in class with more elements according to random_seed if balanced
+       Returns:
+           X_with_class (ndarray): input regions with the selected class 
+           Y_with_class (ndarray): target regions with the selected class 
+           X_without_class (ndarray): input regions without the selected class 
+           Y_without_class (ndarray): target regions without the selected class 
+    """
+
+    with_class = []
+    without_class = []
+    no_class = []
+
+    for i, img in enumerate(Y):
+        m = img.mean()
+        if m == 0:
+            without_class.append(i)
+        else:
+            if m > thr:
+                with_class.append(i)
+            else:
+                no_class.append(i)
+
+    N = len(with_class)
+    M = len(without_class)
+    random.seed(random_seed)
+    if N <= M:
+        random.shuffle(without_class)
+        with_class.extend(without_class[:M - N])
+    else:
+        random.shuffle(with_class)
+        without_class.extend(with_class[:N - M])
+
+    X_with_class = np.delete(X, without_class + no_class, 0)
+    X_without_class = np.delete(X, with_class + no_class, 0)
+    Y_with_class = np.delete(Y, without_class + no_class, 0)
+    Y_without_class = np.delete(Y, with_class + no_class, 0)
+
+    return X_with_class, X_without_class, Y_with_class, Y_without_class
+
+
+def make_dataloader(dataset, batch_size, shuffle=True):
+
+    X, Y = dataset
+
+    X, Y = np2torch(X), np2torch(Y)
+
+    dataset = TensorDataset(X, Y)
+    dataset = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
+
+    return dataset
+
+
+def check_image_folder_consistency(images, masks):
+    file_type_images = images[0].split('.')[-1].lower()
+    file_type_masks = masks[0].split('.')[-1].lower()
+    assert len(images) == len(masks), "images / masks length mismatch"
+    for img_file, mask_file in zip(images, masks):
+        img_name = img_file.split('/')[-1].split('.')[0]
+        assert img_name in mask_file, f"image {img_file} corresponds to {mask_file}?"
+        assert img_file.split('.')[-1].lower() == file_type_images, \
+            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

utils/debug.py

@@ -0,0 +1,371 @@
+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

@@ -0,0 +1,193 @@
+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()

utils/pytorch_ssim.py

@@ -0,0 +1,75 @@
+"""https://github.com/Po-Hsun-Su/pytorch-ssim"""
+
+import torch
+import torch.nn.functional as F
+from torch.autograd import Variable
+import numpy as np
+from math import exp
+
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+    return gauss/gauss.sum()
+
+def create_window(window_size, channel):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
+    return window
+
+def _ssim(img1, img2, window, window_size, channel, size_average = True):
+    mu1 = F.conv2d(img1, window, padding = window_size//2, groups = channel)
+    mu2 = F.conv2d(img2, window, padding = window_size//2, groups = channel)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1*mu2
+
+    sigma1_sq = F.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq
+    sigma2_sq = F.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq
+    sigma12 = F.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2
+
+    C1 = 0.01**2
+    C2 = 0.03**2
+
+    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
+
+    if size_average:
+        return ssim_map.mean()
+    else:
+        return ssim_map.mean(1).mean(1).mean(1)
+
+class SSIM(torch.nn.Module):
+    def __init__(self, window_size = 11, size_average = True):
+        super(SSIM, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.channel = 1
+        self.window = create_window(window_size, self.channel)
+
+    def forward(self, img1, img2):
+        (_, channel, _, _) = img1.size()
+
+        if channel == self.channel and self.window.data.type() == img1.data.type():
+            window = self.window
+        else:
+            window = create_window(self.window_size, channel)
+            
+            if img1.is_cuda:
+                window = window.cuda(img1.get_device())
+            window = window.type_as(img1)
+            
+            self.window = window
+            self.channel = channel
+
+
+        return _ssim(img1, img2, window, self.window_size, channel, self.size_average)
+
+def ssim(img1, img2, window_size = 11, size_average = True):
+    (_, channel, _, _) = img1.size()
+    window = create_window(window_size, channel)
+    
+    if img1.is_cuda:
+        window = window.cuda(img1.get_device())
+    window = window.type_as(img1)
+    
+    return _ssim(img1, img2, window, window_size, channel, size_average)

utils/show_dataset.ipynb

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

utils/splitting.py

@@ -0,0 +1,137 @@
+"""
+Split images in blocks and vice versa
+"""
+
+import random
+import numpy as np
+
+import torch
+
+from skimage.util.shape import view_as_windows
+
+
+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]))
+        return splitted.reshape((-1, ROIs[0], ROIs[1]))
+    
+    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))
+            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]))
+            return splitted.reshape((-1, ROIs[0], ROIs[1]))
+    
+    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))
+        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]
+        for c in range(1, columns):
+            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)
+    
+    return joined_img
+
+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))
+    
+    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], :]
+    return X_cut, Y_cut
+
+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))
+
+    for i in np.arange(datasize):
+        X_cut[i], Y_cut[i] = random_ROI(X, Y, ROIs)
+    return X_cut, Y_cut