From 2966df9a39fa0ffc645e1738f9af4dcbf086ea84 Mon Sep 17 00:00:00 2001
From: lukas-heiligenbrunner <lukas.heiligenbrunner@gmail.com>
Date: Wed, 1 Jun 2022 12:27:58 +0200
Subject: [PATCH] implement basic structure of project

---
 DataLoader.py   | 31 +++++++++++++++++++++++++++++++
 ImageImpaint.py | 32 ++++++++++++++++++++++++++++++++
 Net.py          | 11 +++++++++++
 main.py         |  4 ++++
 4 files changed, 78 insertions(+)
 create mode 100644 DataLoader.py
 create mode 100644 ImageImpaint.py
 create mode 100644 Net.py

diff --git a/DataLoader.py b/DataLoader.py
new file mode 100644
index 0000000..7499ea4
--- /dev/null
+++ b/DataLoader.py
@@ -0,0 +1,31 @@
+import glob
+import os
+
+import numpy as np
+from PIL import Image
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+
+
+class ImageDataset(Dataset):
+    def __init__(self, image_dir):
+        self.image_files = sorted(glob.glob(os.path.join(image_dir, "**", "*.jpg"), recursive=True))
+        # Mean and std arrays could also be defined as class attributes
+        self.norm_mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+        self.norm_std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+    def __getitem__(self, index):
+        # Open image file, convert to numpy array and scale to [0, 1]
+        image = np.array(Image.open(self.image_files[index]), dtype=np.float32) / 255
+        # Perform normalization for each channel
+        image = (image - self.norm_mean) / self.norm_std
+        return image, index
+
+    def __len__(self):
+        return len(self.image_files)
+
+
+def get_image_loader(path: str):
+    image_dataset = ImageDataset(path)
+    image_loader = DataLoader(image_dataset, shuffle=True, batch_size=10)
+    return image_loader
diff --git a/ImageImpaint.py b/ImageImpaint.py
new file mode 100644
index 0000000..1eb5919
--- /dev/null
+++ b/ImageImpaint.py
@@ -0,0 +1,32 @@
+import torch
+
+from DataLoader import get_image_loader
+from Net import ImageNN
+
+
+def train_model():
+    image_loader = get_image_loader("my/supercool/image/dir")
+    # todo split to train and test (maybe evaluation sets)
+    nn = ImageNN() # todo pass size ason.
+
+    optimizer = torch.optim.SGD(nn.parameters(), lr=0.1) # todo adjust parameters and lr
+    loss_function = torch.nn.CrossEntropyLoss()
+    n_epochs = 15 # todo epcchs here
+
+    # Training
+    losses = []
+    for epoch in range(n_epochs):
+        for input_tensor, target_tensor in image_loader:
+            output = nn(input_tensor)  # get model output (forward pass)
+            loss = loss_function(output, target_tensor)  # compute loss given model output and true target
+            loss.backward()  # compute gradients (backward pass)
+            optimizer.step()  # perform gradient descent update step
+            optimizer.zero_grad()  # reset gradients
+            losses.append(loss.item())
+
+    # todo evaluate trained model
+    # todo save trained model to blob file
+
+
+def apply_model():
+    pass
diff --git a/Net.py b/Net.py
new file mode 100644
index 0000000..6b66150
--- /dev/null
+++ b/Net.py
@@ -0,0 +1,11 @@
+import torch
+
+
+class ImageNN(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        # todo implement the nn structure
+
+    def forward(self, x: torch.Tensor):
+        pass
+        # todo implement forward
diff --git a/main.py b/main.py
index 142cf01..269f2be 100644
--- a/main.py
+++ b/main.py
@@ -5,3 +5,7 @@
 
 
 # See PyCharm help at https://www.jetbrains.com/help/pycharm/
+from ImageImpaint import train_model
+
+if __name__ == '__main__':
+    train_model()