From 0f0c7899813143a7ad9fba5b1bb4ba281636e2b4 Mon Sep 17 00:00:00 2001 From: lukas-heiligenbrunner Date: Sun, 10 Jul 2022 01:22:28 +0200 Subject: [PATCH] add mask to training export trainpickle file in correct format --- ApplyModel.py | 2 +- DataLoader.py | 2 +- ImageImpaint.py | 30 ++++++++++++++++-------------- Net.py | 6 +++++- netio.py | 18 +++++++++++++----- 5 files changed, 36 insertions(+), 22 deletions(-) diff --git a/ApplyModel.py b/ApplyModel.py index eec2a97..e990cee 100644 --- a/ApplyModel.py +++ b/ApplyModel.py @@ -60,6 +60,6 @@ def test(): Compress.compress(PICKEL_PATH + "_target.pkl") if __name__ == '__main__': - apply_model("training/000/000017.jpg") + # apply_model("training/000/000017.jpg") eval_evalset() # test() diff --git a/DataLoader.py b/DataLoader.py index a8110bf..5d71143 100644 --- a/DataLoader.py +++ b/DataLoader.py @@ -37,7 +37,7 @@ class ImageDataset(Dataset): spacing = (random.randint(*self.spacingrange), random.randint(*self.spacingrange)) doomed_image = ex4.ex4(target_image, offset, spacing) - return doomed_image[0], np.transpose(target_image, (2, 0, 1)) + return doomed_image[0], doomed_image[1], np.transpose(target_image, (2, 0, 1)) def __len__(self): return len(self.image_files) diff --git a/ImageImpaint.py b/ImageImpaint.py index e24f7b9..65fcff4 100644 --- a/ImageImpaint.py +++ b/ImageImpaint.py @@ -3,7 +3,6 @@ import sys import PIL import numpy as np -import packaging import torch from matplotlib import pyplot as plt from packaging.version import Version @@ -32,7 +31,8 @@ def train_model(): # Load datasets train_loader, test_loader = get_image_loader("training/", precision=np.float32) - nn = ImageNN(n_in_channels=3, precision=np.float32) # todo net params + + nn = ImageNN(n_in_channels=6, precision=np.float32) # todo net params nn.train() # init with train modeAdam nn.to(device) # send net to device available @@ -47,8 +47,9 @@ def train_model(): for epoch in range(n_epochs): print(f"Epoch {epoch}/{n_epochs}\n") i = 0 - for input_tensor, target_tensor in train_loader: + for input_tensor, mask, target_tensor in train_loader: optimizer.zero_grad() # reset gradients + input_tensor = torch.cat((input_tensor, mask), 1) output = nn(input_tensor.to(device)) # get model output (forward pass) @@ -77,8 +78,8 @@ def train_model(): # Plot output if i % 100 == 0: - plot(input_tensor.detach().cpu().numpy()[:1], target_tensor.detach().cpu().numpy()[:1], - output.detach().cpu().numpy()[:1], + plot(input_tensor.detach().cpu().numpy()[0], target_tensor.detach().cpu().numpy()[0], + output.detach().cpu().numpy()[0], plotpath, i, epoch) # evaluate model with submission pkl file @@ -93,10 +94,12 @@ def eval_model(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, # disable gradient calculations with torch.no_grad(): i = 0 - for input, target in dataloader: + for input, mask, target in dataloader: input = input.to(device) target = target.to(device) + mask = mask.to(device) + input = torch.cat((input, mask), 1) out = model(input) loss += loss_fn(out, target).item() print(f'\rEval prog[{i}/{len(dataloader) * dataloader.batch_size}]', end='') @@ -111,14 +114,13 @@ def plot(inputs, targets, predictions, path, update, epoch): os.makedirs(path, exist_ok=True) fig, axes = plt.subplots(ncols=3, figsize=(15, 5)) - for i in range(len(inputs)): - for ax, data, title in zip(axes, [inputs, targets, predictions], ["Input", "Target", "Prediction"]): - ax.clear() - ax.set_title(title) - ax.imshow(DataLoader.postprocess(np.transpose(data[i], (1, 2, 0))), interpolation="none") - # ax.imshow(np.transpose((data[i]), (1, 2, 0)), interpolation="none") - ax.set_axis_off() - fig.savefig(os.path.join(path, f"{epoch:02d}_{update:07d}_{i:02d}.png"), dpi=100) + for ax, data, title in zip(axes, [inputs, targets, predictions], ["Input", "Target", "Prediction"]): + ax.clear() + ax.set_title(title) + ax.imshow(DataLoader.postprocess(np.transpose(data[:3, :, :], (1, 2, 0))), interpolation="none") + # ax.imshow(np.transpose((data[i]), (1, 2, 0)), interpolation="none") + ax.set_axis_off() + fig.savefig(os.path.join(path, f"{epoch:02d}_{update:07d}.png"), dpi=100) plt.close(fig) diff --git a/Net.py b/Net.py index dd0b388..60dd649 100644 --- a/Net.py +++ b/Net.py @@ -1,9 +1,13 @@ +import math + import numpy as np import torch +from torch import nn class ImageNN(torch.nn.Module): - 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): + 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__() diff --git a/netio.py b/netio.py index d9f8e30..6da3a45 100644 --- a/netio.py +++ b/netio.py @@ -15,6 +15,9 @@ def save_model(model: torch.nn.Module): print(f"Saved raw model to {MODEL_PATH}") torch.save(model, MODEL_PATH) + dummy_input = torch.randn(1, 6, 100, 100) + torch.onnx.export(model, dummy_input, MODEL_PATH + ".onnx", verbose=False, opset_version=11) + def eval_evalset(): # read the provided testing pickle file @@ -28,14 +31,19 @@ def eval_evalset(): model.eval() with open('testing/inputs.pkl', 'rb') as handle: b: dict = pickle.load(handle) - outarr = np.zeros(dtype=np.uint8, shape=(len(b['input_arrays']), 3, 100, 100)) + outarr = () 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)) + for input_array, known_array in zip(b['input_arrays'], b['known_arrays']): + input_array = DataLoader.preprocess(input_array, precision=np.float32) + input_tensor = torch.cat((torch.from_numpy(input_array), torch.from_numpy(known_array)), 0) + out = model(input_tensor) out = DataLoader.postprocess(out.cpu().detach().numpy()) - outarr[i] = out + + rest = out * (1 - known_array) + rest = rest[1 - known_array > 0] + + outarr = (*outarr, rest) print(f'\rApplying model [{i}/{piclen}]', end='') i += 1