Initial commit

.gitattributes

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

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 aiaudit.org
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,129 @@
+[![MIT License](https://img.shields.io/apm/l/atomic-design-ui.svg?)](https://github.com/tterb/atomic-design-ui/blob/master/LICENSEs) [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.5235536.svg)](https://doi.org/10.5281/zenodo.5235536)
+
+# From Lens to Logit - Addressing Camera Hardware-Drift Using Raw Sensor Data
+
+*This repository hosts the code for the project ["From Lens to Logit: Addressing Camera Hardware-Drift Using Raw Sensor Data"](https://openreview.net/forum?id=DRAywM1BhU), submitted to the NeurIPS 2021 Datasets and Benchmarks Track.*
+
+In order to address camera hardware-drift we require two ingredients: raw sensor data and an image processing model. This code repository contains the materials for the second ingredient, the image processing model, as well as scripts to load lada and run experiments. For a conceptual overview of the project we reocommend the [project site](https://aiaudit.org/lens2logit/) or the [full paper](https://openreview.net/forum?id=DRAywM1BhU).
+
+## A short introduction
+![L2L Overview](https://user-images.githubusercontent.com/38631399/131536063-585cf9b0-e76e-4e41-a05e-2fcf4902f539.png)
+
+
+To create an image, raw sensor data traverses complex image signal processing pipelines. These pipelines are used by cameras and scientific instruments to produce the images fed into machine learning systems. The processing pipelines vary by device, influencing the resulting image statistics and ultimately contributing to what is known as hardware-drift. However, this processing is rarely considered in machine learning modelling, because available benchmark data sets are generally not in raw format. Here we show that pairing qualified raw sensor data with an explicit, differentiable model of the image processing pipeline allows to tackle camera hardware-drift. 
+
+Specifically, we demonstrate 
+1. the **controlled synthesis of hardware-drift test cases**
+2. modular **hardware-drift forensics**, as well as 
+3. **image processing customization**. 
+
+We make available two data sets. 
+1. **Raw-Microscopy**, contains 
+   * **940 raw bright-field microscopy images** of human blood smear slides for leukocyte classification alongside 
+   * **5,640 variations measured at six different intensities** and twelve additional sets totalling 
+   * **11,280 images of the raw sensor data processed through different pipelines**.
+3. **Raw-Drone**, contains 
+   * **548 raw drone camera images** for car segmentation, alongside 
+   * **3,288 variations measured at six different intensities** and also twelve additional sets totalling 
+   * **6,576 images of the raw sensor data processed through different pipelines**.
+## Data access
+If you use our code you can use the convenient cloud storage integration. Data will be loaded automatically from a cloud storage bucket and stored to your working machine. You can find the code snippet doing that [here](https://github.com/aiaudit-org/lens2logit/blob/f8a165a0c094456f68086167f0bef14c3b311a4e/utils/base.py#L130)
+
+```python
+def get_b2_bucket():
+    bucket_name = 'perturbed-minds'
+    application_key_id = '003d6b042de536a0000000008'
+    application_key = 'K003HMNxnoa91Dy9c0V8JVCKNUnwR9U'
+    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
+```
+We also maintain a copy of the entire dataset with a permanent identifier at Zenodo which you can find here [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.5235536.svg)](https://doi.org/10.5281/zenodo.5235536).
+## Code
+### Dependencies
+#### Conda environment and dependencies
+To run this code out-of-the-box you can install the latest project conda environment stored in `environment.yml`
+```console
+$ conda env create -f environment.yml
+```
+#### segmentation_models_pytorch newest version
+We noticed that PyPi package for `segmentation_models_pytorch` is sometimes behind the project's github repository. If you encounter `smp` related problems we reccomend installing directly from the `smp`  reposiroty via
+```console
+$ python -m pip install git+https://github.com/qubvel/segmentation_models.pytorch
+```
+### Recreate experiments
+The central file for using the **Lens2Logit** framework for experiments as in the paper is `train.py` which provides a rich set of arguments to experiment with raw image data, different image processing models and task models for regression or classification. Below we provide three example prompts for the type of experiments reported in the [paper](https://openreview.net/forum?id=DRAywM1BhU)
+#### Controlled synthesis of hardware-drift test cases
+```console
+$ python train.py \
+--experiment_name YOUR-EXPERIMENT-NAME \
+--run_name YOUR-RUN-NAME \
+--dataset Microscopy \
+--lr 1e-5 \
+--n_splits 5 \
+--epochs 5 \
+--classifier_pretrained \
+--processing_mode static \
+--augmentation weak \
+--log_model True \
+--iso 0.01 \
+--freeze_processor \
+--processor_uri "$processor_uri" \
+--track_processing \
+--track_every_epoch \
+--track_predictions \
+--track_processing_gradients \
+--track_save_tensors \
+```
+#### Modular hardware-drift forensics
+```console
+$ python train.py \
+--experiment_name YOUR-EXPERIMENT-NAME \
+--run_name YOUR-RUN-NAME \
+--dataset Microscopy \
+--adv_training
+--lr 1e-5 \
+--n_splits 5 \
+--epochs 5 \
+--classifier_pretrained \
+--processing_mode parametrized \
+--augmentation weak \
+--log_model True \
+--iso 0.01 \
+--track_processing \
+--track_every_epoch \
+--track_predictions \
+--track_processing_gradients \
+--track_save_tensors \
+```
+#### Image processing customization
+```console
+$ python train.py \
+--experiment_name YOUR-EXPERIMENT-NAME \
+--run_name YOUR-RUN-NAME \
+--dataset Microscopy \
+--lr 1e-5 \
+--n_splits 5 \
+--epochs 5 \
+--classifier_pretrained \
+--processing_mode parametrized \
+--augmentation weak \
+--log_model True \
+--iso 0.01 \
+--track_processing \
+--track_every_epoch \
+--track_predictions \
+--track_processing_gradients \
+--track_save_tensors \
+```
+## Virtual lab log
+We maintain a collaborative virtual lab log at [this address](http://deplo-mlflo-1ssxo94f973sj-890390d809901dbf.elb.eu-central-1.amazonaws.com/#/). There you can browse experiment runs, analyze results through SQL queries and download trained processing and task models.
+![mlflow](https://user-images.githubusercontent.com/38631399/131536233-f6b6e0ae-35f2-4ee0-a5e2-d04f8efb8d73.png)
+
+
+### Review our experiments
+Experiments are listed in the left column. You can select individual runs or compare metrics and parameters across different runs. For runs where we tracked images of intermediate processing steps and images of the gradients at these processing steps you can find at the bottom of a run page in the *results* folder for each epoch.
+### Use our trained models
+When selecting a run and a model was saved you can find the model files, state dictionary and instructions to load at the bottom of a run page under *models*. In the menu bar at the top of the virtual lab log you can also access models via the *Model Registry*. Our code is well integrated with the *mlflow* autologging and -loading package for PyTorch. So when using our code you can just specify the *model uri* as an argument and models will be fetched from the model registry automatically.

dataset.py

@@ -0,0 +1,573 @@
+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.dataset_utils import split_img, list_images_in_dir, load_image
+from utils.base import np2torch, torch2np, b2_download_folder
+
+IMAGE_FILE_TYPES = ['dng', 'png', 'tif', 'tiff']
+
+
+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)
+
+
+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}"

demo.py

@@ -0,0 +1,54 @@
+import gradio as gr
+#import tensorflow as tf
+import numpy as np
+import json
+from os.path import dirname, realpath, join
+import processing.pipeline_numpy as ppn
+
+
+# Load human-readable labels for ImageNet.
+current_dir = dirname(realpath(__file__))
+
+    
+def process(RawImage, CameraParameters, Debayer, Sharpening, Denoising):
+    raw_img = RawImage
+    if CameraParameters == "Microscope":
+        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]
+    elif CameraParameters == "Drone":
+        #drone
+        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]
+    else:
+        print("No valid camera parameter")
+    debayer = Debayer
+    sharpening = Sharpening
+    denoising = Denoising
+    print(np.max(raw_img))
+    raw_img = (raw_img[:,:,0].astype(np.float64)/255.)
+    img = ppn.processing(raw_img, black_level, white_balance, colour_matrix,
+                        debayer=debayer, sharpening=sharpening, denoising=denoising)
+    print(np.max(img))
+    return img
+
+
+iface = gr.Interface(
+    process, 
+    [gr.inputs.Image(),gr.inputs.Radio(["Microscope", "Drone"]),gr.inputs.Dropdown(["bilinear", "malvar2004", "menon2007"]), 
+    gr.inputs.Dropdown(["sharpening_filter", "unsharp_masking"]),
+    gr.inputs.Dropdown(["gaussian_denoising", "median_denoising"])], 
+    "image",
+    capture_session=True,
+    examples=[
+        ["demo-files/car.png"],
+        ["demo-files/micro.png"]
+    ],
+    title="Lens2Logit - Static processing demo",
+    description="You can select a sample raw image, the camera parameters and the pipeline configuration to process the raw image.",)
+
+if __name__ == "__main__":
+    iface.launch(share=True)
+

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

figures/ABtesting.py

@@ -0,0 +1,831 @@
+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 dataset import get_dataset, Subset
+from utils.base import get_mlflow_model_by_name, SmartFormatter
+from processing.pipeline_numpy import RawProcessingPipeline
+
+from utils.hendrycks_robustness 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()

figures/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 \

figures/figure2.sh

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

figures/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

figures/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

model.py

@@ -0,0 +1,305 @@
+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,
+                 adv_parameters='all',
+                 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
+
+        self.unfreeze()
+        if freeze_classifier:
+            pl.LightningModule.freeze(self.classifier)
+        if freeze_processor:
+            pl.LightningModule.freeze(self.processor)
+
+        if adv_training and adv_parameters != 'all':
+            if adv_parameters != 'all':
+                pl.LightningModule.freeze(self.processor)
+                for name, p in self.processor.named_parameters():
+                    if adv_parameters in name:
+                        p.requires_grad = True
+
+    def forward(self, x):
+        x = self.processor(x)
+        apply_augmentation_step = self.training or self.augmentation_on_eval
+        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')
+                elif metric_name == 'accuracy' or not self.training or self.metrics_on_training:
+                    m = metric(y_hat.cpu().detach(), y.cpu())
+                    self.log(f'{step_name}_{metric_name}', m, on_step=False, on_epoch=True,
+                             prog_bar=self.training or metric_name == 'accuracy')
+
+        return loss
+
+    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
+
+        # don't update batchnorm in adversarial training
+        self.processor.train(mode=mode and not self.freeze_processor and not self.adv_training)
+        self.classifier.train(mode=mode and not self.freeze_classifier)
+        return self
+
+    def configure_optimizers(self):
+        self.optimizer = torch.optim.Adam(self.parameters(), self.lr, weight_decay=self.weight_decay)
+        return self.optimizer
+
+    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, reference_processor=None, track_every_epoch=False, track_processing=True, track_gradients=True, track_predictions=True, save_tensors=True):
+        super().__init__()
+        self.data_loader = data_loader
+
+        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
+
+        self.reference_processor = reference_processor
+
+    def callback_track_images(self, model, save_loc):
+        track_images(model,
+                     self.data_loader,
+                     reference_processor=self.reference_processor,
+                     track_processing=self.track_processing,
+                     track_gradients=self.track_gradients,
+                     track_predictions=self.track_predictions,
+                     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.model, save_loc)
+
+    def on_train_epoch_end(self, trainer, pl_module, outputs):
+        if self.track_every_epoch:
+            save_loc = f'results/epoch_{trainer.current_epoch + 1:04d}'
+            self.callback_track_images(trainer.model, 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, reference_processor=None, 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)
+    diffs_full = defaultdict(list)
+
+    track_differences = reference_processor is not None
+
+    for inputs, labels in data_loader:
+
+        inputs, labels = inputs.to(device), labels.to(device)
+        inputs.requires_grad = True
+
+        processed_rgb = processor(inputs)
+
+        if track_differences:
+            # debug(processor)
+            processed_rgb_ref = reference_processor(inputs)
+
+        if track_gradients or track_predictions:
+            logits = classifier(processed_rgb)
+
+            # 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_differences:
+                diffs_full[stage].append((reference_processor.stages[stage] - batch).cpu().detach())
+            if track_gradients:
+                grads_full[stage].append(batch.grad.cpu().detach())
+
+    with torch.no_grad():
+
+        stages = stages_full
+        grads = grads_full
+        diffs = diffs_full
+
+        if track_processing:
+            for stage, batch in stages.items():
+                stages[stage] = torch.cat(batch)
+
+        if track_differences:
+            for stage, batch in diffs.items():
+                diffs[stage] = torch.cat(batch)
+
+        if track_gradients:
+            for stage, batch in grads.items():
+                grads[stage] = torch.cat(batch)
+
+        for stage_nr, stage_name in enumerate(stages):
+            if track_processing:
+                batch = stages[stage_name]
+                log_tensor(batch, os.path.join(save_loc, f'processing_{stage_nr}_{stage_name}.pt'), save_tensors)
+
+            if track_differences:
+                batch = diffs[stage_name]
+                log_tensor(batch, os.path.join(save_loc, f'diffs_{stage_nr}_{stage_name}.pt'), False)
+
+            if track_gradients:
+                batch_grad = grads[stage_name]
+                batch_grad = batch_grad.abs()
+                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)

processing/pipeline_numpy.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 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

processing/pipeline_torch.py

@@ -0,0 +1,314 @@
+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 processing.pipeline_numpy 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, affine=False)
+        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
+
+
+def add_additive_layer(processor):
+    processor.additive_layer = nn.Parameter(torch.zeros((1, 3, 256, 256)))
+    # processor.additive_layer = nn.Parameter(0.001 * torch.randn((1, 3, 256, 256)))
+
+
+class ParametrizedProcessing(nn.Module):
+    def __init__(self, camera_parameters, track_stages=False, batch_norm_output=True):
+        super().__init__()
+        self.stages = None
+        self.buffer = None
+        self.track_stages = track_stages
+
+        black_level, white_balance, colour_matrix = camera_parameters
+
+        self.black_level = nn.Parameter(torch.as_tensor(black_level))
+        self.white_balance = nn.Parameter(torch.as_tensor(white_balance).reshape(1, 3))
+        self.colour_correction = nn.Parameter(torch.as_tensor(colour_matrix).reshape(3, 3))
+
+        self.gamma_correct = nn.Parameter(torch.Tensor([2.2]))
+
+        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, affine=False) if batch_norm_output else None
+
+        self.register_buffer('M_RGB_2_YUV', M_RGB_2_YUV.clone())
+        self.register_buffer('M_YUV_2_RGB', M_YUV_2_RGB.clone())
+
+        self.additive_layer = None  # this can be added in later
+
+    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->bchw', rgb, self.white_balance).contiguous()
+        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 + 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 dataset import get_dataset
+    from utils.base import np2torch, torch2np
+
+    from utils.debug import debug
+    from processing.pipeline_numpy 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)

