Initial commit

app.py

@@ -0,0 +1,75 @@
+from utils.configuration import Configuration
+import tensorflow as tf
+from utils.model import ModelLoss
+from utils.model import LFUNet
+from utils.architectures import UNet
+
+import gradio as gr
+
+configuration = Configuration()
+filters = (64, 128, 128, 256, 256, 512)
+kernels = (7, 7, 7, 3, 3, 3)
+input_image_size = (256, 256, 3)
+architecture = UNet.RESIDUAL_ATTENTION_UNET_SEPARABLE_CONV
+
+trained_model = LFUNet.build_model(architecture=architecture, input_size=input_image_size, filters=filters,
+                                   kernels=kernels, configuration=configuration)
+trained_model.compile(
+            loss=ModelLoss.ms_ssim_l1_perceptual_loss,
+            optimizer=tf.keras.optimizers.Adam(1e-4),
+            metrics=["acc", tf.keras.metrics.Recall(), tf.keras.metrics.Precision()]
+        )
+
+weights_path = "model_weights/model_epochs-40_batch-20_loss-ms_ssim_l1_perceptual_loss_20230210_15_45_38.ckpt"
+trained_model.load_weights(weights_path)
+
+def main(input_img):
+    try:
+        print(input_img)
+        predicted_image = trained_model.predict(input_img)
+        return predicted_image
+    except Exception as e:
+        raise gr.Error("Sorry, something went wrong. Please try again!")
+
+demo = gr.Interface(
+    title= "Lightweight network for face unmasking",
+    description= "This is a demo of a <b>Lightweight network for face unmasking</b> \
+        designed to provide a powerful and efficient solution for restoring facial details obscured by masks.<br> \
+        To use it, simply upload your image, or click one of the examples to load them. Inference needs some time since this demo uses CPU.",
+    fn = main, 
+    inputs= gr.Image(type="filepath").style(height=256),
+    outputs=gr.Image(type='numpy',shape=(256, 256, 3)).style(height=256),
+    # allow_flagging='never',
+    examples=[
+        ["examples/1.png"],
+        ["examples/2.png"],
+        ["examples/3.png"],
+        ["examples/4.png"],
+        ["examples/5.png"],
+        ["examples/6.png"],
+        ["examples/7.png"],
+        ["examples/8.png"],
+    ],
+    css = """
+        .svelte-mppz8v {
+            text-align: -webkit-center;
+        }
+
+        .gallery {
+            display: flex;
+            flex-wrap: wrap;
+            width: 100%;
+        }
+
+        p {
+            font-size: medium;
+        }
+
+        h1 {
+            font-size: xx-large;
+        }
+        """,
+    theme= 'EveryPizza/Cartoony-Gradio-Theme',
+    # article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2204.04676' target='_blank'>Simple Baselines for Image Restoration</a> | <a href='https://arxiv.org/abs/2204.08714' target='_blank'>NAFSSR: Stereo Image Super-Resolution Using NAFNet</a>  | <a href='https://github.com/megvii-research/NAFNet' target='_blank'> Github Repo</a></p>"
+)
+demo.launch(show_error=True, share= True)

configuration.json

@@ -0,0 +1,25 @@
+{
+  "input_images_path": "data/lfw-deepfunneled",
+  "dataset_archive_download_url": "http://vis-www.cs.umass.edu/lfw/lfw-deepfunneled.tgz",
+  "path_to_patterns": "data/mask_patterns",
+  "train_data_path": "data/train",
+  "test_data_path": "data/test",
+  "landmarks_predictor_path": "shape_predictor_68_face_landmarks.dat",
+  "landmarks_predictor_download_url": "http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2",
+  "minimal_confidence": 0.8,
+  "hyp_ratio": 0.3333,
+  "coordinates_range": [-10, 10],
+  "test_image_count": 100,
+  "train_image_count": 15000,
+  "image_size": [256, 256],
+  "mask_type" : "random",
+  "mask_color" : null,
+  "mask_patter" : null,
+  "mask_pattern_weight" : 0.9,
+  "mask_color_weight" : 0.8,
+  "mask_filter_output" : false,
+  "mask_filter_radius" : 2,
+  "test_results_dir": "data/results/",
+  "train_data_limit": 10000,
+  "test_data_limit": 1000
+}

environment.yaml

@@ -0,0 +1,216 @@
+name: unmask3
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - atk-1.0=2.36.0=ha1a6a79_0
+  - backcall=0.2.0=pyhd3eb1b0_0
+  - ca-certificates=2022.10.11=h06a4308_0
+  - cairo=1.16.0=h19f5f5c_2
+  - certifi=2022.9.24=py37h06a4308_0
+  - cudatoolkit=11.2.2=hbe64b41_10
+  - cudnn=8.1.0.77=h90431f1_0
+  - debugpy=1.5.1=py37h295c915_0
+  - decorator=5.1.1=pyhd3eb1b0_0
+  - entrypoints=0.4=py37h06a4308_0
+  - expat=2.4.9=h6a678d5_0
+  - font-ttf-dejavu-sans-mono=2.37=hd3eb1b0_0
+  - font-ttf-inconsolata=2.001=hcb22688_0
+  - font-ttf-source-code-pro=2.030=hd3eb1b0_0
+  - font-ttf-ubuntu=0.83=h8b1ccd4_0
+  - fontconfig=2.13.1=h6c09931_0
+  - fonts-anaconda=1=h8fa9717_0
+  - fonts-conda-ecosystem=1=hd3eb1b0_0
+  - freetype=2.11.0=h70c0345_0
+  - fribidi=1.0.10=h7b6447c_0
+  - gdk-pixbuf=2.42.8=h433bba3_0
+  - glib=2.69.1=h4ff587b_1
+  - gobject-introspection=1.72.0=py37hbb6d50b_0
+  - graphite2=1.3.14=h295c915_1
+  - graphviz=2.50.0=h3cd0ef9_0
+  - gtk2=2.24.33=h73c1081_2
+  - gts=0.7.6=hb67d8dd_3
+  - harfbuzz=4.3.0=hd55b92a_0
+  - icu=58.2=he6710b0_3
+  - ipython=7.31.1=py37h06a4308_1
+  - jedi=0.18.1=py37h06a4308_1
+  - jpeg=9e=h7f8727e_0
+  - jupyter_client=7.3.4=py37h06a4308_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - lerc=3.0=h295c915_0
+  - libdeflate=1.8=h7f8727e_5
+  - libffi=3.3=he6710b0_2
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgd=2.3.3=h695aa2c_1
+  - libgomp=11.2.0=h1234567_1
+  - libpng=1.6.37=hbc83047_0
+  - librsvg=2.54.4=h19fe530_0
+  - libsodium=1.0.18=h7b6447c_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libtiff=4.4.0=hecacb30_0
+  - libtool=2.4.6=h295c915_1008
+  - libuuid=1.0.3=h7f8727e_2
+  - libwebp-base=1.2.4=h5eee18b_0
+  - libxcb=1.15=h7f8727e_0
+  - libxml2=2.9.14=h74e7548_0
+  - lz4-c=1.9.3=h295c915_1
+  - matplotlib-inline=0.1.6=py37h06a4308_0
+  - ncurses=6.3=h5eee18b_3
+  - nest-asyncio=1.5.5=py37h06a4308_0
+  - ninja=1.10.2=h06a4308_5
+  - ninja-base=1.10.2=hd09550d_5
+  - openssl=1.1.1q=h7f8727e_0
+  - pango=1.50.7=h05da053_0
+  - parso=0.8.3=pyhd3eb1b0_0
+  - pcre=8.45=h295c915_0
+  - pexpect=4.8.0=pyhd3eb1b0_3
+  - pickleshare=0.7.5=pyhd3eb1b0_1003
+  - pixman=0.40.0=h7f8727e_1
+  - ptyprocess=0.7.0=pyhd3eb1b0_2
+  - pygments=2.11.2=pyhd3eb1b0_0
+  - python=3.7.13=h12debd9_0
+  - python-dateutil=2.8.2=pyhd3eb1b0_0
+  - pyzmq=23.2.0=py37h6a678d5_0
+  - readline=8.1.2=h7f8727e_1
+  - setuptools=63.4.1=py37h06a4308_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - sqlite=3.39.2=h5082296_0
+  - tk=8.6.12=h1ccaba5_0
+  - tornado=6.2=py37h5eee18b_0
+  - wcwidth=0.2.5=pyhd3eb1b0_0
+  - wheel=0.37.1=pyhd3eb1b0_0
+  - xz=5.2.5=h7f8727e_1
+  - zeromq=4.3.4=h2531618_0
+  - zlib=1.2.12=h7f8727e_2
+  - zstd=1.5.2=ha4553b6_0
+  - pip:
+    - absl-py==1.2.0
+    - anyio==3.6.2
+    - argon2-cffi==21.3.0
+    - argon2-cffi-bindings==21.2.0
+    - astunparse==1.6.3
+    - attrs==22.2.0
+    - beautifulsoup4==4.11.1
+    - black==22.10.0
+    - bleach==6.0.0
+    - cachetools==5.2.0
+    - cffi==1.15.1
+    - charset-normalizer==2.1.1
+    - click==8.1.3
+    - cmake==3.24.1
+    - cycler==0.11.0
+    - defusedxml==0.7.1
+    - dlib==19.24.1
+    - dotmap==1.3.30
+    - fastjsonschema==2.16.2
+    - flatbuffers==1.12
+    - fonttools==4.37.1
+    - gast==0.4.0
+    - google-auth==2.11.0
+    - google-auth-oauthlib==0.4.6
+    - google-pasta==0.2.0
+    - grpcio==1.48.1
+    - h5py==3.7.0
+    - idna==3.3
+    - imageio==2.27.0
+    - importlib-metadata==4.12.0
+    - importlib-resources==5.10.2
+    - imutils==0.5.4
+    - ipykernel==6.16.2
+    - ipython-genutils==0.2.0
+    - ipywidgets==8.0.2
+    - jinja2==3.1.2
+    - joblib==1.1.0
+    - jsonschema==4.17.3
+    - jupyter-console==6.6.3
+    - jupyter-core==4.12.0
+    - jupyter-server==1.23.5
+    - jupyterlab-pygments==0.2.2
+    - jupyterlab-widgets==3.0.3
+    - jupyterthemes==0.20.0
+    - keras==2.9.0
+    - keras-applications==1.0.8
+    - keras-preprocessing==1.1.2
+    - keras-vggface==0.6
+    - kiwisolver==1.4.4
+    - lesscpy==0.15.1
+    - libclang==14.0.6
+    - markdown==3.4.1
+    - markupsafe==2.1.1
+    - matplotlib==3.5.3
+    - mistune==2.0.4
+    - mtcnn==0.1.1
+    - mypy-extensions==0.4.3
+    - nbclassic==0.4.8
+    - nbclient==0.7.2
+    - nbconvert==7.2.9
+    - nbformat==5.7.3
+    - networkx==2.6.3
+    - notebook==6.5.2
+    - notebook-shim==0.2.2
+    - numpy==1.21.6
+    - oauthlib==3.2.0
+    - opencv-contrib-python==4.6.0.66
+    - opencv-python==4.6.0.66
+    - opt-einsum==3.3.0
+    - packaging==21.3
+    - pandas==1.3.5
+    - pandocfilters==1.5.0
+    - pathspec==0.10.1
+    - pillow==9.2.0
+    - pip==22.3.1
+    - pkgutil-resolve-name==1.3.10
+    - platformdirs==2.5.2
+    - ply==3.11
+    - prometheus-client==0.16.0
+    - prompt-toolkit==3.0.38
+    - protobuf==3.19.4
+    - psutil==5.9.5
+    - pyasn1==0.4.8
+    - pyasn1-modules==0.2.8
+    - pycparser==2.21
+    - pydot==1.4.2
+    - pyparsing==3.0.9
+    - pyrsistent==0.19.3
+    - pytz==2022.2.1
+    - pywavelets==1.3.0
+    - pyyaml==6.0
+    - requests==2.28.1
+    - requests-oauthlib==1.3.1
+    - rsa==4.9
+    - scikit-image==0.19.3
+    - scikit-learn==1.0.2
+    - scipy==1.7.3
+    - seaborn==0.11.2
+    - send2trash==1.8.0
+    - sniffio==1.3.0
+    - soupsieve==2.3.2.post1
+    - tensorboard==2.9.1
+    - tensorboard-data-server==0.6.1
+    - tensorboard-plugin-wit==1.8.1
+    - tensorflow==2.9.2
+    - tensorflow-addons==0.19.0
+    - tensorflow-estimator==2.9.0
+    - tensorflow-io-gcs-filesystem==0.26.0
+    - termcolor==1.1.0
+    - terminado==0.17.1
+    - threadpoolctl==3.1.0
+    - tifffile==2021.11.2
+    - tinycss2==1.2.1
+    - tomli==2.0.1
+    - tqdm==4.64.1
+    - traitlets==5.8.1
+    - trianglesolver==1.2
+    - typed-ast==1.5.4
+    - typeguard==2.13.3
+    - typing-extensions==4.3.0
+    - urllib3==1.26.12
+    - webencodings==0.5.1
+    - websocket-client==1.4.2
+    - werkzeug==2.2.2
+    - widgetsnbextension==4.0.3
+    - wrapt==1.14.1
+    - zipp==3.8.1
+prefix: /home/suresh/miniconda3/envs/unmask

model_weights/checkpoint

@@ -0,0 +1,2 @@
+model_checkpoint_path: "model_epochs-40_batch-20_loss-ms_ssim_l1_perceptual_loss_20230210_15_45_38.ckpt"
+all_model_checkpoint_paths: "model_epochs-40_batch-20_loss-ms_ssim_l1_perceptual_loss_20230210_15_45_38.ckpt"

utils/__init__.py

@@ -0,0 +1,34 @@
+import numpy as np
+from PIL import Image
+import requests
+import functools
+from tqdm.notebook import tqdm
+import shutil
+
+def image_to_array(image: Image) -> np.ndarray:
+    """Convert Image to array"""
+    return np.asarray(image).astype('uint8')
+
+
+def load_image(img_path: str) -> Image:
+    """Load image to array"""
+    return Image.open(img_path)
+
+
+def download_data(url, save_path, file_size=None):
+    """Downloads data from `url` to `save_path`"""
+    r = requests.get(url, stream=True, allow_redirects=True)
+    if r.status_code != 200:
+        r.raise_for_status()
+        raise RuntimeError(f'Request to {url} returned status code {r.status_code}')
+
+    if file_size is None:
+        file_size = int(r.headers.get('content-length', 0))
+
+    r.raw.read = functools.partial(r.raw.read, decode_content=True)  # Decompress if needed
+    with tqdm.wrapattr(r.raw, 'read', total=file_size, desc='') as r_raw:
+        with open(save_path, 'wb') as f:
+            shutil.copyfileobj(r_raw, f)
+
+def plot_image_triple():
+    pass

utils/architectures.py

