lukas/ImageImpaint_Python_II
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

saving of model

Browse Source
This commit is contained in:
lukas-heiligenbrunner 2022-07-01 15:35:12 +02:00
parent 79734a8d75
commit 11640a6494
6 changed files with 99 additions and 77 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
ApplyModel.py Normal file
View File

@ -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")

44
DataLoader.py
View File

@ -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

40
ImageImpaint.py
View File

@ -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()

6
Net.py
View File

@ -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)

11
main.py
View File

@ -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()

34
netio.py
View File

@ -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 = []
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)
pic = DataLoader.preprocess(pic, precision=np.float32)
out = model(torch.from_numpy(pic))
out = DataLoader.postprocess(out.detach().numpy())
pickle.dump(out, writehandle, protocol=pickle.HIGHEST_PROTOCOL)
out = DataLoader.postprocess(out.cpu().detach().numpy())
outarr[i] = out
print(
f'\rApplying model [{i}/{piclen}] {sys.getsizeof(outarr)}',end='')
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)