requirements.txt

@@ -0,0 +1,18 @@
+scikit_image
+pandas
+scipy
+numpy
+matplotlib
+b2sdk
+colour_demosaicing
+gradio
+ipython
+mlflow
+Pillow
+pytorch_toolbelt
+rawpy
+scikit_learn
+segmentation_models_pytorch
+tifffile
+torch==1.9.0
+torchvision==0.10.0

train.py

@@ -0,0 +1,426 @@
+import os
+import sys
+import copy
+import argparse
+
+import torch
+from torch import optim
+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.dataset_utils import k_fold
+from utils.augmentation import get_augmentation
+from dataset import Subset, get_dataset
+
+from processing.pipeline_numpy import RawProcessingPipeline
+from processing.pipeline_torch import add_additive_layer, raw2rgb, RawToRGB, ParametrizedProcessing, NNProcessing
+
+from model import log_tensor, resnet_model, LitModel, TrackImagesCallback
+
+import segmentation_models_pytorch as smp
+
+from utils.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("--adv_noise_layer", action='store_true', help="Adds an additive layer to Parametrized Processing")
+parser.add_argument("--adv_track_differences", action='store_true', help='Save difference to default pipeline')
+parser.add_argument('--adv_parameters', choices=['all', 'black_level', 'white_balance',
+                                                 'colour_correction', 'gamma_correct', 'sharpening_filter', 'gaussian_blur', 'additive_layer'])
+
+parser.add_argument("--cache_downloaded_models", type=str2bool, default=True)
+
+parser.add_argument('--test_run', action='store_true')
+
+
+args = parser.parse_args()
+
+os.makedirs('results', exist_ok=True)
+
+
+def run_train(args):
+
+    print(args)
+
+    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+    training_mode = 'adversarial' if args.adv_training else 'default'
+
+    # set tracking uri, this is the address of the mlflow server where light experimental data will be stored
+    mlflow.set_tracking_uri(args.tracking_uri)
+    mlflow.set_experiment(args.experiment_name)
+    os.environ["AWS_ACCESS_KEY_ID"] = "#TODO: fill in your aws access key id for mlflow server here"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "#TODO: fill in your aws secret access key for mlflow server here"
+
+    # 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_noise_layer, # this has to be added manually afterwards for when a model is loaded that doesn't have one yet
+                    )
+
+                elif args.processing_mode == 'neural_network':
+                    processor = NNProcessing(track_stages=track_stages,
+                                             normalize_mosaic=normalize_mosaic, batch_norm_output=True)
+                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 = args.track_processing
+                processor_default.eval()
+                processor_default.to(DEVICE)
+                # debug(processor_default)
+                for p in processor_default.parameters():
+                    p.requires_grad = False
+
+                if args.adv_noise_layer:
+                    add_additive_layer(processor)
+
+                def l2_regularization(x, y):
+                    return ((x - y) ** 2).sum()
+                    # 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):
+                        with torch.no_grad():
+                            x_reference = processor_default(x)
+                        x_processed = processor.buffer['processed_rgb']
+                        return self.weight * self.loss_aux(x_reference, x_processed)
+
+                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=loss, 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,
+                lr=args.lr,
+                loss_aux=loss_aux,
+                adv_training=args.adv_training,
+                adv_parameters=args.adv_parameters,
+                metrics=metrics,
+                augmentation=augmentation,
+                is_segmentation_task=dataset.task == 'segmentation',
+                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
+
+                reference_processor = processor_default if args.adv_training and args.adv_track_differences else None
+
+                callbacks = []
+                if args.track_processing:
+                    callbacks += [TrackImagesCallback(track_loader,
+                                                      reference_processor,
+                                                      track_every_epoch=args.track_every_epoch,
+                                                      track_processing=args.track_processing,
+                                                      track_gradients=args.track_processing_gradients,
+                                                      track_predictions=args.track_predictions,
+                                                      save_tensors=args.track_save_tensors)]
+
+                # if True: #args.save_best:
+                #    if 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,
+                )
+
+    globals().update(locals())  # for convenient access
+
+    return model
+
+
+if __name__ == '__main__':
+    model = run_train(args)

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 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,335 @@
+"""
+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 = '003d6b042de536a0000000008'
+    application_key = 'K003HMNxnoa91Dy9c0V8JVCKNUnwR9U'
+    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:
+            # returns a pandas data frame where each row is a run (if several exist under that name)
+            runs = fetch_runs_list_mlflow(experiment)
+    else:
+        # returns a pandas data frame where each row is a run (if several exist under that name)
+        runs = fetch_runs_list_mlflow(experiment)
+
+    # 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_utils.py

@@ -0,0 +1,190 @@
+
+import random
+import numpy as np
+
+import torch
+
+from skimage.util.shape import view_as_windows
+
+
+def load_image(path):
+    file_type = path.split('.')[-1].lower()
+    if file_type == 'dng':
+        img = rawpy.imread(path).raw_image_visible
+    elif file_type == 'tiff' or file_type == 'tif':
+        img = np.array(tiff.imread(path), dtype=np.float32)
+    else:
+        img = np.array(Image.open(path), dtype=np.float32)
+    return img
+
+
+def list_images_in_dir(path):
+    image_list = [os.path.join(path, img_name)
+                  for img_name in sorted(os.listdir(path))
+                  if img_name.split('.')[-1].lower() in IMAGE_FILE_TYPES]
+    return image_list
+
+
+def k_fold(dataset, n_splits: int, seed: int, train_size: float):
+    """Split dataset in subsets for cross-validation 
+
+       Args:
+            dataset (class): dataset to split
+            n_split (int): Number of re-shuffling & splitting iterations.
+            seed (int): seed for k_fold splitting
+            train_size (float): should be between 0.0 and 1.0 and represent the proportion of the dataset to include in the train split.
+       Returns:
+           idxs (list): indeces for splitting the dataset. The list contain n_split pair of train/test indeces.
+    """
+    if hasattr(dataset, 'labels'):
+        x = dataset.images
+        y = dataset.labels
+    elif hasattr(dataset, 'masks'):
+        x = dataset.images
+        y = dataset.masks
+
+    idxs = []
+
+    if dataset.task == 'classification':
+        sss = StratifiedShuffleSplit(n_splits=n_splits, train_size=train_size, random_state=seed)
+
+        for idxs_train, idxs_test in sss.split(x, y):
+            idxs.append((idxs_train.tolist(), idxs_test.tolist()))
+
+    elif dataset.task == 'segmentation':
+        for n in range(n_splits):
+            split_idx = int(len(dataset) * train_size)
+            indices = np.random.permutation(len(dataset))
+            idxs.append((indices[:split_idx].tolist(), indices[split_idx:].tolist()))
+
+    return idxs
+
+
+def split_img(imgs, ROIs=(3, 3), step=(1, 1)):
+    """Split the imgs in regions of size ROIs.
+
+       Args:
+            imgs (ndarray): images which you want to split 
+            ROIs (tuple): size of sub-regions splitted (ROIs=region of interests)
+            step (tuple): step path from one sub-region to the next one (in the x,y axis)
+
+        Returns:
+            ndarray: splitted subimages. 
+                     The size is (x_num_subROIs*y_num_subROIs, **) where:
+                     x_num_subROIs = ( imgs.shape[1]-int(ROIs[1]/2)*2 )/step[1]
+                     y_num_subROIs = ( imgs.shape[0]-int(ROIs[0]/2)*2 )/step[0]                     
+
+       Example:
+            >>> from dataset_generator import split
+            >>> imgs_splitted = split(imgs, ROI_size = (5,5), step=(2,3))
+    """
+
+    if len(ROIs) > 2:
+        return print("ROIs is a 2 element list")
+
+    if len(step) > 2:
+        return print("step is a 2 element list")
+
+    if type(imgs) != type(np.array(1)):
+        return print("imgs should be a ndarray")
+
+    if len(imgs.shape) == 2:  # Single image with one channel (HxW)
+        splitted = view_as_windows(imgs, (ROIs[0], ROIs[1]), (step[0], step[1]))
+        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

utils/hendrycks_robustness.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/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)