@@ -0,0 +1,344 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Tuple, Optional
+
+import tensorflow as tf
+from tensorflow.keras.layers import *
+from tensorflow.keras.models import *
+
+class BaseUNet(ABC):
+    """
+    Base Interface for UNet
+    """
+
+    def __init__(self, model: Model):
+        self.model: Model = model
+
+    def get_model(self):
+        return self.model
+
+    @staticmethod
+    @abstractmethod
+    def build_model(input_size: Tuple[int, int, int], filters: Tuple, kernels: Tuple):
+        pass
+
+
+class UNet(Enum):
+    """
+    Enum class defining different architecture types available
+    """
+    DEFAULT = 0
+    DEFAULT_IMAGENET_EMBEDDING = 1
+    RESNET = 3
+    RESIDUAL_ATTENTION_UNET_SEPARABLE_CONV = 4
+
+    def build_model(self, input_size: Tuple[int, int, int], filters: Optional[Tuple] = None,
+                    kernels: Optional[Tuple] = None) -> BaseUNet:
+
+        # set default filters
+        if filters is None:
+            filters = (16, 32, 64, 128, 256)
+
+        # set default kernels
+        if kernels is None:
+            kernels = list(3 for _ in range(len(filters)))
+
+        # check kernels and filters
+        if len(filters) != len(kernels):
+            raise Exception('Kernels and filter count has to match.')
+
+        if self == UNet.DEFAULT_IMAGENET_EMBEDDING:
+            print('Using default UNet model with imagenet embedding')
+            return UNetDefault.build_model(input_size, filters, kernels, use_embedding=True)
+        elif self == UNet.RESNET:
+            print('Using UNet Resnet model')
+            return UNet_resnet.build_model(input_size, filters, kernels)
+        elif self == UNet.RESIDUAL_ATTENTION_UNET_SEPARABLE_CONV:
+            print('Using UNet Resnet model with attention mechanism and separable convolutions')
+            return UNet_ResNet_Attention_SeparableConv.build_model(input_size, filters, kernels)
+
+        print('Using default UNet model')
+        return UNetDefault.build_model(input_size, filters, kernels, use_embedding=False)
+    
+
+class Attention(Layer):
+    def __init__(self, **kwargs):
+        super(Attention, self).__init__(**kwargs)
+
+    def build(self, input_shape):
+        # Create a trainable weight variable for this layer.
+        self.kernel = self.add_weight(name='kernel', 
+                                    shape=(input_shape[-1], 1),
+                                    initializer='glorot_normal',
+                                    trainable=True)
+        self.bias = self.add_weight(name='bias',
+                                   shape=(1,),
+                                   initializer='zeros',
+                                   trainable=True)
+        super(Attention, self).build(input_shape)  # Be sure to call this at the end
+
+    def call(self, x):
+        attention = tf.nn.softmax(tf.matmul(x, self.kernel) + self.bias, axis=-1)
+        return tf.multiply(x, attention)
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+
+class UNet_ResNet_Attention_SeparableConv(BaseUNet):
+    """
+    UNet architecture with resnet blocks, attention mechanism and separable convolutions
+    """
+    @staticmethod
+    def build_model(input_size: Tuple[int, int, int], filters: Tuple, kernels: Tuple):
+
+        p0 = Input(shape=input_size)
+        conv_outputs = []
+        first_layer = SeparableConv2D(filters[0], kernels[0], padding='same')(p0)
+        int_layer = first_layer
+        for i, f in enumerate(filters):
+            int_layer, skip = UNet_ResNet_Attention_SeparableConv.down_block(int_layer, f, kernels[i])
+            conv_outputs.append(skip)
+
+        int_layer = UNet_ResNet_Attention_SeparableConv.bottleneck(int_layer, filters[-1], kernels[-1])
+
+        conv_outputs = list(reversed(conv_outputs))
+        reversed_filter = list(reversed(filters))
+        reversed_kernels = list(reversed(kernels))
+        for i, f in enumerate(reversed_filter):
+            if i + 1 < len(reversed_filter):
+                num_filters_next = reversed_filter[i + 1]
+                num_kernels_next = reversed_kernels[i + 1]
+            else:
+                num_filters_next = f
+                num_kernels_next = reversed_kernels[i]
+            int_layer = UNet_ResNet_Attention_SeparableConv.up_block(int_layer, conv_outputs[i], f, num_filters_next, num_kernels_next)
+            int_layer = Attention()(int_layer)
+            
+        # concat. with the first layer
+        int_layer = Concatenate()([first_layer, int_layer])
+        int_layer = SeparableConv2D(filters[0], kernels[0], padding="same", activation="relu")(int_layer)
+        outputs = SeparableConv2D(3, (1, 1), padding="same", activation="sigmoid")(int_layer)
+        model = Model(p0, outputs)
+        return UNet_ResNet_Attention_SeparableConv(model)
+
+    @staticmethod
+    def down_block(x, num_filters: int = 64, kernel: int = 3):
+        # down-sample inputs
+        x = SeparableConv2D(num_filters, kernel, padding='same', strides=2, dilation_rate = 2)(x)
+
+        # inner block
+        out = SeparableConv2D(num_filters, kernel, padding='same')(x)
+        # out = BatchNormalization()(out)
+        out = Activation('relu')(out)
+        out = SeparableConv2D(num_filters, kernel, padding='same')(out)
+
+        # merge with the skip connection
+        out = Add()([out, x])
+        # out = BatchNormalization()(out)
+        return Activation('relu')(out), x
+
+    @staticmethod
+    def up_block(x, skip, num_filters: int = 64, num_filters_next: int = 64, kernel: int = 3):
+        # add U-Net skip connection - before up-sampling
+        concat = Concatenate()([x, skip])
+
+        # inner block
+        out = SeparableConv2D(num_filters, kernel, padding='same', dilation_rate = 2)(concat)
+        # out = BatchNormalization()(out)
+        out = Activation('relu')(out)
+        out = SeparableConv2D(num_filters, kernel, padding='same')(out)
+
+        # merge with the skip connection
+        out = Add()([out, x])
+        # out = BatchNormalization()(out)
+        out = Activation('relu')(out)
+
+        # up-sample
+        out = UpSampling2D((2, 2))(out)
+        out = SeparableConv2D(num_filters_next, kernel, padding='same')(out)
+        # out = BatchNormalization()(out)
+        return Activation('relu')(out)
+
+    @staticmethod
+    def bottleneck(x, num_filters: int = 64, kernel: int = 3):
+        # inner block
+        out = SeparableConv2D(num_filters, kernel, padding='same', dilation_rate = 2)(x)
+        # out = BatchNormalization()(out)
+        out = Activation('relu')(out)
+        out = SeparableConv2D(num_filters, kernel, padding='same')(out)
+        out = Add()([out, x])
+        # out = BatchNormalization()(out)
+        return Activation('relu')(out)
+
+    
+
+
+
+
+
+
+# Class for UNet with Resnet blocks
+
+class UNet_resnet(BaseUNet):
+    """
+    UNet architecture with resnet blocks
+    """
+
+    @staticmethod
+    def build_model(input_size: Tuple[int, int, int], filters: Tuple, kernels: Tuple):
+
+        p0 = Input(shape=input_size)
+        conv_outputs = []
+        first_layer = Conv2D(filters[0], kernels[0], padding='same')(p0)
+        int_layer = first_layer
+        for i, f in enumerate(filters):
+            int_layer, skip = UNet_resnet.down_block(int_layer, f, kernels[i])
+            conv_outputs.append(skip)
+
+        int_layer = UNet_resnet.bottleneck(int_layer, filters[-1], kernels[-1])
+
+        conv_outputs = list(reversed(conv_outputs))
+        reversed_filter = list(reversed(filters))
+        reversed_kernels = list(reversed(kernels))
+        for i, f in enumerate(reversed_filter):
+            if i + 1 < len(reversed_filter):
+                num_filters_next = reversed_filter[i + 1]
+                num_kernels_next = reversed_kernels[i + 1]
+            else:
+                num_filters_next = f
+                num_kernels_next = reversed_kernels[i]
+            int_layer = UNet_resnet.up_block(int_layer, conv_outputs[i], f, num_filters_next, num_kernels_next)
+
+        # concat. with the first layer
+        int_layer = Concatenate()([first_layer, int_layer])
+        int_layer = Conv2D(filters[0], kernels[0], padding="same", activation="relu")(int_layer)
+        outputs = Conv2D(3, (1, 1), padding="same", activation="sigmoid")(int_layer)
+        model = Model(p0, outputs)
+        return UNet_resnet(model)
+
+    @staticmethod
+    def down_block(x, num_filters: int = 64, kernel: int = 3):
+        # down-sample inputs
+        x = Conv2D(num_filters, kernel, padding='same', strides=2)(x)
+
+        # inner block
+        out = Conv2D(num_filters, kernel, padding='same')(x)
+        # out = BatchNormalization()(out)
+        out = Activation('relu')(out)
+        out = Conv2D(num_filters, kernel, padding='same')(out)
+
+        # merge with the skip connection
+        out = Add()([out, x])
+        # out = BatchNormalization()(out)
+        return Activation('relu')(out), x
+
+    @staticmethod
+    def up_block(x, skip, num_filters: int = 64, num_filters_next: int = 64, kernel: int = 3):
+
+        # add U-Net skip connection - before up-sampling
+        concat = Concatenate()([x, skip])
+
+        # inner block
+        out = Conv2D(num_filters, kernel, padding='same')(concat)
+        # out = BatchNormalization()(out)
+        out = Activation('relu')(out)
+        out = Conv2D(num_filters, kernel, padding='same')(out)
+
+        # merge with the skip connection
+        out = Add()([out, x])
+        # out = BatchNormalization()(out)
+        out = Activation('relu')(out)
+
+        # add U-Net skip connection - before up-sampling
+        concat = Concatenate()([out, skip])
+
+        # up-sample
+        # out = UpSampling2D((2, 2))(concat)
+        out = Conv2DTranspose(num_filters_next, kernel, padding='same', strides=2)(concat)
+        out = Conv2D(num_filters_next, kernel, padding='same')(out)
+        # out = BatchNormalization()(out)
+        return Activation('relu')(out)
+
+    @staticmethod
+    def bottleneck(x, filters, kernel: int = 3):
+        x = Conv2D(filters, kernel, padding='same', name='bottleneck')(x)
+        # x = BatchNormalization()(x)
+        return Activation('relu')(x)
+
+
+class UNetDefault(BaseUNet):
+    """
+     UNet architecture from following github notebook for image segmentation:
+    https://github.com/nikhilroxtomar/UNet-Segmentation-in-Keras-TensorFlow/blob/master/unet-segmentation.ipynb
+    https://github.com/nikhilroxtomar/Polyp-Segmentation-using-UNET-in-TensorFlow-2.0
+    """
+
+    @staticmethod
+    def build_model(input_size: Tuple[int, int, int], filters: Tuple, kernels: Tuple, use_embedding: bool = True):
+
+        p0 = Input(input_size)
+
+        if use_embedding:
+            mobilenet_model = tf.keras.applications.MobileNetV2(
+                input_shape=input_size, include_top=False, weights='imagenet'
+            )
+            mobilenet_model.trainable = False
+            mn1 = mobilenet_model(p0)
+            mn1 = Reshape((16, 16, 320))(mn1)
+
+        conv_outputs = []
+        int_layer = p0
+
+        for f in filters:
+            conv_output, int_layer = UNetDefault.down_block(int_layer, f)
+            conv_outputs.append(conv_output)
+
+        int_layer = UNetDefault.bottleneck(int_layer, filters[-1])
+
+        if use_embedding:
+            int_layer = Concatenate()([int_layer, mn1])
+
+        conv_outputs = list(reversed(conv_outputs))
+        for i, f in enumerate(reversed(filters)):
+            int_layer = UNetDefault.up_block(int_layer, conv_outputs[i], f)
+
+        int_layer = Conv2D(filters[0] // 2, 3, padding="same", activation="relu")(int_layer)
+        outputs = Conv2D(3, (1, 1), padding="same", activation="sigmoid")(int_layer)
+        model = Model(p0, outputs)
+        return UNetDefault(model)
+
+    @staticmethod
+    def down_block(x, filters, kernel_size=(3, 3), padding="same", strides=1):
+        c = Conv2D(filters, kernel_size, padding=padding, strides=strides, activation="relu")(x)
+        # c = BatchNormalization()(c)
+        p = MaxPool2D((2, 2), (2, 2))(c)
+        return c, p
+
+    @staticmethod
+    def up_block(x, skip, filters, kernel_size=(3, 3), padding="same", strides=1):
+        us = UpSampling2D((2, 2))(x)
+        c = Conv2D(filters, kernel_size, padding=padding, strides=strides, activation="relu")(us)
+        # c = BatchNormalization()(c)
+        concat = Concatenate()([c, skip])
+        c = Conv2D(filters, kernel_size, padding=padding, strides=strides, activation="relu")(concat)
+        # c = BatchNormalization()(c)
+        return c
+
+    @staticmethod
+    def bottleneck(x, filters, kernel_size=(3, 3), padding="same", strides=1):
+        c = Conv2D(filters, kernel_size, padding=padding, strides=strides, activation="relu")(x)
+        # c = BatchNormalization()(c)
+        return c
+
+
+if __name__ == "__main__":
+    filters = (64, 128, 128, 256, 256, 512)
+    kernels = (7, 7, 7, 3, 3, 3)
+    input_image_size = (256, 256, 3)
+    # model = UNet_resnet()
+    # model = model.build_model(input_size=input_image_size,filters=filters,kernels=kernels)
+    # print(model.summary())
+    # __init__() missing 1 required positional argument: 'model'
+    model = UNetDefault.build_model(input_size=input_image_size, filters=filters, kernels=kernels)
+    print(model.summary())

utils/configuration.py

@@ -0,0 +1,22 @@
+import os
+import json
+from dataclasses import dataclass
+
+
+@dataclass
+class Configuration:
+    def __init__(self, config_file_path: str = "configuration.json"):
+        self.config_file_path = config_file_path
+        self.config_json = None
+        if os.path.exists(config_file_path):
+            with open(self.config_file_path, 'r') as json_file:
+                self.config_json = json.load(json_file)
+        else:
+            print(f'ERROR: Configuration JSON {config_file_path} does not exist.')
+
+    def get(self, key: str):
+        if key in self.config_json:
+            return self.config_json[key]
+        else:
+            print(f'ERROR: Key \'{key}\' is not in configuration JSON.')
+            return None

utils/data_generator.py

@@ -0,0 +1,151 @@
+import copy
+import dlib
+import os
+import bz2
+import random
+from tqdm.notebook import tqdm
+import shutil
+from utils import image_to_array, load_image, download_data
+from utils.face_detection import crop_face, get_face_keypoints_detecting_function
+from mask_utils.mask_utils import mask_image
+
+
+class DataGenerator:
+    def __init__(self, configuration):
+        self.configuration = configuration
+        self.path_to_data = configuration.get('input_images_path')
+        self.path_to_patterns = configuration.get('path_to_patterns')
+        self.minimal_confidence = configuration.get('minimal_confidence')
+        self.hyp_ratio = configuration.get('hyp_ratio')
+        self.coordinates_range = configuration.get('coordinates_range')
+        self.test_image_count = configuration.get('test_image_count')
+        self.train_image_count = configuration.get('train_image_count')
+        self.train_data_path = configuration.get('train_data_path')
+        self.test_data_path = configuration.get('test_data_path')
+        self.predictor_path = configuration.get('landmarks_predictor_path')
+        self.check_predictor()
+
+        self.valid_image_extensions = ('png', 'jpg', 'jpeg')
+        self.face_keypoints_detecting_fun = get_face_keypoints_detecting_function(self.minimal_confidence)
+
+    def check_predictor(self):
+        """ Check if predictor exists. If not downloads it. """
+        if not os.path.exists(self.predictor_path):
+            print('Downloading missing predictor.')
+            url = self.configuration.get('landmarks_predictor_download_url')
+            download_data(url, self.predictor_path + '.bz2', 64040097)
+            print(f'Decompressing downloaded file into {self.predictor_path}')
+            with bz2.BZ2File(self.predictor_path + '.bz2') as fr, open(self.predictor_path, 'wb') as fw:
+                shutil.copyfileobj(fr, fw)
+
+    def get_face_landmarks(self, image):
+        """Compute 68 facial landmarks"""
+        landmarks = []
+        image_array = image_to_array(image)
+        detector = dlib.get_frontal_face_detector()
+        predictor = dlib.shape_predictor(self.predictor_path)
+        face_rectangles = detector(image_array)
+        if len(face_rectangles) < 1:
+            return None
+        dlib_shape = predictor(image_array, face_rectangles[0])
+        for i in range(0, dlib_shape.num_parts):
+            landmarks.append([dlib_shape.part(i).x, dlib_shape.part(i).y])
+        return landmarks
+
+    def get_files_faces(self):
+        """Get path of all images in dataset"""
+        image_files = []
+        for dirpath, dirs, files in os.walk(self.path_to_data):
+            for filename in files:
+                fname = os.path.join(dirpath, filename)
+                if fname.endswith(self.valid_image_extensions):
+                    image_files.append(fname)
+
+        return image_files
+
+    def generate_images(self, image_size=None, test_image_count=None, train_image_count=None):
+        """Generate test and train data (images with and without the mask)"""
+        if image_size is None:
+            image_size = self.configuration.get('image_size')
+        if test_image_count is None:
+            test_image_count = self.test_image_count
+        if train_image_count is None:
+            train_image_count = self.train_image_count
+
+        if not os.path.exists(self.train_data_path):
+            os.mkdir(self.train_data_path)
+            os.mkdir(os.path.join(self.train_data_path, 'inputs'))
+            os.mkdir(os.path.join(self.train_data_path, 'outputs'))
+
+        if not os.path.exists(self.test_data_path):
+            os.mkdir(self.test_data_path)
+            os.mkdir(os.path.join(self.test_data_path, 'inputs'))
+            os.mkdir(os.path.join(self.test_data_path, 'outputs'))
+
+        print('Generating testing data')
+        self.generate_data(test_image_count,
+                           image_size=image_size,
+                           save_to=self.test_data_path)
+        print('Generating training data')
+        self.generate_data(train_image_count,
+                           image_size=image_size,
+                           save_to=self.train_data_path)
+
+    def generate_data(self, number_of_images, image_size=None, save_to=None):
+        """ Add masks on `number_of_images` images
+            if save_to is valid path to folder images are saved there otherwise generated data are just returned in list
+        """
+        inputs = []
+        outputs = []
+
+        if image_size is None:
+            image_size = self.configuration.get('image_size')
+
+        for i, file in tqdm(enumerate(random.sample(self.get_files_faces(), number_of_images)), total=number_of_images):
+            # Load images
+            image = load_image(file)
+
+            # Detect keypoints and landmarks on face
+            face_landmarks = self.get_face_landmarks(image)
+            if face_landmarks is None:
+                continue
+            keypoints = self.face_keypoints_detecting_fun(image)
+
+            # Generate mask
+            image_with_mask = mask_image(copy.deepcopy(image), face_landmarks, self.configuration)
+
+            # Crop images
+            cropped_image = crop_face(image_with_mask, keypoints)
+            cropped_original = crop_face(image, keypoints)
+
+            # Resize all images to NN input size
+            res_image = cropped_image.resize(image_size)
+            res_original = cropped_original.resize(image_size)
+
+            # Save generated data to lists or to folder
+            if save_to is None:
+                inputs.append(res_image)
+                outputs.append(res_original)
+            else:
+                res_image.save(os.path.join(save_to, 'inputs', f"{i:06d}.png"))
+                res_original.save(os.path.join(save_to, 'outputs', f"{i:06d}.png"))
+
+        if save_to is None:
+            return inputs, outputs
+
+    def get_dataset_examples(self, n=10, test_dataset=False):
+        """
+        Returns `n` random images form dataset. If `test_dataset` parameter
+        is not provided or False it will return images from training part of dataset.
+        If `test_dataset` parameter is True it will return images from testing part of dataset.
+        """
+        if test_dataset:
+            data_path = self.test_data_path
+        else:
+            data_path = self.train_data_path
+
+        images = os.listdir(os.path.join(data_path, 'inputs'))
+        images = random.sample(images, n)
+        inputs = [os.path.join(data_path, 'inputs', img) for img in images]
+        outputs = [os.path.join(data_path, 'outputs', img) for img in images]
+        return inputs, outputs

utils/face_detection.py

@@ -0,0 +1,111 @@
+"""Functions for face detection"""
+from math import pi
+from typing import Tuple, Optional, Dict
+
+import tensorflow as tf
+import matplotlib.patches as patches
+import matplotlib.pyplot as plt
+from PIL import Image
+from mtcnn import MTCNN
+from trianglesolver import solve
+
+from utils import image_to_array
+
+
+def compute_slacks(height, width, hyp_ratio) -> Tuple[float, float]:
+    """Compute slacks to add to bounding box on each site"""
+
+    # compute angle and side for hypotenuse
+    _, b, _, A, _, _ = solve(c=width, a=height, B=pi / 2)
+
+    # compute new height and width
+    a, _, c, _, _, _ = solve(b=b * (1.0 + hyp_ratio), B=pi / 2, A=A)
+
+    # compute slacks
+    return c - width, a - height
+
+
+def get_face_keypoints_detecting_function(minimal_confidence: float = 0.8):
+    """Create function for face keypoints detection"""
+
+    # face detector
+    detector = MTCNN()
+
+    # detect faces and their keypoints
+    def get_keypoints(image: Image) -> Optional[Dict]:
+
+        # run inference to detect faces (on CPU only)
+        with tf.device("/cpu:0"):
+            detection = detector.detect_faces(image_to_array(image))
+
+        # run detection and keep results with certain confidence only
+        results = [item for item in detection if item['confidence'] > minimal_confidence]
+
+        # nothing found
+        if len(results) == 0:
+            return None
+
+        # return result with highest confidence and size
+        return max(results, key=lambda item: item['confidence'] * item['box'][2] * item['box'][3])
+
+    # return function
+    return get_keypoints
+
+
+def plot_face_detection(image: Image, ax, face_keypoints: Optional, hyp_ratio: float = 1 / 3):
+    """Plot faces with keypoints and bounding boxes"""
+
+    # make annotations
+    if face_keypoints is not None:
+
+        # get bounding box
+        x, y, width, height = face_keypoints['box']
+
+        # add rectangle patch for detected face
+        rectangle = patches.Rectangle((x, y), width, height, linewidth=1, edgecolor='r', facecolor='none')
+        ax.add_patch(rectangle)
+
+        # add rectangle patch with slacks
+        w_s, h_s = compute_slacks(height, width, hyp_ratio)
+        rectangle = patches.Rectangle((x - w_s, y - h_s), width + 2 * w_s, height + 2 * h_s, linewidth=1, edgecolor='r',
+                                      facecolor='none')
+        ax.add_patch(rectangle)
+
+        # add keypoints
+        for coordinates in face_keypoints['keypoints'].values():
+            circle = plt.Circle(coordinates, 3, color='r')
+            ax.add_artist(circle)
+
+    # add image
+    ax.imshow(image)
+
+
+def get_crop_points(image: Image, face_keypoints: Optional, hyp_ratio: float = 1 / 3) -> Image:
+    """Find position where to crop face from image"""
+    if face_keypoints is None:
+        return 0, 0, image.width, image.height
+
+    # get bounding box
+    x, y, width, height = face_keypoints['box']
+
+    # compute slacks
+    w_s, h_s = compute_slacks(height, width, hyp_ratio)
+
+    # compute coordinates
+    left = min(max(0, x - w_s), image.width)
+    upper = min(max(0, y - h_s), image.height)
+    right = min(x + width + w_s, image.width)
+    lower = min(y + height + h_s, image.height)
+
+    return left, upper, right, lower
+
+
+def crop_face(image: Image, face_keypoints: Optional, hyp_ratio: float = 1 / 3) -> Image:
+    """Crop input image to just the face"""
+    if face_keypoints is None:
+        print("No keypoints detected on image")
+        return image
+
+    left, upper, right, lower = get_crop_points(image, face_keypoints, hyp_ratio)
+
+    return image.crop((left, upper, right, lower))

utils/model.py

@@ -0,0 +1,495 @@
+import os
+from datetime import datetime
+from glob import glob
+from typing import Tuple, Optional
+from utils import load_image
+import random
+import cv2
+import numpy as np
+import tensorflow as tf
+from PIL import Image
+from sklearn.model_selection import train_test_split
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.utils import CustomObjectScope
+from utils.face_detection import get_face_keypoints_detecting_function, crop_face, get_crop_points
+from utils.architectures import UNet
+from tensorflow.keras.losses import MeanSquaredError, mean_squared_error
+from keras_vggface.vggface import VGGFace
+import tensorflow.keras.backend as K
+from tensorflow.keras.applications import VGG19
+
+# # VGG19 model for perceptual loss
+# vgg = VGG19(include_top=False, weights='imagenet')
+
+# def preprocess_image(image):
+#     image = tf.image.resize(image, (224, 224))
+#     image = tf.keras.applications.vgg19.preprocess_input(image)
+#     return image
+
+# def perceptual_loss(y_true, y_pred):
+#     y_true = preprocess_image(y_true)
+#     y_pred = preprocess_image(y_pred)
+#     y_true_c = vgg(y_true)
+#     y_pred_c = vgg(y_pred)
+#     loss = K.mean(K.square(y_pred_c - y_true_c))
+#     return loss
+
+vgg_face_model = VGGFace(model='resnet50', include_top=False, input_shape=(256, 256, 3), pooling='avg')
+
+
+class ModelLoss:
+    @staticmethod
+    @tf.function
+    def ms_ssim_l1_perceptual_loss(gt, y_pred, max_val=1.0, l1_weight=1.0):
+        """
+        Computes MS-SSIM and perceptual loss
+        @param gt: Ground truth image
+        @param y_pred: Predicted image
+        @param max_val: Maximal MS-SSIM value
+        @param l1_weight: Weight of L1 normalization
+        @return: MS-SSIM and perceptual loss
+        """
+
+        # Compute SSIM loss
+        ssim_loss = 1 - tf.reduce_mean(tf.image.ssim(gt, y_pred, max_val=max_val))
+
+        # Compute perceptual loss
+        vgg_face_outputs = vgg_face_model(y_pred)
+        vgg_face_loss = tf.reduce_mean(tf.losses.mean_squared_error(vgg_face_outputs,vgg_face_model(gt)))
+        
+
+        # Combine both losses with l1 normalization
+        l1 = mean_squared_error(gt, y_pred)
+        l1_casted = tf.cast(l1 * l1_weight, tf.float32)
+        return ssim_loss + l1_casted + vgg_face_loss
+
+
+
+class LFUNet(tf.keras.models.Model):
+    """
+    Model for Mask2Face - removes mask from people faces using U-net neural network
+    """
+    def __init__(self, model: tf.keras.models.Model, configuration=None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.model: tf.keras.models.Model = model
+        self.configuration = configuration
+        self.face_keypoints_detecting_fun = get_face_keypoints_detecting_function(0.8)
+        self.mse = MeanSquaredError()
+
+    def call(self, x, **kwargs):
+        return self.model(x)
+
+    @staticmethod
+    @tf.function
+    def ssim_loss(gt, y_pred, max_val=1.0):
+        """
+        Computes standard SSIM loss
+        @param gt: Ground truth image
+        @param y_pred: Predicted image
+        @param max_val: Maximal SSIM value
+        @return: SSIM loss
+        """
+        return 1 - tf.reduce_mean(tf.image.ssim(gt, y_pred, max_val=max_val))
+
+    @staticmethod
+    @tf.function
+    def ssim_l1_loss(gt, y_pred, max_val=1.0, l1_weight=1.0):
+        """
+        Computes SSIM loss with L1 normalization
+        @param gt: Ground truth image
+        @param y_pred: Predicted image
+        @param max_val: Maximal SSIM value
+        @param l1_weight: Weight of L1 normalization
+        @return: SSIM L1 loss
+        """
+        ssim_loss = 1 - tf.reduce_mean(tf.image.ssim(gt, y_pred, max_val=max_val))
+        l1 = mean_squared_error(gt, y_pred)
+        return ssim_loss + tf.cast(l1 * l1_weight, tf.float32)
+
+    
+
+    # @staticmethod
+    # @tf.function
+    # def ms_ssim_l1_perceptual_loss(gt, y_pred, max_val=1.0, l1_weight=1.0, perceptual_weight=1.0):
+    #     """
+    #     Computes MS-SSIM loss, L1 loss, and perceptual loss
+    #     @param gt: Ground truth image
+    #     @param y_pred: Predicted image
+    #     @param max_val: Maximal SSIM value
+    #     @param l1_weight: Weight of L1 normalization
+    #     @param perceptual_weight: Weight of perceptual loss
+    #     @return: MS-SSIM L1 perceptual loss
+    #     """
+    #     y_pred = tf.clip_by_value(y_pred, 0, float("inf"))
+    #     y_pred = tf.debugging.check_numerics(y_pred, message='y_pred has NaN values')
+    #     ms_ssim_loss = 1 - tf.reduce_mean(tf.image.ssim_multiscale(gt, y_pred, max_val=max_val))
+    #     l1_loss = tf.losses.mean_absolute_error(gt, y_pred)
+    #     vgg_face_outputs = vgg_face_model(y_pred)
+    #     vgg_face_loss = tf.reduce_mean(tf.losses.mean_squared_error(vgg_face_outputs,vgg_face_model(gt)))
+    #     return ms_ssim_loss + tf.cast(l1_loss * l1_weight, tf.float32) + perceptual_weight*vgg_face_loss
+
+    # Function for ms-ssim loss + l1 loss
+    @staticmethod
+    @tf.function
+    def ms_ssim_l1_loss(gt, y_pred, max_val=1.0, l1_weight=1.0):
+        """
+        Computes MS-SSIM loss and L1 loss
+        @param gt: Ground truth image
+        @param y_pred: Predicted image
+        @param max_val: Maximal SSIM value
+        @param l1_weight: Weight of L1 normalization
+        @return: MS-SSIM L1 loss
+        """
+        # Replace NaN values with 0
+        y_pred = tf.clip_by_value(y_pred, 0, float("inf"))
+
+        ms_ssim_loss = 1 - tf.reduce_mean(tf.image.ssim_multiscale(gt, y_pred, max_val=max_val))
+        l1_loss = tf.losses.mean_absolute_error(gt, y_pred)
+        return ms_ssim_loss + tf.cast(l1_loss * l1_weight, tf.float32)
+
+
+    @staticmethod
+    def load_model(model_path, configuration=None):
+        """
+        Loads saved h5 file with trained model.
+        @param configuration: Optional instance of Configuration with config JSON
+        @param model_path: Path to h5 file
+        @return: LFUNet
+        """
+        with CustomObjectScope({'ssim_loss': LFUNet.ssim_loss, 'ssim_l1_loss': LFUNet.ssim_l1_loss, 'ms_ssim_l1_perceptual_loss': ModelLoss.ms_ssim_l1_perceptual_loss, 'ms_ssim_l1_loss': LFUNet.ms_ssim_l1_loss}):
+            model = tf.keras.models.load_model(model_path)
+        return LFUNet(model, configuration)
+
+    @staticmethod
+    def build_model(architecture: UNet, input_size: Tuple[int, int, int], filters: Optional[Tuple] = None,
+                    kernels: Optional[Tuple] = None, configuration=None):
+        """
+        Builds model based on input arguments
+        @param architecture: utils.architectures.UNet architecture
+        @param input_size: Size of input images
+        @param filters: Tuple with sizes of filters in U-net
+        @param kernels: Tuple with sizes of kernels in U-net. Must be the same size as filters.
+        @param configuration: Optional instance of Configuration with config JSON
+        @return: LFUNet
+        """
+        return LFUNet(architecture.build_model(input_size, filters, kernels).get_model(), configuration)
+
+    def train(self, epochs=20, batch_size=20, loss_function='mse', learning_rate=1e-4,
+              predict_difference: bool = False):
+        """
+        Train the model.
+        @param epochs: Number of epochs during training
+        @param batch_size: Batch size
+        @param loss_function: Loss function. Either standard tensorflow loss function or `ssim_loss` or `ssim_l1_loss`
+        @param learning_rate: Learning rate
+        @param predict_difference: Compute prediction on difference between input and output image
+        @return: History of training
+        """
+        # get data
+        (train_x, train_y), (valid_x, valid_y) = self.load_train_data()
+        (test_x, test_y) = self.load_test_data()
+
+        train_dataset = LFUNet.tf_dataset(train_x, train_y, batch_size, predict_difference)
+        valid_dataset = LFUNet.tf_dataset(valid_x, valid_y, batch_size, predict_difference, train=False)
+        test_dataset = LFUNet.tf_dataset(test_x, test_y, batch_size, predict_difference, train=False)
+
+        # select loss
+        if loss_function == 'ssim_loss':
+            loss = LFUNet.ssim_loss
+        elif loss_function == 'ssim_l1_loss':
+            loss = LFUNet.ssim_l1_loss
+        elif loss_function == 'ms_ssim_l1_perceptual_loss':
+            loss = ModelLoss.ms_ssim_l1_perceptual_loss
+        elif loss_function == 'ms_ssim_l1_loss':
+            loss = LFUNet.ms_ssim_l1_loss
+        else:
+            loss = loss_function
+
+        # compile loss with selected loss function
+        self.model.compile(
+            loss=loss,
+            optimizer=tf.keras.optimizers.Adam(learning_rate),
+            metrics=["acc", tf.keras.metrics.Recall(), tf.keras.metrics.Precision()]
+        )
+
+        # define callbacks
+        callbacks = [
+            ModelCheckpoint(
+                f'models/model_epochs-{epochs}_batch-{batch_size}_loss-{loss_function}_{LFUNet.get_datetime_string()}.h5'),
+            EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True)
+        ]
+
+        # evaluation before training
+        results = self.model.evaluate(test_dataset)
+        print("- TEST -> LOSS: {:10.4f}, ACC: {:10.4f}, RECALL: {:10.4f}, PRECISION: {:10.4f}".format(*results))
+
+        # fit the model
+        history = self.model.fit(train_dataset, validation_data=valid_dataset, epochs=epochs, callbacks=callbacks)
+
+        # evaluation after training
+        results = self.model.evaluate(test_dataset)
+        print("- TEST -> LOSS: {:10.4f}, ACC: {:10.4f}, RECALL: {:10.4f}, PRECISION: {:10.4f}".format(*results))
+
+        # use the model for inference on several test images
+        self._test_results(test_x, test_y, predict_difference)
+
+        # return training history
+        return history
+
+    def _test_results(self, test_x, test_y, predict_difference: bool):
+        """
+        Test trained model on testing dataset. All images in testing dataset are processed and result image triples
+        (input with mask, ground truth, model output) are stored to `data/results` into folder with time stamp
+        when this method was executed.
+        @param test_x: List of input images
+        @param test_y: List of ground truth output images
+        @param predict_difference: Compute prediction on difference between input and output image
+        @return: None
+        """
+        if self.configuration is None:
+            result_dir = f'data/results/{LFUNet.get_datetime_string()}/'
+        else:
+            result_dir = os.path.join(self.configuration.get('test_results_dir'), LFUNet.get_datetime_string())
+        os.makedirs(result_dir, exist_ok=True)
+
+        for i, (x, y) in enumerate(zip(test_x, test_y)):
+            x = LFUNet.read_image(x)
+            y = LFUNet.read_image(y)
+
+            y_pred = self.model.predict(np.expand_dims(x, axis=0))
+            if predict_difference:
+                y_pred = (y_pred * 2) - 1
+                y_pred = np.clip(x - y_pred.squeeze(axis=0), 0.0, 1.0)
+            else:
+                y_pred = y_pred.squeeze(axis=0)
+            h, w, _ = x.shape
+            white_line = np.ones((h, 10, 3)) * 255.0
+
+            all_images = [
+                x * 255.0, white_line,
+                y * 255.0, white_line,
+                y_pred * 255.0
+            ]
+            image = np.concatenate(all_images, axis=1)
+            cv2.imwrite(os.path.join(result_dir, f"{i}.png"), image)
+
+    def summary(self):
+        """
+        Prints model summary
+        """
+        self.model.summary()
+
+    def predict(self, img_path, predict_difference: bool = False):
+        """
+        Use trained model to take down the mask from image with person wearing the mask.
+        @param img_path: Path to image to processed
+        @param predict_difference: Compute prediction on difference between input and output image
+        @return: Image without the mask on the face
+        """
+        # Load image into RGB format
+        image = load_image(img_path)
+        image = image.convert('RGB')
+
+        # Find facial keypoints and crop the image to just the face
+        keypoints = self.face_keypoints_detecting_fun(image)
+        cropped_image = crop_face(image, keypoints)
+        print(cropped_image.size)
+
+        # Resize image to input recognized by neural net
+        resized_image = cropped_image.resize((256, 256))
+        image_array = np.array(resized_image)
+
+        # Convert from RGB to BGR (open cv format)
+        image_array = image_array[:, :, ::-1].copy()
+        image_array = image_array / 255.0
+
+        # Remove mask from input image
+        y_pred = self.model.predict(np.expand_dims(image_array, axis=0))
+        h, w, _ = image_array.shape
+
+        if predict_difference:
+            y_pred = (y_pred * 2) - 1
+            y_pred = np.clip(image_array - y_pred.squeeze(axis=0), 0.0, 1.0)
+        else:
+            y_pred = y_pred.squeeze(axis=0)
+
+        # Convert output from model to image and scale it back to original size
+        y_pred = y_pred * 255.0
+        im = Image.fromarray(y_pred.astype(np.uint8)[:, :, ::-1])
+        im = im.resize(cropped_image.size)
+        left, upper, _, _ = get_crop_points(image, keypoints)
+
+        # Combine original image with output from model
+        image.paste(im, (int(left), int(upper)))
+        return image
+
+    @staticmethod
+    def get_datetime_string():
+        """
+        Creates date-time string
+        @return: String with current date and time
+        """
+        now = datetime.now()
+        return now.strftime("%Y%m%d_%H_%M_%S")
+
+    def load_train_data(self, split=0.2):
+        """
+        Loads training data (paths to training images)
+        @param split: Percentage of training data used for validation as float from 0.0 to 1.0. Default 0.2.
+        @return: Two tuples - first with training data (tuple with (input images, output images)) and second
+                    with validation data (tuple with (input images, output images))
+        """
+        if self.configuration is None:
+            train_dir = 'data/train/'
+            limit = None
+        else:
+            train_dir = self.configuration.get('train_data_path')
+            limit = self.configuration.get('train_data_limit')
+        print(f'Loading training data from {train_dir} with limit of {limit} images')
+        return LFUNet.load_data(os.path.join(train_dir, 'inputs'), os.path.join(train_dir, 'outputs'), split, limit)
+
+    def load_test_data(self):
+        """
+        Loads testing data (paths to testing images)
+        @return: Tuple with testing data - (input images, output images)
+        """
+        if self.configuration is None:
+            test_dir = 'data/test/'
+            limit = None
+        else:
+            test_dir = self.configuration.get('test_data_path')
+            limit = self.configuration.get('test_data_limit')
+        print(f'Loading testing data from {test_dir} with limit of {limit} images')
+        return LFUNet.load_data(os.path.join(test_dir, 'inputs'), os.path.join(test_dir, 'outputs'), None, limit)
+
+    @staticmethod
+    def load_data(input_path, output_path, split=0.2, limit=None):
+        """
+        Loads data (paths to images)
+        @param input_path: Path to folder with input images
+        @param output_path: Path to folder with output images
+        @param split: Percentage of data used for validation as float from 0.0 to 1.0. Default 0.2.
+                      If split is None it expects you are loading testing data, otherwise expects training data.
+        @param limit: Maximal number of images loaded from data folder. Default None (no limit).
+        @return: If split is not None: Two tuples - first with training data (tuple with (input images, output images))
+                    and second with validation data (tuple with (input images, output images))
+                 Else: Tuple with testing data - (input images, output images)
+        """
+        images = sorted(glob(os.path.join(input_path, "*.png")))
+        masks = sorted(glob(os.path.join(output_path, "*.png")))
+        if len(images) == 0:
+            raise TypeError(f'No images found in {input_path}')
+        if len(masks) == 0:
+            raise TypeError(f'No images found in {output_path}')
+
+        if limit is not None:
+            images = images[:limit]
+            masks = masks[:limit]
+
+        if split is not None:
+            total_size = len(images)
+            valid_size = int(split * total_size)
+            train_x, valid_x = train_test_split(images, test_size=valid_size, random_state=42)
+            train_y, valid_y = train_test_split(masks, test_size=valid_size, random_state=42)
+            return (train_x, train_y), (valid_x, valid_y)
+
+        else:
+            return images, masks
+
+    @staticmethod
+    def read_image(path):
+        """
+        Loads image, resize it to size 256x256 and normalize to float values from 0.0 to 1.0.
+        @param path: Path to image to be loaded.
+        @return: Loaded image in open CV format.
+        """
+        x = cv2.imread(path, cv2.IMREAD_COLOR)
+        x = cv2.resize(x, (256, 256))
+        x = x / 255.0
+        return x
+
+    @staticmethod
+    def tf_parse(x, y):
+        """
+        Mapping function for dataset creation. Load and resize images.
+        @param x: Path to input image
+        @param y: Path to output image
+        @return: Tuple with input and output image with shape (256, 256, 3)
+        """
+        def _parse(x, y):
+            x = LFUNet.read_image(x.decode())
+            y = LFUNet.read_image(y.decode())
+            return x, y
+
+        x, y = tf.numpy_function(_parse, [x, y], [tf.float64, tf.float64])
+        x.set_shape([256, 256, 3])
+        y.set_shape([256, 256, 3])
+        return x, y
+
+    @staticmethod
+    def tf_dataset(x, y, batch=8, predict_difference: bool = False, train: bool = True):
+        """
+        Creates standard tensorflow dataset.
+        @param x: List of paths to input images
+        @param y: List of paths to output images
+        @param batch: Batch size
+        @param predict_difference: Compute prediction on difference between input and output image
+        @param train: Flag if training dataset should be generated
+        @return: Dataset with loaded images
+        """
+        dataset = tf.data.Dataset.from_tensor_slices((x, y))
+        dataset = dataset.map(LFUNet.tf_parse)
+        random_seed = random.randint(0, 999999999)
+
+        if predict_difference:
+            def map_output(img_in, img_target):
+                return img_in, (img_in - img_target + 1.0) / 2.0
+
+            dataset = dataset.map(map_output)
+
+        if train:
+            # for the train set, we want to apply data augmentations and shuffle data to different batches
+
+            # random flip
+            def flip(img_in, img_out):
+                return tf.image.random_flip_left_right(img_in, random_seed), \
+                       tf.image.random_flip_left_right(img_out, random_seed)
+
+            # augmenting quality - parameters
+            hue_delta = 0.05
+            saturation_low = 0.2
+            saturation_up = 1.3
+            brightness_delta = 0.1
+            contrast_low = 0.2
+            contrast_up = 1.5
+
+            # augmenting quality
+            def color(img_in, img_out):
+                # Augmentations applied are:
+                # - random hue
+                # - random saturation
+                # - random brightness
+                # - random contrast
+                # - random flip left right
+                # - random flip up down
+                img_in = tf.image.random_hue(img_in, hue_delta, random_seed) 
+                img_in = tf.image.random_saturation(img_in, saturation_low, saturation_up, random_seed)
+                img_in = tf.image.random_brightness(img_in, brightness_delta, random_seed)
+                img_in = tf.image.random_contrast(img_in, contrast_low, contrast_up, random_seed)
+                img_out = tf.image.random_hue(img_out, hue_delta, random_seed)
+                img_out = tf.image.random_saturation(img_out, saturation_low, saturation_up, random_seed)
+                img_out = tf.image.random_brightness(img_out, brightness_delta, random_seed)
+                img_out = tf.image.random_contrast(img_out, contrast_low, contrast_up, random_seed)
+                return img_in, img_out
+
+            # shuffle data and create batches
+            dataset = dataset.shuffle(5000)
+            dataset = dataset.batch(batch)
+
+            # apply augmentations
+            dataset = dataset.map(flip)
+            dataset = dataset.map(color)
+        else:
+            dataset = dataset.batch(batch)
+
+        return dataset.prefetch(tf.data.experimental.AUTOTUNE)