def sketch2anime(img, load_size=512):
img, aus_resize = read_img_path(img.name, load_size)
aus_tensor = model(img)
aus_img = tensor_to_img(aus_tensor)
image_pil = Image.fromarray(aus_img)
image_pil = image_pil.resize(aus_resize, Image.BICUBIC)
return image_pil
def test(dataroot='test_samples/',
load_size=512,
output_dir='results/',
gpu_ids=[]):
# create model
model = create_model(
gpu_ids) # create a model given model and other options
model.eval()
# get input data
if os.path.isdir(dataroot):
test_list = get_image_list(dataroot)
elif os.path.isfile(dataroot):
test_list = [dataroot]
else:
raise Exception(
"{} is not a valid directory or image file.".format(dataroot))
# save outputs
save_dir = 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, load_size)
aus_tensor = model(img)
aus_img = tensor_to_img(aus_tensor)
save_image(aus_img, aus_path, aus_resize)
def anime2sketch(img_input, load_size=512):
img, aus_resize = read_img_path(c, load_size)
model = load_model()
aus_tensor = model(img)
aus_img = tensor_to_img(aus_tensor)
image_pil = Image.fromarray(aus_img)
image_pil = image_pil.resize(aus_resize, Image.BICUBIC)
return image_pil
default=[],
help="gpu ids: e.g. 0 0,1,2 0,2.")
opt = parser.parse_args()
# create model
gpu_list = ','.join(str(x) for x in opt.gpu_ids)
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
device = torch.device('cuda' if len(opt.gpu_ids) > 0 else 'cpu')
model = create_model().to(
device) # 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.to(device))
aus_img = tensor_to_img(aus_tensor)
save_image(aus_img, aus_path, aus_resize)