def evaluate(args):
content_image = utils.tensor_load_rgbimage(args.content_image,
size=args.content_size,
keep_asp=True)
content_image = content_image.unsqueeze(0)
style = utils.tensor_load_rgbimage(args.style_image, size=args.style_size)
style = style.unsqueeze(0)
style = utils.preprocess_batch(style)
style_model = Net(ngf=args.ngf)
style_model.load_state_dict(torch.load(args.model), False)
if args.cuda:
style_model.cuda()
content_image = content_image.cuda()
style = style.cuda()
style_v = Variable(style)
content_image = Variable(utils.preprocess_batch(content_image))
style_model.setTarget(style_v)
output = style_model(content_image)
#output = utils.color_match(output, style_v)
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def evaluate(args):
content_image = utils.tensor_load_rgbimage(args.content_image,
size=args.content_size,
keep_asp=True)
content_image = content_image.unsqueeze(0)
style = utils.tensor_load_rgbimage(args.style_image, size=args.style_size)
style = style.unsqueeze(0)
style = utils.preprocess_batch(style)
vgg = Vgg16()
utils.init_vgg16(args.vgg_model_dir)
vgg.load_state_dict(
torch.load(os.path.join(args.vgg_model_dir, "vgg16.weight")))
style_model = HangSNetV1()
style_model.load_state_dict(torch.load(args.model))
if args.cuda:
style_model.cuda()
vgg.cuda()
content_image = content_image.cuda()
style = style.cuda()
style_v = Variable(style, volatile=True)
utils.subtract_imagenet_mean_batch(style_v)
features_style = vgg(style_v)
gram_style = [utils.gram_matrix(y) for y in features_style]
content_image = Variable(utils.preprocess_batch(content_image))
target = Variable(gram_style[2].data, requires_grad=False)
style_model.setTarget(target)
output = style_model(content_image)
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def evaluate(my_content_image, my_content_size, my_style_image,my_style_size, my_ngf, my_cuda, my_output_image, my_model):
content_image = utils.tensor_load_rgbimage(my_content_image, size=my_content_size, keep_asp=True)
content_image = content_image.unsqueeze(0)
style = utils.tensor_load_rgbimage(my_style_image, size=my_style_size)
style = style.unsqueeze(0)
style = utils.preprocess_batch(style)
style_model = Net(ngf=my_ngf)
model_dict = torch.load(my_model)
model_dict_clone = model_dict.copy()
for key, value in model_dict_clone.items():
if key.endswith(('running_mean', 'running_var')):
del model_dict[key]
style_model.load_state_dict(model_dict, False)
if my_cuda:
style_model.cuda()
content_image = content_image.cuda()
style = style.cuda()
style_v = Variable(style)
content_image = Variable(utils.preprocess_batch(content_image))
style_model.setTarget(style_v)
output = style_model(content_image)
#output = utils.color_match(output, style_v)
utils.tensor_save_bgrimage(output.data[0], my_output_image, my_cuda)
return utils.tensor_return_bgrimage(output.data[0], my_cuda)
def evaluate():
if mx.context.num_gpus() > 0:
ctx = mx.gpu()
else:
ctx = mx.cpu(0)
# loading configs
args = Options().parse()
cfg = Configs(args.config_path)
# set logging level
logging.basicConfig(level=logging.INFO)
# images
content_image = tensor_load_rgbimage(cfg.content_image,
ctx,
size=cfg.val_img_size,
keep_asp=True)
style_image = tensor_load_rgbimage(cfg.style_image,
ctx,
size=cfg.val_style_size)
style_image = preprocess_batch(style_image)
# model
style_model = Net(ngf=cfg.ngf)
style_model.collect_params().load(cfg.val_model, ctx=ctx)
# forward
output = style_model(content_image, style_image)
# save img
tensor_save_bgrimage(output[0], cfg.output_img)
logging.info("Save img to {}".format(cfg.output_img))
def stylize(args):
#content_image = utils.tensor_load_rgbimage(args.content_image, scale = args.content_scale)
#content_image = content_image.unsqueeze(0)
content_image = None
if args.srcnn:
content_image = utils.tensor_load_rgbimage(args.content_image,
scale=args.upsample)
else:
content_image = utils.tensor_load_rgbimage(args.content_image)
content_image.unsqueeze_(0)
if args.cuda:
content_image = content_image.cuda()
content_image = Variable(utils.preprocess_batch(content_image),
volatile=True)
style_model = None
if args.srcnn:
style_model = SRCNN()
else:
style_model = TransformerNet(args.arch)
##style_model = TransformerNet()
style_model.load_state_dict(torch.load(args.model))
if args.cuda:
style_model.cuda()
output = style_model(content_image)
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def do_style(msgnet, style, content, image):
with torch.no_grad():
if type(content) is str:
content = utils.tensor_load_rgbimage(content, 256, keep_asp=True)
msgnet.setTarget(utils.tensor_load_rgbimage(style, 256).unsqueeze(0))
im = msgnet(content.detach().unsqueeze(0))[0]
utils.tensor_save_rgbimage(im, image)
def predict(args):
style_model = Net(ngf=args.ngf)
model_dict = torch.load(args.model)
model_dict_clone = model_dict.copy()
for key, value in model_dict_clone.items():
if key.endswith(('running_mean', 'running_var')):
del model_dict[key]
style_model.load_state_dict(model_dict, False)
style_model.eval()
if args.cuda:
style_model.cuda()
img = utils.tensor_load_rgbimage(args.image, (args.w, args.h))
img = img.unsqueeze(0).float()
if args.cuda:
img = img.cuda()
img = Variable(img)
style = utils.tensor_load_rgbimage(args.style_image, args.style_size)
style = style.unsqueeze(0)
style = utils.preprocess_batch(style)
if args.cuda:
style = style.cuda()
style = Variable(style, requires_grad=False)
style_model.set_target(style)
img = style_model(img)
utils.tensor_save_rgbimage(img, args.out_image, args.cuda)
def optimize(args):
""" Gatys et al. CVPR 2017
ref: Image Style Transfer Using Convolutional Neural Networks
"""
# load the content and style target
content_image = utils.tensor_load_rgbimage(args.content_image,
size=args.content_size,
keep_asp=True)
content_image = content_image.unsqueeze(0)
content_image = Variable(utils.preprocess_batch(content_image),
requires_grad=False)
content_image = utils.subtract_imagenet_mean_batch(content_image)
style_image = utils.tensor_load_rgbimage(args.style_image,
size=args.style_size)
style_image = style_image.unsqueeze(0)
style_image = Variable(utils.preprocess_batch(style_image),
requires_grad=False)
style_image = utils.subtract_imagenet_mean_batch(style_image)
# load the pre-trained vgg-16 and extract features
vgg = Vgg16()
utils.init_vgg16(args.vgg_model_dir)
vgg.load_state_dict(
torch.load(os.path.join(args.vgg_model_dir, "vgg16.weight")))
if args.cuda:
content_image = content_image.cuda()
style_image = style_image.cuda()
vgg.cuda()
features_content = vgg(content_image)
f_xc_c = Variable(features_content[1].data, requires_grad=False)
features_style = vgg(style_image)
gram_style = [utils.gram_matrix(y) for y in features_style]
# init optimizer
output = Variable(content_image.data, requires_grad=True)
optimizer = Adam([output], lr=args.lr)
mse_loss = torch.nn.MSELoss()
# optimizing the images
for e in range(args.iters):
utils.imagenet_clamp_batch(output, 0, 255)
optimizer.zero_grad()
features_y = vgg(output)
content_loss = args.content_weight * mse_loss(features_y[1], f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_y = utils.gram_matrix(features_y[m])
gram_s = Variable(gram_style[m].data, requires_grad=False)
style_loss += args.style_weight * mse_loss(gram_y, gram_s)
total_loss = content_loss + style_loss
if (e + 1) % args.log_interval == 0:
print(total_loss.data.cpu().numpy()[0])
total_loss.backward()
optimizer.step()
# save the image
output = utils.add_imagenet_mean_batch(output)
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def optimize(args):
""" Gatys et al. CVPR 2017
ref: Image Style Transfer Using Convolutional Neural Networks
"""
if args.cuda:
ctx = mx.gpu(0)
else:
ctx = mx.cpu(0)
# load the content and style target
content_image = utils.tensor_load_rgbimage(args.content_image,
ctx,
size=args.content_size,
keep_asp=True)
content_image = utils.subtract_imagenet_mean_preprocess_batch(
content_image)
style_image = utils.tensor_load_rgbimage(args.style_image,
ctx,
size=args.style_size)
style_image = utils.subtract_imagenet_mean_preprocess_batch(style_image)
# load the pre-trained vgg-16 and extract features
vgg = net.Vgg16()
utils.init_vgg_params(vgg, 'models', ctx=ctx)
# content feature
f_xc_c = vgg(content_image)[1]
# style feature
features_style = vgg(style_image)
gram_style = [net.gram_matrix(y) for y in features_style]
# output
output = Parameter('output', shape=content_image.shape)
output.initialize(ctx=ctx)
output.set_data(content_image)
# optimizer
trainer = gluon.Trainer([output], 'adam', {'learning_rate': args.lr})
mse_loss = gluon.loss.L2Loss()
# optimizing the images
for e in range(args.iters):
utils.imagenet_clamp_batch(output.data(), 0, 255)
# fix BN for pre-trained vgg
with autograd.record():
features_y = vgg(output.data())
content_loss = 2 * args.content_weight * mse_loss(
features_y[1], f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_y = net.gram_matrix(features_y[m])
gram_s = gram_style[m]
style_loss = style_loss + 2 * args.style_weight * mse_loss(
gram_y, gram_s)
total_loss = content_loss + style_loss
total_loss.backward()
trainer.step(1)
if (e + 1) % args.log_interval == 0:
print('loss:{:.2f}'.format(total_loss.asnumpy()[0]))
# save the image
output = utils.add_imagenet_mean_batch(output.data())
utils.tensor_save_bgrimage(output[0], args.output_image, args.cuda)
def optimize(args):
""" Gatys et al. CVPR 2017
ref: Image Style Transfer Using Convolutional Neural Networks
"""
if args.cuda:
ctx = mx.gpu(0)
else:
ctx = mx.cpu(0)
# load the content and style target
content_image = utils.tensor_load_rgbimage(args.content_image,ctx, size=args.content_size, keep_asp=True)
content_image = utils.subtract_imagenet_mean_preprocess_batch(content_image)
style_image = utils.tensor_load_rgbimage(args.style_image, ctx, size=args.style_size)
style_image = utils.subtract_imagenet_mean_preprocess_batch(style_image)
# load the pre-trained vgg-16 and extract features
vgg = net.Vgg16()
utils.init_vgg_params(vgg, 'models', ctx=ctx)
# content feature
f_xc_c = vgg(content_image)[1]
# style feature
features_style = vgg(style_image)
gram_style = [net.gram_matrix(y) for y in features_style]
# output
output = Parameter('output', shape=content_image.shape)
output.initialize(ctx=ctx)
output.set_data(content_image)
# optimizer
trainer = gluon.Trainer([output], 'adam',
{'learning_rate': args.lr})
mse_loss = gluon.loss.L2Loss()
# optimizing the images
for e in range(args.iters):
utils.imagenet_clamp_batch(output.data(), 0, 255)
# fix BN for pre-trained vgg
with autograd.record():
features_y = vgg(output.data())
content_loss = 2 * args.content_weight * mse_loss(features_y[1], f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_y = net.gram_matrix(features_y[m])
gram_s = gram_style[m]
style_loss = style_loss + 2 * args.style_weight * mse_loss(gram_y, gram_s)
total_loss = content_loss + style_loss
total_loss.backward()
trainer.step(1)
if (e + 1) % args.log_interval == 0:
print('loss:{:.2f}'.format(total_loss.asnumpy()[0]))
# save the image
output = utils.add_imagenet_mean_batch(output.data())
utils.tensor_save_bgrimage(output[0], args.output_image, args.cuda)
def fast_evaluate(args, basedir, contents, idx=0):
# basedir to save the data
style_model = Net(ngf=args.ngf)
style_model.load_state_dict(torch.load(args.model), False)
style_model.eval()
if args.cuda:
style_model.cuda()
style_loader = StyleLoader(args.style_folder,
args.style_size,
cuda=args.cuda)
for content_image in contents:
idx += 1
content_image = utils.tensor_load_rgbimage(content_image,
size=args.content_size,
keep_asp=True).unsqueeze(0)
if args.cuda:
content_image = content_image.cuda()
content_image = Variable(utils.preprocess_batch(content_image))
for isx in range(style_loader.size()):
style_v = Variable(style_loader.get(isx).data)
style_model.setTarget(style_v)
output = style_model(content_image)
filename = os.path.join(basedir, "{}_{}.png".format(idx, isx + 1))
utils.tensor_save_bgrimage(output.data[0], filename, args.cuda)
print(filename)
def main():
args = Options().parse()
style_model = Net(ngf=args.ngf)
model_dict = torch.load(args.model)
model_dict_clone = model_dict.copy()
for key, value in model_dict_clone.items():
if key.endswith(('running_mean', 'running_var')):
del model_dict[key]
style_model.load_state_dict(model_dict, False)
style_loaders = utils.StyleLoader(args.style_folder, args.style_size)
content_image = utils.tensor_load_rgbimage(args.content_image,
size=args.style_size,
keep_asp=True)
content_image = content_image.unsqueeze(0)
if args.cuda:
style_model.cuda()
content_image = content_image.cuda()
content_image = Variable(utils.preprocess_batch(content_image))
# for i, style_loader in enumerate(style_loaders):
for i in range(style_loaders.size()):
print(i)
style_v = style_loaders.get(i)
style_model.setTarget(style_v)
output = style_model(content_image)
filepath = "out/output" + str(i + 1) + '.jpg'
print(filepath)
utils.tensor_save_bgrimage(output.data[0], filepath, args.cuda)
def stylize_onnx_caffe2(args):
"""
Read ONNX model and run it using Caffe2
"""
assert not args.export_onnx
# TODO: change tools to run without PyTorch
content_image = utils.tensor_load_rgbimage(args.content_image,
scale=args.content_scale)
content_image = content_image.unsqueeze(0)
content_image = utils.preprocess_batch(content_image).numpy()
import onnx
import onnx_caffe2.backend
model = onnx.load(args.model)
prepared_backend = onnx_caffe2.backend.prepare(
model, device='CUDA' if args.cuda else 'CPU')
if args.half:
for op in prepared_backend.predict_net.op:
op.engine = 'CUDNN'
content_image = content_image.astype(np.float16)
inp = {model.graph.input[0].name: content_image}
c2_out = prepared_backend.run(inp)[0]
if args.half:
c2_out = c2_out.astype(np.float32)
output = torch.from_numpy(c2_out)
utils.tensor_save_rgbimage(output[0], args.output_image, args.cuda)
def evaluate(args):
if args.cuda:
ctx = mx.gpu(0)
else:
ctx = mx.cpu(0)
# images
content_image = utils.tensor_load_rgbimage(args.content_image,ctx, size=args.content_size, keep_asp=True)
style_image = utils.tensor_load_rgbimage(args.style_image, ctx, size=args.style_size)
style_image = utils.preprocess_batch(style_image)
# model
style_model = net.Net(ngf=args.ngf)
style_model.collect_params().load(args.model, ctx=ctx)
# forward
style_model.setTarget(style_image)
output = style_model(content_image)
utils.tensor_save_bgrimage(output[0], args.output_image, args.cuda)
def stylize(args):
img = None
content_image = utils.tensor_load_rgbimage(args.content_image,
scale=args.content_scale)
content_image = content_image.unsqueeze(0)
style_model = TransformerNet()
style_model.load_state_dict(torch.load(args.model))
cam = cv2.VideoCapture(0)
for x in range(0, 150):
ret_val, img13 = cam.read()
content_image = utils.tensor_load_rgbimage_cam(
img13, scale=args.content_scale)
content_image = content_image.unsqueeze(0)
if args.cuda:
content_image = content_image.cuda()
content_image2 = Variable(utils.preprocess_batch(content_image),
volatile=True)
if args.cuda:
style_model.cuda()
output = style_model(content_image2)
im = utils.tensor_ret_bgrimage(output.data[0], args.output_image,
args.cuda)
if img is None:
img = pl.imshow(im)
else:
img.set_data(im)
pl.pause(.1)
pl.draw()
def stylize(args):
content_image = utils.tensor_load_rgbimage(args.content_image,
scale=args.content_scale)
content_image = content_image.unsqueeze(0)
content_image = Variable(utils.preprocess_batch(content_image))
style_model = torch.load(args.model)
output = style_model(content_image)
utils.tensor_save_bgrimage(output.data[0], args.output_image)
def get(self, i): idx = i%len(self.files) filepath = os.path.join(self.folder, self.files[idx]) style = utils.tensor_load_rgbimage(filepath, self.style_size) style = style.unsqueeze(0) style = utils.preprocess_batch(style) if self.cuda: style = style.cuda() style_v = Variable(style, requires_grad=False) return style_v
def evaluate(raw_content_image, raw_content_size, style_image, style_size,
cuda, output_name):
content_image = utils.tensor_load_rgbimage(raw_content_image,
size=raw_content_size,
keep_asp=True)
content_image = content_image.unsqueeze(0)
style = utils.tensor_load_rgbimage(style_image, size=style_size)
style = style.unsqueeze(0)
style = utils.preprocess_batch(style)
style_v = Variable(style)
content_image = Variable(utils.preprocess_batch(content_image))
style_model.setTarget(style_v)
output = style_model(content_image)
transfer_image = utils.tensor_save_bgrimage(output.data[0], output_name,
cuda)
return transfer_image
def evaluate(model_dir, c_img, s_img, img_size, out_img):
content_image = utils.tensor_load_rgbimage(c_img, size=img_size, keep_asp=True)
content_image = content_image.unsqueeze(0).to(device)
style = utils.tensor_load_rgbimage(s_img, size=img_size)
style = style.unsqueeze(0)
style = utils.preprocess_batch(style).to(device)
style_model = Net(ngf=FILTER_CHANNEL, dv=device).to(device)
style_model.load_state_dict(torch.load(model_dir), False)
style_v = Variable(style)
content_image = Variable(utils.preprocess_batch(content_image))
style_model.setTarget(style_v)
output = style_model(content_image)
#output = utils.color_match(output, style_v)
utils.tensor_save_bgrimage(output.data[0], out_img)
print ('Done')
def stylize(args):
if args.model.endswith(".onnx"):
return stylize_onnx_caffe2(args)
content_image = utils.tensor_load_rgbimage(args.content_image,
scale=args.content_scale)
content_image = content_image.unsqueeze(0)
if args.cuda:
content_image = content_image.cuda()
content_image = Variable(utils.preprocess_batch(content_image),
requires_grad=False)
style_model = TransformerNet()
state_dict = torch.load(args.model)
# removed_modules = ['in2']
in_names = [
"in1.scale", "in1.shift", "in2.scale", "in2.shift", "in3.scale",
"in3.shift", "res1.in1.scale", "res1.in1.shift", "res1.in2.scale",
"res1.in2.shift", "res2.in1.scale", "res2.in1.shift", "res2.in2.scale",
"res2.in2.shift", "res3.in1.scale", "res3.in1.shift", "res3.in2.scale",
"res3.in2.shift", "res4.in1.scale", "res4.in1.shift", "res4.in2.scale",
"res4.in2.shift", "res5.in1.scale", "res5.in1.shift", "res5.in2.scale",
"res5.in2.shift", "in4.scale", "in4.shift", "in5.scale", "in5.shift"
]
# kl = list(state_dict.keys())
# for k in kl:
for k in in_names:
state_dict[k.replace("scale",
"weight").replace("shift",
"bias")] = state_dict.pop(k)
style_model.load_state_dict(state_dict)
if args.cuda:
style_model.cuda()
if args.half:
style_model.half()
content_image = content_image.half()
if args.export_onnx:
assert args.export_onnx.endswith(
".onnx"), "Export model file should end with .onnx"
output = torch.onnx._export(style_model, content_image,
args.export_onnx)
else:
output = style_model(content_image)
if args.half:
output = output.float()
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def run_demo(eval_args):
## load parameters
#content_image = 'images/content/xjtlu.jpg'
#style_image = 'images/styles/starry_night.jpg'
#eval_args = program_args(content_image,content_image,style_image,128,128,0)
#eval_args = camero_args(style_image)
if eval_args.cuda == 0:
ctx = mx.cpu()
else:
ctx = mx.gpu()
## Change the content and style image using Style Loader
#content_image = utils.tensor_load_rgbimage(eval_args.contentImage, ctx, size=eval_args.size, keep_asp=True)
style_image = utils.tensor_load_rgbimage(eval_args.styleImage,
ctx,
size=eval_args.size)
style_image = utils.preprocess_batch(style_image)
style_model = net.Net(ngf=eval_args.ngf)
style_model.load_parameters(eval_args.model, ctx=ctx)
style_model.set_target(style_image)
cam = cv2.VideoCapture(0)
while True:
## read frame
ret, frame = cam.read()
# read content image (cimg)
#cimg = img.copy()
#img = np.array(img).transpose(2, 0, 1)
content_img = load_image(frame, ctx, eval_args.size)
output = style_model(content_img)
tensor = output[0]
#(b, g, r) = F.split(tensor, num_outputs=3, axis=0)
#tensor = F.concat(r, g, b, dim=0)
img = F.clip(tensor, 0, 255).asnumpy()
img = img.transpose(1, 2, 0).astype('uint8')
img = Image.fromarray(img)
image = np.array(
img.resize((frame.shape[1], frame.shape[0]), Image.ANTIALIAS))
#print(frame.shape,image.shape)
numpy_horizontal = np.hstack((frame, image))
#cv2.imshow("Content Window",frame)
#cv2.imshow("Style Window",grey)
cv2.imshow("Test Window Shape", numpy_horizontal)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cam.release()
cv2.destroyAllWindows()
def evaluate(args):
# set output_image
dirname = os.path.dirname(args.content_image)
style_ = os.path.basename(args.style_image).split('.')[0]
basename = style_ + '_' + os.path.basename(args.content_image)
args.output_image = os.path.join(dirname, basename)
content_image = utils.tensor_load_rgbimage(args.content_image,
size=args.content_size,
keep_asp=True)
content_image = content_image.unsqueeze(0)
style = utils.tensor_load_rgbimage(args.style_image, size=args.style_size)
style = style.unsqueeze(0)
style = utils.preprocess_batch(style)
model_dict = torch.load(args.model)
model_dict_clone = model_dict.copy() # We can't mutate while iterating
for key, value in model_dict_clone.items():
if key.endswith(('running_mean', 'running_var')):
del model_dict[key]
style_model = Net(ngf=args.ngf)
style_model.load_state_dict(model_dict, False)
# style_model = Net(ngf=args.ngf)
# style_model.load_state_dict(torch.load(args.model), False)
if args.cuda:
style_model.cuda()
content_image = content_image.cuda()
style = style.cuda()
style_v = Variable(style)
content_image = Variable(utils.preprocess_batch(content_image))
style_model.setTarget(style_v)
output = style_model(content_image)
#output = utils.color_match(output, style_v)
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def do_gan(generator, content, image):
with torch.no_grad():
if type(content) is str:
content = utils.tensor_load_rgbimage(content,
256,
keep_asp=False,
need_normalize=False)
content = gan_normalize(content)
content = content.unsqueeze(0)
content = cat_embeddings(content, zero_embeddings(content))
im = split_only_batch(generator(content))[0]
im = gan_unnormalize(im)
utils.tensor_save_rgbimage(im, image, need_unnormalize=False)
def evaluate_test():
content_image = 'images/content/xjtlu.jpg'
style_image = 'images/styles/starry_night.jpg'
output_image = 'test_output.jpg'
cuda = 0
eval_args = program_args(content_image,content_image,style_image,128,128,0)
if cuda == 0:
ctx = mx.cpu(0)
else:
ctx = mx.gpu(0)
content_image = utils.tensor_load_rgbimage(eval_args.contentImage, ctx, size=eval_args.size, keep_asp=True)
style_image = utils.tensor_load_rgbimage(eval_args.styleImage, ctx, size=eval_args.size)
style_image = utils.preprocess_batch(style_image)
# model
style_model = net.Net(ngf=eval_args.ngf)
style_model.load_parameters(eval_args.model, ctx=ctx)
# forward
style_model.set_target(style_image)
output = style_model(content_image)
tensor = output[0]
#(b, g, r) = F.split(tensor, num_outputs=3, axis=0)
#tensor = F.concat(r, g, b, dim=0)
img = F.clip(tensor, 0, 255).asnumpy()
img = img.transpose(1, 2, 0).astype('uint8')
img = Image.fromarray(img)
#print(type(output[0]),output[0].asnumpy().shape)
#np_array = output.asnumpy()
#img_array = np.reshape(np_array,(np_array.shape[2],np_array.shape[3],np_array.shape[1]))
#print(img_array.shape)
#img = Image.fromarray(np.uint8(img_array))
original = Image.open(content_image_path)
style = Image.open(style_image_path)
style.show()
original.show()
img.show()
def evaluate(args):
content_image = utils.tensor_load_rgbimage(args.content_image, size=args.content_size, keep_asp=True)
content_image = content_image.unsqueeze(0)
style = utils.tensor_load_rgbimage(args.style_image, size=args.style_size)
style = style.unsqueeze(0)
style = utils.preprocess_batch(style)
style_model = Net(ngf=args.ngf) # comment out for PyTorch 0.4
style_model.load_state_dict(torch.load(args.model), False) # comment out for PyTorch 0.4
# https://github.com/zhanghang1989/PyTorch-Multi-Style-Transfer/issues/21
# Compatibility shim for PyTorch 0.4
# model_dict = torch.load('models/icons2.model') # uncomment for PyTorch 0.4
# model_dict_clone = model_dict.copy() # uncomment for PyTorch 0.4
# for key, value in model_dict_clone.items(): # uncomment for PyTorch 0.4
# if key.endswith(('running_mean', 'running_var')): # uncomment for PyTorch 0.4
# del model_dict[key] # uncomment for PyTorch 0.4
### Next cell
#style_model = Net(ngf=128)
#style_model.load_state_dict(model_dict, False)
if args.cuda:
style_model.cuda()
content_image = content_image.cuda()
style = style.cuda()
style_v = Variable(style)
content_image = Variable(utils.preprocess_batch(content_image))
style_model.setTarget(style_v)
output = style_model(content_image)
#output = utils.color_match(output, style_v)
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def evaluate(args):
content_image = utils.tensor_load_rgbimage(args.content_image,
size=args.content_size,
keep_asp=True)
content_image = content_image.unsqueeze(0)
style = utils.tensor_load_rgbimage(args.style_image, size=args.style_size)
style = style.unsqueeze(0)
style = utils.preprocess_batch(style)
# style_model = Net(ngf=args.ngf)
model_dict = torch.load(
args.model
) # or args.resume, matching what's in the line with style_model.load_state_dict
model_dict_clone = model_dict.copy() # We can't mutate while iterating
for key, value in model_dict_clone.items():
if key.endswith(('running_mean', 'running_var')):
del model_dict[key]
style_model = Net(ngf=128) # to run with torch-0.3.0.post4
# style_model.load_state_dict(torch.load(args.model), False)
style_model.load_state_dict(model_dict, False)
if args.cuda:
style_model.cuda()
content_image = content_image.cuda()
style = style.cuda()
style_v = Variable(style)
content_image = Variable(utils.preprocess_batch(content_image))
style_model.setTarget(style_v)
output = style_model(content_image)
#output = utils.color_match(output, style_v)
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def evaluate(args):
if args.cuda:
ctx = mx.gpu(0)
else:
ctx = mx.cpu(0)
# images
content_image = utils.tensor_load_rgbimage(args.content_image,
ctx,
size=args.content_size,
keep_asp=True)
style_image = utils.tensor_load_rgbimage(args.style_image,
ctx,
size=args.style_size)
style_image = utils.preprocess_batch(style_image)
# model
WIDTH = content_image.shape[2]
HEIGHT = content_image.shape[3]
style_model = net.Net(ngf=args.ngf, width=WIDTH, height=HEIGHT)
style_model.collect_params().load(args.model, ctx=ctx)
# forward
style_model.setTarget(style_image)
output = style_model(content_image)
utils.tensor_save_bgrimage(output[0], args.output_image, args.cuda)
#Added to save and visualise model
x = mx.sym.var('data')
a = mx.sym.var('width')
b = mx.sym.var('height')
y = style_model(x)
#y_json = y.tojson()
y.save("MODEL.json")
y_json = y.tojson()
#print( "Network\n%s" % y_json)
style_model.save_params("MODEL.params")
graph = mx.viz.plot_network(y, save_format='pdf')
graph.render()
mx.visualization.print_summary(y, {'data': (1, 3, WIDTH, HEIGHT)})
def stylize(args):
content_image = utils.tensor_load_rgbimage(args.content_image,
scale=args.content_scale)
content_image = content_image.unsqueeze(0)
if args.cuda:
content_image = content_image.cuda()
content_image = Variable(utils.preprocess_batch(content_image))
style_model = TransformerNet()
style_model.load_state_dict(torch.load(args.model))
if args.cuda:
style_model.cuda()
output = style_model(content_image)
utils.tensor_save_bgrimage(output.data[0], args.output_image, args.cuda)
def stylize(content_image, model, content_scale=None, cuda=0):
content_image = utils.tensor_load_rgbimage(content_image,
scale=content_scale)
content_image = content_image.unsqueeze(0)
if cuda:
content_image = content_image.cuda()
content_image = Variable(utils.preprocess_batch(content_image),
volatile=True)
style_model = TransformerNet()
style_model.load_state_dict(torch.load(model))
if cuda:
style_model.cuda()
output = style_model(content_image)
return utils.tensor_to_Image(output, cuda)
def fast_evaluate(args, basedir, contents, idx=0):
# basedir to save the data
# style_model = Net(ngf=args.ngf)
model_dict = torch.load(
args.model
) # or args.resume, matching what's in the line with style_model.load_state_dict
model_dict_clone = model_dict.copy() # We can't mutate while iterating
for key, value in model_dict_clone.items():
if key.endswith(('running_mean', 'running_var')):
del model_dict[key]
style_model = Net(ngf=128) # to run with torch-0.3.0.post4
# style_model.load_state_dict(torch.load(args.model), False)
style_model.load_state_dict(model_dict, False)
style_model.eval()
if args.cuda:
style_model.cuda()
style_loader = StyleLoader(args.style_folder,
args.style_size,
cuda=args.cuda)
for content_image in contents:
idx += 1
content_image = utils.tensor_load_rgbimage(content_image,
size=args.content_size,
keep_asp=True).unsqueeze(0)
if args.cuda:
content_image = content_image.cuda()
content_image = Variable(utils.preprocess_batch(content_image))
for isx in range(style_loader.size()):
style_v = Variable(style_loader.get(isx).data)
style_model.setTarget(style_v)
output = style_model(content_image)
filename = os.path.join(basedir, "{}_{}.png".format(idx, isx + 1))
utils.tensor_save_bgrimage(output.data[0], filename, args.cuda)
print(filename)
def train(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
kwargs = {'num_workers': 0, 'pin_memory': False}
else:
kwargs = {}
transform = transforms.Compose([
transforms.Scale(args.image_size),
transforms.CenterCrop(args.image_size),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))
])
train_dataset = datasets.ImageFolder(args.dataset, transform)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
**kwargs)
transformer = TransformerNet()
optimizer = Adam(transformer.parameters(), args.lr)
mse_loss = torch.nn.MSELoss()
vgg = Vgg16()
utils.init_vgg16(args.vgg_model_dir)
vgg.load_state_dict(
torch.load(os.path.join(args.vgg_model_dir, "vgg16.weight")))
if args.cuda:
transformer.cuda()
vgg.cuda()
style = utils.tensor_load_rgbimage(args.style_image, size=args.style_size)
style = style.repeat(args.batch_size, 1, 1, 1)
style = utils.preprocess_batch(style)
if args.cuda:
style = style.cuda()
style_v = Variable(style, volatile=True)
utils.subtract_imagenet_mean_batch(style_v)
features_style = vgg(style_v)
gram_style = [utils.gram_matrix(y) for y in features_style]
for e in range(args.epochs):
transformer.train()
agg_content_loss = 0.
agg_style_loss = 0.
count = 0
for batch_id, (x, _) in enumerate(train_loader):
n_batch = len(x)
count += n_batch
optimizer.zero_grad()
x = Variable(utils.preprocess_batch(x))
if args.cuda:
x = x.cuda()
y = transformer(x)
xc = Variable(x.data.clone(), volatile=True)
utils.subtract_imagenet_mean_batch(y)
utils.subtract_imagenet_mean_batch(xc)
features_y = vgg(y)
features_xc = vgg(xc)
f_xc_c = Variable(features_xc[1].data, requires_grad=False)
content_loss = args.content_weight * mse_loss(
features_y[1], f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_s = Variable(gram_style[m].data, requires_grad=False)
gram_y = utils.gram_matrix(features_y[m])
style_loss += args.style_weight * mse_loss(
gram_y, gram_s[:n_batch, :, :])
total_loss = content_loss + style_loss
total_loss.backward()
optimizer.step()
agg_content_loss += content_loss.data[0]
agg_style_loss += style_loss.data[0]
if (batch_id + 1) % args.log_interval == 0:
mesg = "{}\tEpoch {}:\t[{}/{}]\tcontent: {:.6f}\tstyle: {:.6f}\ttotal: {:.6f}".format(
time.ctime(), e + 1, count, len(train_dataset),
agg_content_loss / (batch_id + 1),
agg_style_loss / (batch_id + 1),
(agg_content_loss + agg_style_loss) / (batch_id + 1))
print(mesg)
# save model
transformer.eval()
transformer.cpu()
save_model_filename = "epoch_" + str(args.epochs) + "_" + str(
time.ctime()).replace(' ', '_') + "_" + str(
args.content_weight) + "_" + str(args.style_weight) + ".model"
save_model_path = os.path.join(args.save_model_dir, save_model_filename)
torch.save(transformer.state_dict(), save_model_path)
print("\nDone, trained model saved at", save_model_path)
