diff --git a/ApplyModel.py b/ApplyModel.py new file mode 100644 index 0000000..bf53f94 --- /dev/null +++ b/ApplyModel.py @@ -0,0 +1,29 @@ +import numpy as np +import torch +from PIL import Image + +import DataLoader +import ex4 +from ImageImpaint import get_train_device +from netio import load_model + + +def apply_model(filepath: str): + device = get_train_device() + + img = Image.open(filepath) + model = load_model() + model.to(device) + + pic = DataLoader.preprocess(img, precision=np.float32) + pic = ex4.ex4(pic, (5, 5), (4, 4))[0] + Image.fromarray((np.transpose(pic * 255.0, (1, 2, 0)).astype(np.uint8))).save("filename_grid.jpg") + out = model(torch.from_numpy(pic).to(device)) + out = DataLoader.postprocess(out.cpu().detach().numpy()) + out = np.transpose(out, (1, 2, 0)) + im = Image.fromarray(out) + im.save("filename.jpg", format="jpeg") + + +if __name__ == '__main__': + apply_model("training/000/000017.jpg") diff --git a/DataLoader.py b/DataLoader.py index 26df549..edb0f28 100644 --- a/DataLoader.py +++ b/DataLoader.py @@ -14,22 +14,15 @@ IMG_SIZE = 100 class ImageDataset(Dataset): - def __init__(self, image_dir): + def __init__(self, image_dir, precision: np.float32 or np.float64): 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) + self.precision = precision def __getitem__(self, index): # Open image file, convert to numpy array and scale to [0, 1] target_image = Image.open(self.image_files[index]) - # image = np.array(Image.open(self.image_files[index]), dtype=np.float32) / 255 - resize_transforms = transforms.Compose([ - transforms.Resize(size=IMG_SIZE), - transforms.CenterCrop(size=(IMG_SIZE, IMG_SIZE)), - ]) - target_image = resize_transforms(target_image) - target_image = preprocess(target_image) + + target_image = preprocess(target_image, self.precision) # calculate image with black grid doomed_image = ex4.ex4(target_image, (5, 5), (4, 4)) @@ -43,9 +36,16 @@ class ImageDataset(Dataset): return len(self.image_files) -def preprocess(input: np.array) -> np.array: +def preprocess(input: np.array, precision: np.float32 or np.float64) -> np.array: + # image = np.array(Image.open(self.image_files[index]), dtype=np.float32) / 255 + resize_transforms = transforms.Compose([ + transforms.Resize(size=IMG_SIZE), + transforms.CenterCrop(size=(IMG_SIZE, IMG_SIZE)), + ]) + input = resize_transforms(input) + # normalize image from 0-1 - target_image = np.array(input, dtype=np.float64) / 255.0 + target_image = np.array(input, dtype=precision) / 255.0 # Perform normalization for each channel # image = (image - self.norm_mean) / self.norm_std @@ -59,33 +59,27 @@ def postprocess(input: np.array) -> np.array: return target_image -def get_image_loader(path: str): - image_dataset = ImageDataset(path) +def get_image_loader(path: str, precision: np.float32 or np.float64): + image_dataset = ImageDataset(path, precision) totlen = len(image_dataset) - test_set_size = .001 + test_set_size = .1 trains, tests = torch.utils.data.dataset.random_split(image_dataset, lengths=(totlen - int(totlen * test_set_size), int(totlen * test_set_size)), - generator=torch.Generator().manual_seed(42)) + generator=torch.Generator().manual_seed(0)) train_loader = DataLoader( trains, shuffle=True, # shuffle the order of our samples - batch_size=5, # stack 4 samples to a minibatch + batch_size=25, # stack 4 samples to a minibatch num_workers=4 # no background workers (see comment below) ) test_loader = DataLoader( tests, shuffle=True, # shuffle the order of our samples - batch_size=1, # stack 4 samples to a minibatch + batch_size=5, # stack 4 samples to a minibatch num_workers=0 # no background workers (see comment below) ) return train_loader, test_loader - - -def rgb2gray(rgb_array: np.ndarray): - r, g, b = rgb_array[:, :, 0], rgb_array[:, :, 1], rgb_array[:, :, 2] - gray = 0.2989 * r + 0.5870 * g + 0.1140 * b - return gray diff --git a/ImageImpaint.py b/ImageImpaint.py index 208635b..bba021f 100644 --- a/ImageImpaint.py +++ b/ImageImpaint.py @@ -1,11 +1,13 @@ import numpy as np import torch +from PIL.Image import Image +import DataLoader from DataLoader import get_image_loader from Net import ImageNN # 01.05.22 -- 0.5h -from netio import save_model, load_model +from netio import save_model, load_model, eval_evalset def get_train_device(): @@ -21,17 +23,15 @@ def train_model(): device = get_train_device() # Load datasets - train_loader, test_loader = get_image_loader("training/") - nn = ImageNN(n_in_channels=3) # todo pass size ason. + train_loader, test_loader = get_image_loader("training/", precision=np.float32) + nn = ImageNN(n_in_channels=3, precision=np.float32) # todo net params nn.train() # init with train mode nn.to(device) # send net to device available optimizer = torch.optim.AdamW(nn.parameters(), lr=0.1, weight_decay=1e-5) # todo adjust parameters and lr loss_function = torch.nn.MSELoss() - n_epochs = 10 # todo epcchs here - - # todo look wtf is that - nn.double() + loss_function.to(device) + n_epochs = 7 # todo epcchs here train_sample_size = len(train_loader) losses = [] @@ -40,23 +40,21 @@ def train_model(): print(f"Epoch {epoch}/{n_epochs}\n") i = 0 for input_tensor, target_tensor in train_loader: - input_tensor.to(device) - target_tensor.to(device) + output = nn(input_tensor.to(device)) # get model output (forward pass) - output = nn(input_tensor) # get model output (forward pass) - loss = loss_function(output, target_tensor) # compute loss given model output and true target + loss = loss_function(output.to(device), target_tensor.to(device)) # 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()) + i += train_loader.batch_size print( f'\rTraining epoch {epoch} [{i}/{train_sample_size * train_loader.batch_size}] (curr loss: {loss.item():.3})', end='') - i += train_loader.batch_size # eval model every 3000th sample - if i % 15 == 0: + if i % 3000 == 0: print(f"\nEvaluating model") eval_loss = eval_model(nn, test_loader, loss_function, device) print(f"Evalution loss={eval_loss}") @@ -64,8 +62,10 @@ def train_model(): best_eval_loss = eval_loss save_model(nn) - # switch net to eval mode - print(eval_model(nn, test_loader, loss_function, device=device)) + nn.train() + + # evaluate model with submission pkl file + eval_evalset() # func to evaluate our trained model @@ -77,8 +77,8 @@ def eval_model(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, with torch.no_grad(): i = 0 for input, target in dataloader: - input.to(device) - target.to(device) + input = input.to(device) + target = target.to(device) out = model(input) loss += loss_fn(out, target).item() @@ -86,11 +86,9 @@ def eval_model(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, i += dataloader.batch_size print() loss /= len(dataloader) - model.train() return loss -def apply_model(): - model = load_model() - pass +if __name__ == '__main__': + train_model() diff --git a/Net.py b/Net.py index 4766c9c..dd0b388 100644 --- a/Net.py +++ b/Net.py @@ -1,8 +1,9 @@ +import numpy as np import torch class ImageNN(torch.nn.Module): - def __init__(self, n_in_channels: int = 1, n_hidden_layers: int = 3, n_kernels: int = 32, kernel_size: int = 7): + def __init__(self, precision: np.float32 or np.float64, n_in_channels: int = 1, n_hidden_layers: int = 3, n_kernels: int = 32, kernel_size: int = 7): """Simple CNN with `n_hidden_layers`, `n_kernels`, and `kernel_size` as hyperparameters""" super().__init__() @@ -25,6 +26,9 @@ class ImageNN(torch.nn.Module): padding=int(kernel_size / 2) ) + if precision == np.float64: + self.double() + def forward(self, x): """Apply CNN to input `x` of shape (N, n_channels, X, Y), where N=n_samples and X, Y are spatial dimensions""" cnn_out = self.hidden_layers(x) # apply hidden layers (N, n_in_channels, X, Y) -> (N, n_kernels, X, Y) diff --git a/main.py b/main.py deleted file mode 100644 index 269f2be..0000000 --- a/main.py +++ /dev/null @@ -1,11 +0,0 @@ -# This is a sample Python script. - -# Press Umschalt+F10 to execute it or replace it with your code. -# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings. - - -# See PyCharm help at https://www.jetbrains.com/help/pycharm/ -from ImageImpaint import train_model - -if __name__ == '__main__': - train_model() diff --git a/netio.py b/netio.py index 0a144fb..d9f8e30 100644 --- a/netio.py +++ b/netio.py @@ -1,5 +1,5 @@ +import os import pickle -import sys import numpy as np import torch @@ -15,32 +15,40 @@ def save_model(model: torch.nn.Module): print(f"Saved raw model to {MODEL_PATH}") torch.save(model, MODEL_PATH) + +def eval_evalset(): # read the provided testing pickle file print("Generating pickle file with privided test data") + try: + os.unlink(PICKEL_PATH) + except: + pass + + model = load_model() model.eval() with open('testing/inputs.pkl', 'rb') as handle: - with open(PICKEL_PATH, 'wb') as writehandle: - b: dict = pickle.load(handle) - outarr = [] - i=0 - piclen = len(b['input_arrays']) - for pic in b['input_arrays']: - pic = DataLoader.preprocess(pic) - out = model(torch.from_numpy(pic)) - out = DataLoader.postprocess(out.detach().numpy()) - pickle.dump(out, writehandle, protocol=pickle.HIGHEST_PROTOCOL) + b: dict = pickle.load(handle) + outarr = np.zeros(dtype=np.uint8, shape=(len(b['input_arrays']), 3, 100, 100)) + i = 0 + piclen = len(b['input_arrays']) + for pic in b['input_arrays']: + pic = DataLoader.preprocess(pic, precision=np.float32) + out = model(torch.from_numpy(pic)) + out = DataLoader.postprocess(out.cpu().detach().numpy()) + outarr[i] = out - print( - f'\rApplying model [{i}/{piclen}] {sys.getsizeof(outarr)}',end='') - i += 1 + print(f'\rApplying model [{i}/{piclen}]', end='') + i += 1 + write_to_pickle(PICKEL_PATH, list(outarr)) # compress the generated pickle arr Compress.compress(PICKEL_PATH) -def load_model(): - model = ImageNN() - model.load_state_dict(torch.load(MODEL_PATH)) - model.eval() - return model +def write_to_pickle(filename: str, data): + with open(filename, 'wb') as handle: + pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL) + +def load_model(): + return torch.load(MODEL_PATH)