55 lines
1.4 KiB
Python
55 lines
1.4 KiB
Python
import os
|
|
import pickle
|
|
|
|
import numpy as np
|
|
import torch
|
|
import Compress
|
|
import DataLoader
|
|
from Net import ImageNN
|
|
|
|
MODEL_PATH = 'impaintmodel.pt'
|
|
PICKEL_PATH = 'impaintmodel.pkl'
|
|
|
|
|
|
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:
|
|
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}]', end='')
|
|
i += 1
|
|
|
|
write_to_pickle(PICKEL_PATH, list(outarr))
|
|
# compress the generated pickle arr
|
|
Compress.compress(PICKEL_PATH)
|
|
|
|
|
|
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)
|