add mask to training

export trainpickle file in correct format
This commit is contained in:
lukas-heiligenbrunner 2022-07-10 01:22:28 +02:00
parent 1add9d278f
commit 0f0c789981
5 changed files with 36 additions and 22 deletions

View File

@ -60,6 +60,6 @@ def test():
Compress.compress(PICKEL_PATH + "_target.pkl") Compress.compress(PICKEL_PATH + "_target.pkl")
if __name__ == '__main__': if __name__ == '__main__':
apply_model("training/000/000017.jpg") # apply_model("training/000/000017.jpg")
eval_evalset() eval_evalset()
# test() # test()

View File

@ -37,7 +37,7 @@ class ImageDataset(Dataset):
spacing = (random.randint(*self.spacingrange), random.randint(*self.spacingrange)) spacing = (random.randint(*self.spacingrange), random.randint(*self.spacingrange))
doomed_image = ex4.ex4(target_image, offset, spacing) 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): def __len__(self):
return len(self.image_files) return len(self.image_files)

View File

@ -3,7 +3,6 @@ import sys
import PIL import PIL
import numpy as np import numpy as np
import packaging
import torch import torch
from matplotlib import pyplot as plt from matplotlib import pyplot as plt
from packaging.version import Version from packaging.version import Version
@ -32,7 +31,8 @@ def train_model():
# Load datasets # Load datasets
train_loader, test_loader = get_image_loader("training/", precision=np.float32) 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.train() # init with train modeAdam
nn.to(device) # send net to device available nn.to(device) # send net to device available
@ -47,8 +47,9 @@ def train_model():
for epoch in range(n_epochs): for epoch in range(n_epochs):
print(f"Epoch {epoch}/{n_epochs}\n") print(f"Epoch {epoch}/{n_epochs}\n")
i = 0 i = 0
for input_tensor, target_tensor in train_loader: for input_tensor, mask, target_tensor in train_loader:
optimizer.zero_grad() # reset gradients optimizer.zero_grad() # reset gradients
input_tensor = torch.cat((input_tensor, mask), 1)
output = nn(input_tensor.to(device)) # get model output (forward pass) output = nn(input_tensor.to(device)) # get model output (forward pass)
@ -77,8 +78,8 @@ def train_model():
# Plot output # Plot output
if i % 100 == 0: if i % 100 == 0:
plot(input_tensor.detach().cpu().numpy()[:1], target_tensor.detach().cpu().numpy()[:1], plot(input_tensor.detach().cpu().numpy()[0], target_tensor.detach().cpu().numpy()[0],
output.detach().cpu().numpy()[:1], output.detach().cpu().numpy()[0],
plotpath, i, epoch) plotpath, i, epoch)
# evaluate model with submission pkl file # 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 # disable gradient calculations
with torch.no_grad(): with torch.no_grad():
i = 0 i = 0
for input, target in dataloader: for input, mask, target in dataloader:
input = input.to(device) input = input.to(device)
target = target.to(device) target = target.to(device)
mask = mask.to(device)
input = torch.cat((input, mask), 1)
out = model(input) out = model(input)
loss += loss_fn(out, target).item() loss += loss_fn(out, target).item()
print(f'\rEval prog[{i}/{len(dataloader) * dataloader.batch_size}]', end='') 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) os.makedirs(path, exist_ok=True)
fig, axes = plt.subplots(ncols=3, figsize=(15, 5)) 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"]): for ax, data, title in zip(axes, [inputs, targets, predictions], ["Input", "Target", "Prediction"]):
ax.clear() ax.clear()
ax.set_title(title) ax.set_title(title)
ax.imshow(DataLoader.postprocess(np.transpose(data[i], (1, 2, 0))), interpolation="none") 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.imshow(np.transpose((data[i]), (1, 2, 0)), interpolation="none")
ax.set_axis_off() ax.set_axis_off()
fig.savefig(os.path.join(path, f"{epoch:02d}_{update:07d}_{i:02d}.png"), dpi=100) fig.savefig(os.path.join(path, f"{epoch:02d}_{update:07d}.png"), dpi=100)
plt.close(fig) plt.close(fig)

6
Net.py
View File

@ -1,9 +1,13 @@
import math
import numpy as np import numpy as np
import torch import torch
from torch import nn
class ImageNN(torch.nn.Module): 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""" """Simple CNN with `n_hidden_layers`, `n_kernels`, and `kernel_size` as hyperparameters"""
super().__init__() super().__init__()

View File

@ -15,6 +15,9 @@ def save_model(model: torch.nn.Module):
print(f"Saved raw model to {MODEL_PATH}") print(f"Saved raw model to {MODEL_PATH}")
torch.save(model, 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(): def eval_evalset():
# read the provided testing pickle file # read the provided testing pickle file
@ -28,14 +31,19 @@ def eval_evalset():
model.eval() model.eval()
with open('testing/inputs.pkl', 'rb') as handle: with open('testing/inputs.pkl', 'rb') as handle:
b: dict = pickle.load(handle) b: dict = pickle.load(handle)
outarr = np.zeros(dtype=np.uint8, shape=(len(b['input_arrays']), 3, 100, 100)) outarr = ()
i = 0 i = 0
piclen = len(b['input_arrays']) piclen = len(b['input_arrays'])
for pic in b['input_arrays']: for input_array, known_array in zip(b['input_arrays'], b['known_arrays']):
pic = DataLoader.preprocess(pic, precision=np.float32) input_array = DataLoader.preprocess(input_array, precision=np.float32)
out = model(torch.from_numpy(pic)) 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()) 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='') print(f'\rApplying model [{i}/{piclen}]', end='')
i += 1 i += 1