diff --git a/.DS_Store b/.DS_Store
new file mode 100644
index 0000000..b34f877
Binary files /dev/null and b/.DS_Store differ
diff --git a/create_dataset.py b/create_dataset.py
new file mode 100644
index 0000000..4926e78
--- /dev/null
+++ b/create_dataset.py
@@ -0,0 +1,68 @@
+from glob import glob
+from os import path
+from typing import Optional
+
+import numpy as np
+from PIL import Image
+from torch.utils.data import Dataset
+
+from prep_image import transform_image, to_grayscale, prepare_image
+
+
+class RandomImagePixelationDataset(Dataset):
+
+    def __init__(
+            self,
+            image_dir,
+            width_range: tuple[int, int],
+            height_range: tuple[int, int],
+            size_range: tuple[int, int],
+            dtype: Optional[type] = None
+    ):
+        RandomImagePixelationDataset._check_range(width_range, "width")
+        RandomImagePixelationDataset._check_range(height_range, "height")
+        RandomImagePixelationDataset._check_range(size_range, "size")
+        self.image_files = sorted(path.abspath(f) for f in glob(path.join(image_dir, "**", "*.jpg"), recursive=True))
+        self.width_range = width_range
+        self.height_range = height_range
+        self.size_range = size_range
+        self.dtype = dtype
+
+    @staticmethod
+    def _check_range(r: tuple[int, int], name: str):
+        if r[0] < 2:
+            raise ValueError(f"minimum {name} must be >= 2")
+        if r[0] > r[1]:
+            raise ValueError(f"minimum {name} must be <= maximum {name}")
+
+    def __getitem__(self, index):
+        with Image.open(self.image_files[index]) as img:
+            img = transform_image(img)
+            image = np.array(img, dtype=self.dtype)
+        image = to_grayscale(image)  # Image shape is now (1, H, W)
+        image_width = image.shape[-1]
+        image_height = image.shape[-2]
+
+        # Create RNG in each __getitem__ call to ensure reproducibility even in
+        # environments with multiple threads and/or processes
+        rng = np.random.default_rng(seed=index)
+
+        # Both width and height can be arbitrary, but they must not exceed the
+        # actual image width and height
+        width = min(rng.integers(low=self.width_range[0], high=self.width_range[1], endpoint=True), image_width)
+        height = min(rng.integers(low=self.height_range[0], high=self.height_range[1], endpoint=True), image_height)
+
+        # Ensure that x and y always fit with the randomly chosen width and
+        # height (and not throw an error in "prepare_image")
+        x = rng.integers(image_width - width, endpoint=True)
+        y = rng.integers(image_height - height, endpoint=True)
+
+        # Block size can be arbitrary again
+        size = rng.integers(low=self.size_range[0], high=self.size_range[1], endpoint=True)
+
+        pixelated_image, known_array, target_array = prepare_image(image, x, y, width, height, size)
+
+        return pixelated_image, known_array, target_array, self.image_files[index]
+
+    def __len__(self):
+        return len(self.image_files)
diff --git a/image_augmentation.py b/image_augmentation.py
new file mode 100644
index 0000000..f72ba8a
--- /dev/null
+++ b/image_augmentation.py
@@ -0,0 +1,36 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+import random
+from glob import glob
+from os import path
+
+
+def random_augmented_image(
+        image: Image,
+        seed: int,
+) -> torch.Tensor:
+    candidates = [transforms.RandomRotation(180), transforms.RandomVerticalFlip(),
+                  transforms.RandomHorizontalFlip(), transforms.ColorJitter()]
+    random.seed(seed)
+    x = image
+    random_transforms = random.sample(candidates, k=2)
+    for transform in random_transforms:
+        x = transform(x)
+    x = transforms.ToTensor()(x)
+    x = torch.nn.Dropout(p=0.01)(x)
+    return transforms.ToPILImage()(x)
+
+
+if __name__ == "__main__":
+    image_dir = r"training_raw"
+    augmented_image_dir = r"training_raw/augmented"
+    image_files = sorted(path.abspath(f) for f in glob(path.join(image_dir, "**", "*.jpg"), recursive=True))
+    i = 0
+    for img in image_files:
+        with Image.open(img) as image:
+            augmented_image = random_augmented_image(image, seed=3)
+            filename = path.basename(img)
+            augmented_image_path = path.join(augmented_image_dir, f"augmented_{i}.jpg")
+            augmented_image.save(augmented_image_path)
+            i += 1
diff --git a/model.py b/model.py
new file mode 100644
index 0000000..cafb438
--- /dev/null
+++ b/model.py
@@ -0,0 +1,249 @@
+import torch
+
+
+class SimpleCNN(torch.nn.Module):
+    def __init__(self,
+                 in_channels: int = 1,
+                 out_channels: int = 64,
+                 hidden_channels: int = 3,
+                 kernel_size: int = 3,
+                 dropout_rate: float = 0.01):
+        super().__init__()
+        layers = []
+        for i in range(hidden_channels):
+            layers.append(torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2))
+            layers.append(torch.nn.ReLU())
+            layers.append(torch.nn.BatchNorm2d(num_features=out_channels))
+            layers.append(torch.nn.Dropout2d(p=dropout_rate))
+            layers.append(torch.nn.MaxPool2d(kernel_size=2, stride=2))
+            layers.append(
+                torch.nn.ConvTranspose2d(in_channels=out_channels, out_channels=out_channels, kernel_size=2, stride=2))
+            in_channels = out_channels
+        self.hidden_layers = torch.nn.Sequential(*layers)
+        self.out_layer = torch.nn.Conv2d(in_channels, 1, kernel_size=kernel_size, padding=kernel_size // 2)
+        self.ac_out = torch.nn.Sigmoid()
+
+    def forward(self, x):
+        normalized_x = x.float() / 255
+        img = self.hidden_layers(normalized_x)
+        output = self.out_layer(img)
+        scaled_output = self.ac_out(output) * 255
+        return scaled_output
+
+
+# ---------- Past Models ----------
+
+# Researched alot about possible CNNs - this is Vanilla UNet Architecture,
+# turns out it doesn't generalize well since its to complex
+class Unet(torch.nn.Module):
+    def __init__(self, input_channels: int, padding='same', kernel_size: int = 3):
+        super().__init__()
+
+        # Encoder
+        self.max_pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
+        self.ac = torch.nn.ReLU()
+        # input size = 1
+        # input shape torch.Tensor[3, 1, 64, 64]
+        print(input_channels)
+        self.conv_1 = torch.nn.Conv2d(in_channels=1,
+                                      out_channels=input_channels,
+                                      kernel_size=kernel_size,
+                                      padding=padding)
+        self.conv_1_1 = torch.nn.Conv2d(in_channels=input_channels,
+                                        out_channels=input_channels,
+                                        kernel_size=kernel_size,
+                                        padding=padding)
+        self.batch_norm_1 = torch.nn.BatchNorm2d(input_channels)
+        # apply conv layer 2 times, ReLu after every application, max pooling after every 2 applications of conv layer
+        input_channels *= 2
+        # input size = 128
+        self.conv_2 = torch.nn.Conv2d(in_channels=input_channels // 2,
+                                      out_channels=input_channels,
+                                      kernel_size=kernel_size,
+                                      padding=padding)
+        self.conv_2_1 = torch.nn.Conv2d(in_channels=input_channels,
+                                        out_channels=input_channels,
+                                        kernel_size=kernel_size,
+                                        padding=padding)
+        self.batch_norm_2 = torch.nn.BatchNorm2d(input_channels)
+        input_channels *= 2
+        # input size = 256
+        self.conv_3 = torch.nn.Conv2d(in_channels=input_channels // 2,
+                                      out_channels=input_channels,
+                                      kernel_size=kernel_size,
+                                      padding=padding)
+        self.conv_3_1 = torch.nn.Conv2d(in_channels=input_channels,
+                                        out_channels=input_channels,
+                                        kernel_size=kernel_size,
+                                        padding=padding)
+        self.batch_norm_3 = torch.nn.BatchNorm2d(input_channels)
+        input_channels *= 2
+        # Bridge
+        # input size = 512
+        self.conv_4 = torch.nn.Conv2d(in_channels=input_channels // 2,
+                                      out_channels=input_channels,
+                                      kernel_size=kernel_size,
+                                      padding=padding)
+        self.conv_4_1 = torch.nn.Conv2d(in_channels=input_channels,
+                                        out_channels=input_channels,
+                                        kernel_size=kernel_size,
+                                        padding=padding)
+        self.batch_norm_4 = torch.nn.BatchNorm2d(input_channels)
+        # Decoder
+        self.up_samp = torch.nn.UpsamplingBilinear2d(scale_factor=2)
+        # input size = 256, kernel size = 3 - 1
+        input_channels //= 2
+        self.conv_5 = torch.nn.Conv2d(in_channels=input_channels * 2,
+                                      out_channels=input_channels,
+                                      kernel_size=kernel_size,
+                                      padding=padding)
+        # concatenate in forward function (tensor after last application of conv 3 merged with current one(up sampled))
+        # input size = 512
+        self.conv_6 = torch.nn.Conv2d(in_channels=input_channels * 2,
+                                      out_channels=input_channels,
+                                      kernel_size=kernel_size,
+                                      padding=padding)
+        self.conv_6_1 = torch.nn.Conv2d(in_channels=input_channels,
+                                        out_channels=input_channels,
+                                        kernel_size=kernel_size,
+                                        padding=padding)
+        self.batch_norm_5 = torch.nn.BatchNorm2d(input_channels)
+        # apply conv_6 2 times
+        # up sampling
+        # input size = 128, kernel size = 3 - 1
+        input_channels //= 2
+        self.conv_7 = torch.nn.Conv2d(in_channels=input_channels * 2,
+                                      out_channels=input_channels,
+                                      kernel_size=kernel_size,
+                                      padding=padding)
+        # concatenate in forward function (tensor after conv 2 with current one)
+        # input size = 128, kernel size = 3
+        self.conv_8 = torch.nn.Conv2d(in_channels=input_channels * 2,
+                                      out_channels=input_channels,
+                                      kernel_size=kernel_size,
+                                      padding=padding)
+        self.conv_8_1 = torch.nn.Conv2d(in_channels=input_channels,
+                                        out_channels=input_channels,
+                                        kernel_size=kernel_size,
+                                        padding=padding)
+        self.batch_norm_6 = torch.nn.BatchNorm2d(input_channels)
+        # up sampling
+        # input size = 64, kernel size = 3 - 1
+        input_channels //= 2
+        self.conv_9 = torch.nn.Conv2d(in_channels=input_channels * 2,
+                                      out_channels=input_channels,
+                                      kernel_size=kernel_size,
+                                      padding=padding)
+        # concatenate in forward function (tensor after conv1 with current one)
+        self.conv_10 = torch.nn.Conv2d(in_channels=input_channels * 2,
+                                       out_channels=input_channels,
+                                       kernel_size=kernel_size,
+                                       padding=padding)
+        self.conv_10_1 = torch.nn.Conv2d(in_channels=input_channels,
+                                         out_channels=input_channels,
+                                         kernel_size=kernel_size,
+                                         padding=padding)
+        self.batch_norm_7 = torch.nn.BatchNorm2d(input_channels)
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d):
+                torch.nn.init.xavier_uniform_(m.weight)
+        self.out = torch.nn.Conv2d(64, 1, kernel_size=1)
+        self.ac_out = torch.nn.Sigmoid()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Encoder
+        x = x.float() / 255
+        img = self.ac(self.conv_1(x))
+        img = self.batch_norm_1(img)
+        conv1 = img = self.ac(self.conv_1_1(img))
+        img = self.batch_norm_1(img)
+        img = self.max_pool(img)
+        img = self.ac(self.conv_2(img))
+        img = self.batch_norm_2(img)
+        conv2 = img = self.ac(self.conv_2_1(img))
+        img = self.batch_norm_2(img)
+        img = self.max_pool(img)
+        img = self.ac(self.conv_3(img))
+        img = self.batch_norm_3(img)
+        conv3 = img = self.ac(self.conv_3_1(img))
+        img = self.batch_norm_3(img)
+        img = self.max_pool(img)
+        # Bridge
+        img = self.ac(self.conv_4(img))
+        img = self.batch_norm_4(img)
+        img = torch.nn.Dropout(p=0.001)(img)
+        img = self.ac(self.conv_4_1(img))
+        img = self.batch_norm_4(img)
+        img = torch.nn.Dropout(p=0.001)(img)
+        # Decoder
+        img = self.up_samp(img)
+        img = self.ac(self.conv_5(img))
+        img = self.batch_norm_5(img)
+        img = torch.cat([conv3, img], dim=1)
+        img = self.ac(self.conv_6(img))
+        img = self.batch_norm_5(img)
+        img = self.ac(self.conv_6_1(img))
+        img = self.batch_norm_5(img)
+
+        img = self.up_samp(img)
+        img = self.ac(self.conv_7(img))
+        img = self.batch_norm_6(img)
+        img = torch.concatenate([conv2, img], dim=1)
+        img = self.ac(self.conv_8(img))
+        img = self.batch_norm_6(img)
+        img = self.ac(self.conv_8_1(img))
+        img = self.batch_norm_6(img)
+
+        img = self.up_samp(img)
+        img = self.ac(self.conv_9(img))
+        img = self.batch_norm_7(img)
+        img = torch.cat([conv1, img], dim=1)
+        img = self.ac(self.conv_10(img))
+        img = self.batch_norm_7(img)
+        img = self.ac(self.conv_10_1(img))
+        img = self.batch_norm_7(img)
+
+        # Output
+        output = self.ac_out(self.out(img))
+        return output
+
+
+# Testing around with basic idea of UNet above but decreasing model complexity - did not work to well either
+class Model(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.in_chn = torch.nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, padding=1)
+        self.bn2d_64 = torch.nn.BatchNorm2d(num_features=64)
+        self.ac = torch.nn.ReLU()
+        self.max_pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
+        self.conv_64_128 = torch.nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
+        self.bn2d_128 = torch.nn.BatchNorm2d(num_features=128)
+        self.conv_128_256 = torch.nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1)
+        self.bn2d_256 = torch.nn.BatchNorm2d(num_features=256)
+
+        self.conv_256_512 = torch.nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1)
+        self.bn2d_512 = torch.nn.BatchNorm2d(num_features=512)
+
+        # concat
+        self.conv_256_128 = torch.nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, padding=1)
+        self.conv_128_64 = torch.nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=1)
+        self.out_chn = torch.nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1)
+        self.ac_out = torch.nn.Sigmoid()
+
+    def forward(self, x):
+        img = self.bn2d_64(self.ac(self.in_chn(x)))
+        state_1 = img = self.bn2d_128(self.ac(self.conv_64_128(img)))
+        img = torch.nn.Dropout(p=0.08)(img)
+        img = self.bn2d_256(self.ac(self.conv_128_256(img)))
+        img = torch.nn.Dropout(p=0.01)(img)
+        img = self.bn2d_256(self.ac(self.conv_128_256(img)))
+        img = torch.nn.Dropout(p=0.01)(img)
+        img = self.bn2d_128(self.ac(self.conv_256_128(img)))
+        img = torch.nn.Dropout(p=0.1)(img)
+        img = torch.concat([state_1, img], dim=1)
+        img = self.bn2d_128(self.ac(self.conv_256_128(img)))
+        img = torch.nn.Dropout(p=0.01)(img)
+        img = self.bn2d_64(self.ac(self.conv_128_64(img)))
+        img = torch.nn.Dropout(p=0.08)(img)
+        img = self.out_chn(img)
+        return self.ac_out(img)
diff --git a/models/model_14.pth b/models/model_14.pth
new file mode 100644
index 0000000..d1df5de
Binary files /dev/null and b/models/model_14.pth differ
diff --git a/models/model_15.pth b/models/model_15.pth
new file mode 100644
index 0000000..775b415
Binary files /dev/null and b/models/model_15.pth differ
diff --git a/models/model_16.pth b/models/model_16.pth
new file mode 100644
index 0000000..4b5d500
Binary files /dev/null and b/models/model_16.pth differ
diff --git a/models/model_17.pth b/models/model_17.pth
new file mode 100644
index 0000000..7b710c5
Binary files /dev/null and b/models/model_17.pth differ
diff --git a/models/model_18.pth b/models/model_18.pth
new file mode 100644
index 0000000..e249acb
Binary files /dev/null and b/models/model_18.pth differ
diff --git a/models/model_19.pth b/models/model_19.pth
new file mode 100644
index 0000000..61605fd
Binary files /dev/null and b/models/model_19.pth differ
diff --git a/models/model_20.pth b/models/model_20.pth
new file mode 100644
index 0000000..707e69c
Binary files /dev/null and b/models/model_20.pth differ
diff --git a/prep_image.py b/prep_image.py
new file mode 100644
index 0000000..441b466
--- /dev/null
+++ b/prep_image.py
@@ -0,0 +1,87 @@
+import numpy as np
+from torchvision import transforms
+from PIL import Image
+
+
+def transform_image(image: Image) -> Image:
+    im_shape = 64
+    resize_transforms = transforms.Compose([
+        transforms.Resize(size=im_shape),
+        transforms.CenterCrop(size=(im_shape, im_shape))
+    ])
+    return resize_transforms(image)
+
+
+def to_grayscale(pil_image: np.ndarray) -> np.ndarray:
+    if pil_image.ndim == 2:
+        return pil_image.copy()[None]
+    if pil_image.ndim != 3:
+        raise ValueError("image must have either shape (H, W) or (H, W, 3)")
+    if pil_image.shape[2] != 3:
+        raise ValueError(f"image has shape (H, W, {pil_image.shape[2]}), but it should have (H, W, 3)")
+
+    rgb = pil_image / 255
+    rgb_linear = np.where(
+        rgb < 0.04045,
+        rgb / 12.92,
+        ((rgb + 0.055) / 1.055) ** 2.4
+    )
+    grayscale_linear = 0.2126 * rgb_linear[..., 0] + 0.7152 * rgb_linear[..., 1] + 0.0722 * rgb_linear[..., 2]
+
+    grayscale = np.where(
+        grayscale_linear < 0.0031308,
+        12.92 * grayscale_linear,
+        1.055 * grayscale_linear ** (1 / 2.4) - 0.055
+    )
+    grayscale = grayscale * 255
+
+    if np.issubdtype(pil_image.dtype, np.integer):
+        grayscale = np.round(grayscale)
+    return grayscale.astype(pil_image.dtype)[None]
+
+
+def prepare_image(image: np.ndarray, x: int, y: int, width: int, height: int, size: int) -> \
+        tuple[np.ndarray, np.ndarray, np.ndarray]:
+    if image.ndim < 3 or image.shape[-3] != 1:
+        # This is actually more general than the assignment specification
+        raise ValueError("image must have shape (..., 1, H, W)")
+    if width < 2 or height < 2 or size < 2:
+        raise ValueError("width/height/size must be >= 2")
+    if x < 0 or (x + width) > image.shape[-1]:
+        raise ValueError(f"x={x} and width={width} do not fit into the image width={image.shape[-1]}")
+    if y < 0 or (y + height) > image.shape[-2]:
+        raise ValueError(f"y={y} and height={height} do not fit into the image height={image.shape[-2]}")
+
+    # The (height, width) slices to extract the area that should be pixelated. Since we
+    # need this multiple times, specify the slices explicitly instead of using [:] notation
+    area = (..., slice(y, y + height), slice(x, x + width))
+
+    # This returns already a copy, so we are independent of "image"
+    pixelated_image = pixelate(image, x, y, width, height, size)
+
+    known_array = np.full_like(image, fill_value=False, dtype=bool)
+    known_array[area] = True
+
+    # Create a copy to avoid that "target_array" and "image" point to the same array
+    # target_array = image[area].copy()
+    stacked_target_array = np.full((1, 64, 64), 0)
+    stacked_target_array[area] = image[area].copy()
+    target_array = stacked_target_array
+
+    return pixelated_image, known_array, target_array
+
+
+def pixelate(image: np.ndarray, x: int, y: int, width: int, height: int, size: int) -> np.ndarray:
+    # Need a copy since we overwrite data directly
+    image = image.copy()
+    curr_x = x
+
+    while curr_x < x + width:
+        curr_y = y
+        while curr_y < y + height:
+            block = (..., slice(curr_y, min(curr_y + size, y + height)), slice(curr_x, min(curr_x + size, x + width)))
+            image[block] = image[block].mean()
+            curr_y += size
+        curr_x += size
+
+    return image
diff --git a/stack_with_padding.py b/stack_with_padding.py
new file mode 100644
index 0000000..b5f5766
--- /dev/null
+++ b/stack_with_padding.py
@@ -0,0 +1,26 @@
+import torch
+from typing import List, Tuple
+import numpy as np
+
+
+def stack_with_padding(batch_as_list: List[Tuple[np.ndarray, np.ndarray, np.ndarray, str]]) -> Tuple[
+        torch.Tensor, torch.Tensor, torch.Tensor, List[str]]:
+    batch_size = len(batch_as_list)
+    max_height = max(pixelated_image.shape[1] for pixelated_image, _, _, _ in batch_as_list)
+    max_width = max(pixelated_image.shape[-1] for pixelated_image, _, _, _ in batch_as_list)
+
+    stacked_pixelated_images = np.ones((batch_size, 1, max_height, max_width))
+    stacked_known_arrays = np.ones((batch_size, 1, max_height, max_width))
+    stacked_target_arrays = np.ones((batch_size, 1, max_height, max_width))
+    for i, (img, arr, tar, _) in enumerate(batch_as_list):
+        stacked_pixelated_images[i, :, :img.shape[1], :img.shape[2]] = img
+        stacked_known_arrays[i, :, :arr.shape[1], : arr.shape[2]] = arr
+        stacked_target_arrays[i, :, :arr.shape[1], : arr.shape[2]] = tar
+
+    stacked_pixelated_images = torch.Tensor(np.stack(list(image for image in stacked_pixelated_images)))
+    stacked_known_arrays = torch.Tensor(np.stack(list(arr for arr in stacked_known_arrays)))
+    stacked_target_arrays = torch.Tensor(np.stack(list(tar for tar in stacked_target_arrays)))
+
+    image_files = [image_file for _, _, _, image_file in batch_as_list]
+
+    return stacked_pixelated_images, stacked_known_arrays, stacked_target_arrays, image_files
diff --git a/training_loop.py b/training_loop.py
new file mode 100644
index 0000000..70d7ec6
--- /dev/null
+++ b/training_loop.py
@@ -0,0 +1,131 @@
+import numpy as np
+import torch
+from torch.utils import data
+from typing import Tuple
+import matplotlib.pyplot as plt
+from stack_with_padding import stack_with_padding
+from torch.backends import mps
+from tqdm import tqdm
+
+
+def training_loop(
+        network: torch.nn.Module,
+        train_data: torch.utils.data.Dataset,
+        eval_data: torch.utils.data.Dataset,
+        num_epochs: int,
+        show_progress: bool = False
+) -> Tuple[list, list]:
+    optimizer = torch.optim.Adam(params=network.parameters(), lr=1e-5, weight_decay=1e-5)
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset=train_data,
+        batch_size=64,
+        shuffle=True,
+        collate_fn=stack_with_padding,
+        num_workers=10,
+        pin_memory=True
+    )
+    data_loader_eval = torch.utils.data.DataLoader(
+        dataset=eval_data,
+        batch_size=64,
+        collate_fn=stack_with_padding,
+        num_workers=10,
+        pin_memory=True
+    )
+    # cuda can be added if executed on machine with cuda available
+    use_mps = torch.backends.mps.is_available()
+    device = torch.device("mps" if use_mps else "cpu")
+    network.to(device)
+    epoch_losses = []
+    eval_losses = []
+    # init progress bar
+    progress_bar = None
+    if show_progress:
+        total_iterations = num_epochs * (len(data_loader_train) + len(data_loader_eval))
+        progress_bar = tqdm(total=total_iterations, desc="Training and Evaluation")
+    for epoch in range(num_epochs):
+        mbl = []
+        network.train()
+        for inputs, known, targets, _ in data_loader_train:
+            inputs = inputs.to(device).float()
+            known = known.to(device).float()
+            targets = targets.to(device).float()
+            optimizer.zero_grad()
+            output = network(inputs)
+            known_tensor = known.bool()
+            # Extract only pixelated area from output / targets for loss calculation
+            target = targets[known_tensor]
+            out = output[known_tensor]
+            rmse_loss = torch.nn.MSELoss()(out, target)
+            rmse_loss = torch.sqrt(rmse_loss)
+            rmse_loss.backward()
+            optimizer.step()
+
+            mbl.append(rmse_loss.item())
+            if show_progress:
+                progress_bar.update(1)
+        epoch_losses.append(np.mean(mbl))
+        # Save model state after each epoch, can be optimized, good for now
+        torch.save(network.state_dict(), f"model_{epoch + 1}.pth")
+
+        network.eval()
+        with torch.no_grad():
+            eval_mlb = []
+            for inputs, known, targets, _ in data_loader_eval:
+                inputs = inputs.to(device).float()
+                known = known.to(device).float()
+                targets = targets.to(device).float()
+                output = network(inputs)
+                known_tensor = known.bool()
+                target = targets[known_tensor]
+                out = output[known_tensor]
+                loss = torch.nn.MSELoss()(out, target)
+                loss = torch.sqrt(loss)
+                eval_mlb.append(loss.item())
+                if show_progress:
+                    progress_bar.update(1)
+            eval_losses.append(np.mean(eval_mlb))
+    return epoch_losses, eval_losses
+
+
+if __name__ == "__main__":
+    from model import SimpleCNN
+    from create_dataset import RandomImagePixelationDataset
+
+    torch.random.manual_seed(0)
+    # training_raw dir contains 300 folders from provided images and one folder with all images from these 300 folders,
+    # but augmented (sums up to about 60k images). (also see file image_augmentation)
+    train_data = RandomImagePixelationDataset(
+        r"training_raw",
+        width_range=(4, 32),
+        height_range=(4, 32),
+        size_range=(4, 16)
+    )
+    # eval_raw dir contains the remaining 50 folders for evaluation (~5k images)
+    eval_data = RandomImagePixelationDataset(
+        r"eval_raw",
+        width_range=(4, 32),
+        height_range=(4, 32),
+        size_range=(4, 16)
+    )
+    # working config: in_channels=1, out_channels=128, hidden_channels=6, kernel_size=7, dropout_rate=0.01
+    network = SimpleCNN(in_channels=1, out_channels=128, hidden_channels=6, kernel_size=7, dropout_rate=0.01)
+
+    model_param = filter(lambda p: p.requires_grad, network.parameters())
+    params = sum(np.prod(p.size()) for p in model_param)
+    print(f"Model has {params} parameters.")
+
+    epochs, train_loss, eval_loss = [], [], []
+    train_losses, eval_losses = training_loop(network, train_data, eval_data,
+                                              num_epochs=20, show_progress=True)
+    for epoch, (tl, el) in enumerate(zip(train_losses, eval_losses)):
+        print(f"Epoch: {epoch} --- Train loss: {tl:7.2f} --- Eval loss: {el:7.2f}")
+        epochs.append(epoch)
+        train_loss.append(tl)
+        eval_loss.append(el)
+
+    plt.plot(epochs, train_loss, label='train')
+    plt.plot(epochs, eval_loss, label='eval')
+    plt.xlabel('Epoch')
+    plt.ylabel('Loss')
+    plt.title('Loss per Epoch')
+    plt.show()