import argparse
import os
import torch
from data import get_data, rgb2lab, lab2rgb
from model import get_model, model_load, model_device, valid_epoch, enable_amp
if __name__ == "__main__":
"""Test model."""
parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint',
type=str,
default="output/ImageColor_G.pth",
help="checkpoint file")
parser.add_argument('--bs', type=int, default=32, help="batch size")
args = parser.parse_args()
# get model
model = get_model(trainning=False)
model_load(model.net_G, args.checkpoint)
device = model_device()
model.net_G.to(device)
enable_amp(model.net_G)
print("Start testing ...")
test_dl = get_data(trainning=False, bs=args.bs)
valid_epoch(test_dl, model, device, tag='test')
type=str,
default="dataset/predict/input",
help="input folder")
parser.add_argument('--output',
type=str,
default="dataset/predict/output",
help="output folder")
args = parser.parse_args()
model = get_model()
model_load(model, args.checkpoint)
device = model_device()
model.to(device)
model.eval()
enable_amp(model)
totensor = transforms.ToTensor()
toimage = transforms.ToPILImage()
video = Video()
video.reset(args.input)
progress_bar = tqdm(total=len(video))
h = video.height
w = video.width
scale = 4
count = 1
for index in range(len(video)):
progress_bar.update(1)
from data import get_data
from model import enable_amp, get_model, model_device, model_load, valid_epoch
if __name__ == "__main__":
"""Test model."""
parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint',
type=str,
default="models/VideoColor.pth",
help="checkpoint file")
parser.add_argument('--bs', type=int, default=2, help="batch size")
args = parser.parse_args()
# get model
model_r = get_model("modelR")
model_load(model_r, "modelR", args.checkpoint)
device = model_device()
model_r.to(device)
model_c = get_model("modelC")
model_load(model_c, "modelC", args.checkpoint)
model_c.to(device)
enable_amp(model_r)
enable_amp(model_c)
print("Start testing ...")
test_dl = get_data(trainning=False, bs=args.bs)
valid_epoch(test_dl, model_r, device, tag='test')
