def main(): opt = get_opt() print(opt) model = GMM(opt) load_checkpoint(model, opt.checkpoint) # model.cuda() model.eval() modes = ['train', 'val', 'test'] for mode in modes: print('Run on {} data'.format(mode.upper())) dataset = GMMDataset(opt, mode, data_list=mode+'_pairs.txt', train=False) dataloader = DataLoader(dataset, batch_size=opt.batch_size, num_workers=opt.n_worker, shuffle=False) with torch.no_grad(): run(opt, model, dataloader, mode) print('Successfully completed')
class CPVTON(object): def __init__(self, gmm_path, tom_path): ''' 初始化两个模型的预训练数据 init pretrained models ''' self.gmm = GMM() load_checkpoint(self.gmm, gmm_path) self.gmm.eval() self.tom = UnetGenerator(23, 4, 6, ngf=64, norm_layer=nn.InstanceNorm2d) load_checkpoint(self.tom, tom_path) self.tom.eval() self.gmm.cuda() self.tom.cuda() def predict(self, parse_array, pose_map, human, c): ''' 传入的前四个都是array. shape为(*,256,192) input 4 np array with the shape of (*,256,192) ''' im = transformer(human) c = transformer(c) # [-1,1] # parse -> shape parse_shape = (parse_array > 0).astype(np.float32) # 模糊化,下采样+上采样 # blur, downsample + upsample parse_shape = Image.fromarray((parse_shape * 255).astype(np.uint8)) parse_shape = parse_shape.resize((192 // 16, 256 // 16), Image.BILINEAR) parse_shape = parse_shape.resize((192, 256), Image.BILINEAR) shape = transformer(parse_shape) parse_head = (parse_array == 1).astype(np.float32) + \ (parse_array == 2).astype(np.float32) + \ (parse_array == 4).astype(np.float32) + \ (parse_array == 13).astype(np.float32) + \ (parse_array == 9).astype(np.float32) phead = torch.from_numpy(parse_head) # [0,1] im_h = im * phead - (1 - phead) agnostic = torch.cat([shape, im_h, pose_map], 0) # batch==1 agnostic = agnostic.unsqueeze(0).cuda() c = c.unsqueeze(0).cuda() # warp result grid, theta = self.gmm(agnostic.cuda(), c.cuda()) c_warp = F.grid_sample(c.cuda(), grid, padding_mode='border') tensor = (c_warp.detach().clone() + 1) * 0.5 * 255 tensor = tensor.cpu().clamp(0, 255) array = tensor.numpy().astype('uint8') c_warp = transformer(np.transpose(array[0], axes=(1, 2, 0))) c_warp = c_warp.unsqueeze(0) outputs = self.tom(torch.cat([agnostic.cuda(), c_warp.cuda()], 1)) p_rendered, m_composite = torch.split(outputs, 3, 1) p_rendered = torch.tanh(p_rendered) m_composite = torch.sigmoid(m_composite) p_tryon = c_warp.cuda() * m_composite + p_rendered * (1 - m_composite) return (p_tryon, c_warp)
load_checkpoint(generator, "checkpoints/tom_train_new_2/step_070000.pth") generator.cuda() embedder_model = Embedder() load_checkpoint(embedder_model, "checkpoints/identity_embedding_for_test/step_045000.pth") image_embedder = embedder_model.embedder_b.cuda() prod_embedder = embedder_model.embedder_a.cuda() model = G() if not opt.checkpoint == '' and os.path.exists(opt.checkpoint): load_checkpoint(model, opt.checkpoint) model.cuda() model.eval() gmm_model.eval() image_embedder.eval() generator.eval() prod_embedder.eval() pbar = tqdm(enumerate(data_loader), total=len(data_loader)) product_embeddings = [] outfit_embeddings = [] transfer_embeddings = [] product_embeddings_gt = [] outfit_embeddings_gt = [] for i, (inputs, inputs_2) in pbar: