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)