Spaces:
Runtime error
Runtime error
| """Test script for anime-to-sketch translation | |
| Example: | |
| python3 test.py --dataroot /your_path/dir --load_size 512 | |
| python3 test.py --dataroot /your_path/img.jpg --load_size 512 | |
| """ | |
| import os | |
| from data import get_image_list | |
| from model import create_model | |
| from data import read_img_path, tensor_to_img, save_image | |
| import argparse | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser(description='Anime-to-sketch test options.') | |
| parser.add_argument('--dataroot','-i', default='test_samples/', type=str) | |
| parser.add_argument('--load_size','-s', default=512, type=int) | |
| parser.add_argument('--output_dir','-o', default='results/', type=str) | |
| parser.add_argument('--gpu_ids', '-g', default=[], help="gpu ids: e.g. 0 0,1,2 0,2.") | |
| opt = parser.parse_args() | |
| # create model | |
| model = create_model(opt.gpu_ids) # create a model given opt.model and other options | |
| model.eval() | |
| # get input data | |
| if os.path.isdir(opt.dataroot): | |
| test_list = get_image_list(opt.dataroot) | |
| elif os.path.isfile(opt.dataroot): | |
| test_list = [opt.dataroot] | |
| else: | |
| raise Exception("{} is not a valid directory or image file.".format(opt.dataroot)) | |
| # save outputs | |
| save_dir = opt.output_dir | |
| os.makedirs(save_dir, exist_ok=True) | |
| for test_path in test_list: | |
| basename = os.path.basename(test_path) | |
| aus_path = os.path.join(save_dir, basename) | |
| img, aus_resize = read_img_path(test_path, opt.load_size) | |
| aus_tensor = model(img) | |
| aus_img = tensor_to_img(aus_tensor) | |
| save_image(aus_img, aus_path, aus_resize) | |