Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import os
- import random
- import numpy as np
- import torch
- import torch.nn as nn
- import torch.nn.functional as F
- from torch.optim import Adam
- from torchvision.utils import save_image
- from tensorboardX import SummaryWriter
- from dataloader import get_loader
- from models.image_encoder import ImageEncoder
- from models.image_decoder import ImageDecoder
- from models.modality_fusion import ModalityFusion
- from models.vgg_perceptual_loss import VGGPerceptualLoss
- from models.svg_decoder import SVGLSTMDecoder, SVGMDNTop
- from models.svg_encoder import SVGLSTMEncoder
- from models import util_funcs
- from options import get_parser_main_model
- from data_utils.svg_utils import render
- from models.imgsr.modules import TrainOptions, create_model
- opts = get_parser_main_model().parse_args()
- opts.experiment_name = opts.experiment_name + '_' + opts.model_name
- os.makedirs("experiments", exist_ok=True)
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
- def get_z(temperature=1.0):
- z = torch.randn(1, 128).to(device) * temperature
- return z
- def get_img_feat(temperature=0.5):
- z = torch.randn(1, 1024).to(device) * temperature
- return z
- def get_enc_z(temperature=0.5):
- z = torch.randn(4, 1024).to(device) * temperature
- return z
- svg_dir = "output"
- exp_dir = os.path.join("experiments", opts.experiment_name)
- ckpt_dir = os.path.join(exp_dir, "checkpoints")
- res_dir = os.path.join(exp_dir, "results")
- svg_decoder = SVGLSTMDecoder(char_categories=opts.char_categories,
- bottleneck_bits=opts.bottleneck_bits, mode=opts.mode, max_sequence_length=opts.max_seq_len,
- hidden_size=opts.hidden_size,
- num_hidden_layers=opts.num_hidden_layers,
- feature_dim=opts.seq_feature_dim, ff_dropout=opts.ff_dropout, rec_dropout=opts.rec_dropout)
- mdn_top_layer = SVGMDNTop(num_mixture=opts.num_mixture, seq_len=opts.max_seq_len, hidden_size=opts.hidden_size,
- mode=opts.mode, mix_temperature=opts.mix_temperature,
- gauss_temperature=opts.gauss_temperature, dont_reduce=opts.dont_reduce_loss)
- modality_fusion = ModalityFusion(img_feat_dim = 16 * opts.image_size, hidden_size = opts.hidden_size, ref_nshot = opts.ref_nshot, bottleneck_bits = opts.bottleneck_bits, mode=opts.mode)
- # load parameters
- epoch = opts.test_epoch
- svg_decoder_fpath = os.path.join(ckpt_dir, f"{opts.model_name}_{epoch}.seqdec.pth")
- svg_decoder.load_state_dict(torch.load(svg_decoder_fpath, map_location=device))
- svg_decoder.eval()
- mdn_top_layer_fpath = os.path.join(ckpt_dir, f"{opts.model_name}_{epoch}.mdntl.pth")
- mdn_top_layer.load_state_dict(torch.load(mdn_top_layer_fpath, map_location=device))
- mdn_top_layer.eval()
- modality_fusion_fpath = os.path.join(ckpt_dir, f"{opts.model_name}_{epoch}.modalfuse.pth")
- modality_fusion.load_state_dict(torch.load(modality_fusion_fpath, map_location=device))
- modality_fusion.eval()
- # to device
- modality_fusion = modality_fusion.to(device)
- svg_decoder = svg_decoder.to(device)
- mdn_top_layer = mdn_top_layer.to(device)
- val_img_l1_loss = 0.0
- val_img_pt_loss = 0.0
- mean = np.load('./data/mean.npz')
- std = np.load('./data/stdev.npz')
- mean = torch.from_numpy(mean).to(device).to(torch.float32)
- std = torch.from_numpy(std).to(device).to(torch.float32)
- latent_feat = get_z()
- latent_feat = latent_feat.repeat(opts.char_categories, 1)
- trg_char = F.one_hot(torch.arange(52), num_classes=opts.char_categories)
- sd_init_state = svg_decoder.init_state_input(latent_feat, trg_char)
- hidden, cell = sd_init_state['hidden'], sd_init_state['cell']
- hidden_self, cell_self = hidden, cell
- # outputs_self = torch.zeros(trg_seq.size(0), trg_seq.size(1), opts.hidden_size).to(device)
- tgt_len = 51
- trg_seq_size_1 = 52
- sampled_svg = torch.zeros(tgt_len, trg_seq_size_1, opts.seq_feature_dim).to(device)
- for t in range(0, tgt_len):
- # self sample results
- if t == 0:
- inpt_self = torch.zeros(trg_seq_size_1, opts.seq_feature_dim).to(device)
- else:
- inpt_self = sampled_svg[t - 1]
- decoder_output_self = svg_decoder(inpt_self, hidden_self, cell_self)
- output_self, hidden_self, cell_self = decoder_output_self['output'], decoder_output_self['hidden'], \
- decoder_output_self['cell']
- top_output_self = mdn_top_layer(output_self)
- sampled_step = mdn_top_layer.sample(top_output_self, output_self, opts.mode)
- #print(sampled_step.size())
- sampled_svg[t] = sampled_step
- svg_dec_out = sampled_svg.clone().detach()
- svg_dec_out = svg_dec_out.transpose(0, 1)
- svg_dec_out = svg_dec_out * std + mean
- for i, one_seq in enumerate(svg_dec_out):
- syn_svg_outfile = os.path.join(svg_dir, f"generated{i}.svg")
- syn_svg_f = open(syn_svg_outfile, 'w')
- try:
- svg = render(one_seq.cpu().numpy())
- syn_svg_f.write(svg)
- except:
- continue
- syn_svg_f.close()
- syn_svg_f.close()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement