66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
import pickle
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
import Compress
|
|
import DataLoader
|
|
import ex4
|
|
from ImageImpaint import get_train_device
|
|
from netio import load_model, eval_evalset, write_to_pickle
|
|
|
|
|
|
def apply_model(filepath: str):
|
|
device = get_train_device()
|
|
|
|
img = Image.open(filepath)
|
|
model = load_model()
|
|
model.to(device)
|
|
|
|
pic = DataLoader.crop_image(img)
|
|
pic = DataLoader.preprocess(pic, precision=np.float32)
|
|
pic = ex4.ex4(pic, (5, 5), (4, 4))[0]
|
|
Image.fromarray((np.transpose(DataLoader.postprocess(pic), (1, 2, 0)).astype(np.uint8))).save("filename_grid.jpg")
|
|
out = model(torch.from_numpy(pic).to(device))
|
|
out = out.cpu().detach().numpy()
|
|
out = DataLoader.postprocess(out)
|
|
out = np.transpose(out, (1, 2, 0))
|
|
im = Image.fromarray(out)
|
|
im.save("filename.jpg", format="jpeg")
|
|
|
|
def test():
|
|
# read the provided testing pickle file
|
|
print("Generating pickle file with privided test data")
|
|
PICKEL_PATH = "test"
|
|
|
|
model = load_model()
|
|
model.eval()
|
|
|
|
loader,_ = DataLoader.get_image_loader("training/", np.float32)
|
|
outarr = np.zeros(dtype=np.uint8, shape=(8663, 3, 100, 100))
|
|
targetarr = np.zeros(dtype=np.uint8, shape=(8663, 3, 100, 100))
|
|
|
|
i = 0
|
|
for input, target in loader:
|
|
out = model(input)
|
|
out = DataLoader.postprocess(out.cpu().detach().numpy())
|
|
outarr[i] = out
|
|
targetarr[i] = DataLoader.postprocess(target.cpu().detach().numpy())
|
|
print(f'\rApplying model [{i}/{len(loader)}]', end='')
|
|
i += 1
|
|
if i==8663:
|
|
break
|
|
write_to_pickle(PICKEL_PATH + "_pred.pkl", list(outarr))
|
|
# compress the generated pickle arr
|
|
Compress.compress(PICKEL_PATH + "_pred.pkl")
|
|
|
|
write_to_pickle(PICKEL_PATH + "_target.pkl", list(targetarr))
|
|
# compress the generated pickle arr
|
|
Compress.compress(PICKEL_PATH + "_target.pkl")
|
|
|
|
if __name__ == '__main__':
|
|
# apply_model("training/000/000017.jpg")
|
|
eval_evalset()
|
|
# test()
|