Python多階段框架實現虛擬試衣間,超逼真

作者 | 李秋鍵

責編 | 晉兆雨

頭圖 | CSDN下載自視覺中國

任意姿態下的虛擬試衣因其巨大的應用潛力而引起了人們的廣泛關注。然而，現有的方法在將新穎的服裝和姿勢貼合到一個人身上的同時，很難保留服裝紋理和面部特徵(面孔、毛髮)中的細節。故在論文《Downto the Last Detail: Virtual Try-on with Detail Carving》中提出了一種新的多階段合成框架，可以很好地保留圖像顯著區域的豐富細節。

具體地說，就是提出了一個多階段的框架，將生成分解為空間對齊，然後由粗到細生成。為了更好地保留顯著區域的細節，如服裝和面部區域，我們提出了一個樹塊(樹擴張融合塊)來利用多尺度特徵在發生器網絡。通過多個階段的端到端訓練，可以聯合優化整個框架，最終使得視覺逼真度得到了顯著的提高、同時獲得了細節更為豐富的結果。在標準數據集上進行的大量實驗表明，他們提出的框架實現了最先進的性能，特別是在保存服裝紋理和面部識別的視覺細節方面。

故今天我們將在他們代碼的基礎上，實現虛擬換衣系統。具體流程如下：

實驗前的準備

首先我們使用的python版本是3.6.5所用到的模塊如下：

opencv是將用來進行圖像處理和圖片保存讀取等操作。numpy模塊用來處理矩陣數據的運算。pytorch模塊是常用的用來搭建模型和訓練的深度學習框架,和tensorflow以及Keras等具有相當的地位。json是為了讀取json存儲格式的數據。PIL庫可以完成對圖像進行批處理、生成圖像預覽、圖像格式轉換和圖像處理操作，包括圖像基本處理、像素處理、顏色處理等。argparse 是python自帶的命令行參數解析包，可以用來方便地讀取命令行參數。

網絡模型的定義和訓練

其中已經訓練好的模型地址如下：https://drive.google.com/open?id=1vQo4xNGdYe2uAtur0mDlHY7W2ZR3shWT。其中需要將其中的模型放到"./pretrained_checkpoint"目錄下。

對於數據集的存放，分為cloth_image（用來存儲衣服圖片）,cloth_mask（用來分割衣服的mask，可以使用grabcut的方法進行分割保存）,image（用來存儲人物圖片），parse_cihp（用來衣服語義分析的圖片結果，可以使用[CIHP_PGN](https://github.com/Engineering-Course/CIHP_PGN)的方法獲得）和pose_coco（用來存儲提取到的人物姿態特徵數據，可以使用openpose進行提取保存為josn數據即可）。

對於模型的訓練，我們需要使用到VGG19模型，網絡上可以很容易下載到，然後把它放到vgg_model文件夾下。

其中提出的一種基於目標姿態和店內服裝圖像由粗到細的多階段圖像生成框架，首先是設計了一個解析轉換網絡來預測目標語義圖，該語義圖在空間上對齊相應的身體部位，並提供更多關於軀幹和四肢形狀的結構信息。然後使用一種新的樹擴張融合塊(tree - block)算法，將空間對齊的布料與粗糙的渲染圖像融合在一起，以獲得更合理、更體面的結果。其中這個虛擬試穿網絡不僅不藉助3D信息，可以在任意姿態下將新衣服疊加到人的對應區域上，還保留和增強了顯著區域的豐富細節，如布料紋理、面部特徵等。同時還使用了空間對齊、多尺度上下文特徵聚集和顯著的區域增強，以由粗到細的方式各種難題。

（1）其中網絡主要使用pix2pix模型，其中的部分代碼如下：

class PixelDiscriminator(nn.Module):def__init__(self, input_nc,ndf=64, norm_layer=nn.InstanceNorm2d):super(PixelDiscriminator,self).__init__if type(norm_layer) ==functools.partial:use_bias =norm_layer.func == nn.InstanceNorm2delse:use_bias = norm_layer ==nn.InstanceNorm2dself.net = nn.Sequential(nn.Conv2d(input_nc, ndf,kernel_size=1, stride=1, padding=0),nn.LeakyReLU(0.2, True),nn.Conv2d(ndf, ndf * 2,kernel_size=1, stride=1, padding=0, bias=use_bias),norm_layer(ndf * 2),nn.LeakyReLU(0.2, True),nn.Conv2d(ndf * 2, 1,kernel_size=1, stride=1, padding=0, bias=use_bias),nn.Sigmoid)defforward(self, input):return self.net(input)class PatchDiscriminator(nn.Module):def__init__(self, input_nc,ndf=64, n_layers=3, norm_layer=nn.InstanceNorm2d):super(PatchDiscriminator,self).__init__if type(norm_layer) ==functools.partial: # no need to use biasas BatchNorm2d has affine parametersuse_bias =norm_layer.func == nn.InstanceNorm2delse:use_bias = norm_layer ==nn.InstanceNorm2dkw = 4padw = 1sequence =[nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),nn.LeakyReLU(0.2, True)]nf_mult = 1nf_mult_prev = 1# channel upfor n in range(1,n_layers): # gradually increase thenumber of filtersnf_mult_prev = nf_mult #1,2,4,8nf_mult = min(2 ** n, 8)sequence += [nn.Conv2d(ndf *nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw,bias=use_bias),norm_layer(ndf *nf_mult),nn.LeakyReLU(0.2,True)]# channel downnf_mult_prev = nf_multnf_mult = min(2 ** n_layers,8)sequence += [nn.Conv2d(ndf *nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw,bias=use_bias),norm_layer(ndf *nf_mult),nn.LeakyReLU(0.2, True)]# channel = 1 (bct, 1, x, x)sequence += [nn.Conv2d(ndf *nf_mult, 1, kernel_size=kw, stride=1, padding=padw)] # output 1 channel prediction mapsequence += [nn.Sigmoid()]self.model =nn.Sequential(*sequence)

（2）生成器部分代碼：

class GenerationModel(BaseModel):defname(self):return 'Generation model:pix2pix | pix2pixHD'def__init__(self, opt):self.t0 = timeBaseModel.__init__(self,opt)self.train_mode =opt.train_mode# resume of networksresume_gmm = opt.resume_gmmresume_G_parse =opt.resume_G_parseresume_D_parse =opt.resume_D_parseresume_G_appearance =opt.resume_G_appresume_D_appearance =opt.resume_D_appresume_G_face = opt.resume_G_faceresume_D_face =opt.resume_D_face# define networkself.gmm_model =torch.nn.DataParallel(GMM(opt)).cudaself.generator_parsing =Define_G(opt.input_nc_G_parsing, opt.output_nc_parsing, opt.ndf, opt.netG_parsing,opt.norm,not opt.no_dropout, opt.init_type, opt.init_gain, opt.gpu_ids)self.discriminator_parsing =Define_D(opt.input_nc_D_parsing, opt.ndf, opt.netD_parsing, opt.n_layers_D,opt.norm, opt.init_type, opt.init_gain, opt.gpu_ids)self.generator_appearance =Define_G(opt.input_nc_G_app, opt.output_nc_app, opt.ndf, opt.netG_app,opt.norm,not opt.no_dropout, opt.init_type, opt.init_gain, opt.gpu_ids,with_tanh=False)self.discriminator_appearance = Define_D(opt.input_nc_D_app, opt.ndf,opt.netD_app, opt.n_layers_D,opt.norm, opt.init_type, opt.init_gain, opt.gpu_ids)self.generator_face =Define_G(opt.input_nc_D_face, opt.output_nc_face, opt.ndf, opt.netG_face,opt.norm,notopt.no_dropout, opt.init_type, opt.init_gain, opt.gpu_ids)self.discriminator_face =Define_D(opt.input_nc_D_face, opt.ndf, opt.netD_face, opt.n_layers_D,opt.norm, opt.init_type, opt.init_gain, opt.gpu_ids)if opt.train_mode == 'gmm':setattr(self,'generator', self.gmm_model)else:setattr(self,'generator', getattr(self, 'generator_' + self.train_mode))setattr(self, 'discriminator',getattr(self, 'discriminator_' + self.train_mode))# load networksself.networks_name = ['gmm','parsing', 'parsing', 'appearance', 'appearance', 'face', 'face']self.networks_model =[self.gmm_model, self.generator_parsing, self.discriminator_parsing,self.generator_appearance, self.discriminator_appearance,self.generator_face, self.discriminator_face]self.networks =dict(zip(self.networks_name, self.networks_model))self.resume_path =[resume_gmm, resume_G_parse, resume_D_parse, resume_G_appearance,resume_D_appearance, resume_G_face, resume_D_face]for network, resume inzip(self.networks_model, self.resume_path):if network != andresume != '':assert(osp.exists(resume), 'the resume not exits')print('loading...')self.load_network(network, resume, ifprint=False)# define optimizerself.optimizer_gmm =torch.optim.Adam(self.gmm_model.parameters, lr=opt.lr, betas=(0.5, 0.999))self.optimizer_parsing_G =torch.optim.Adam(self.generator_parsing.parameters, lr=opt.lr,betas=[opt.beta1, 0.999])self.optimizer_parsing_D =torch.optim.Adam(self.discriminator_parsing.parameters, lr=opt.lr,betas=[opt.beta1, 0.999])self.optimizer_appearance_G= torch.optim.Adam(self.generator_appearance.parameters, lr=opt.lr,betas=[opt.beta1, 0.999])self.optimizer_appearance_D= torch.optim.Adam(self.discriminator_appearance.parameters, lr=opt.lr,betas=[opt.beta1, 0.999])self.optimizer_face_G =torch.optim.Adam(self.generator_face.parameters, lr=opt.lr, betas=[opt.beta1,0.999])self.optimizer_face_D =torch.optim.Adam(self.discriminator_face.parameters, lr=opt.lr,betas=[opt.beta1, 0.999])if opt.train_mode == 'gmm':self.optimizer_G =self.optimizer_gmmelif opt.joint_all:self.optimizer_G =[self.optimizer_parsing_G, self.optimizer_appearance_G, self.optimizer_face_G]setattr(self,'optimizer_D', getattr(self, 'optimizer_' + self.train_mode + '_D'))else:setattr(self,'optimizer_G', getattr(self, 'optimizer_' + self.train_mode + '_G'))setattr(self,'optimizer_D', getattr(self, 'optimizer_' + self.train_mode + '_D'))self.t1 = time

模型的使用

在模型訓練完成之後，通過命令「python demo.py --batch_size_v 80--num_workers 4 --forward_save_path 'demo/forward'」生成圖片。

（1）分別定義讀取模型函數和模型調用處理圖片函數

def load_model(model, path):checkpoint = torch.load(path)try:model.load_state_dict(checkpoint)except:model.load_state_dict(checkpoint.state_dict)model = model.cudamodel.evalprint(20*'=')for param in model.parameters:param.requires_grad = Falsedef forward(opt, paths, gpu_ids, refine_path):cudnn.enabled = Truecudnn.benchmark = Trueopt.output_nc = 3gmm = GMM(opt)gmm =torch.nn.DataParallel(gmm).cuda# 'batch'generator_parsing =Define_G(opt.input_nc_G_parsing, opt.output_nc_parsing, opt.ndf,opt.netG_parsing, opt.norm,notopt.no_dropout, opt.init_type, opt.init_gain, opt.gpu_ids)generator_app_cpvton =Define_G(opt.input_nc_G_app, opt.output_nc_app, opt.ndf, opt.netG_app,opt.norm,notopt.no_dropout, opt.init_type, opt.init_gain, opt.gpu_ids, with_tanh=False)generator_face =Define_G(opt.input_nc_D_face, opt.output_nc_face, opt.ndf, opt.netG_face,opt.norm,notopt.no_dropout, opt.init_type, opt.init_gain, opt.gpu_ids)models = [gmm,generator_parsing, generator_app_cpvton, generator_face]for model, path in zip(models,paths):load_model(model, path)print('==>loaded model')augment = {}if '0.4' in torch.__version__:augment['3'] =transforms.Compose([# transforms.Resize(256),transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))]) # change to [C, H, W]augment['1'] = augment['3']else:augment['3'] =transforms.Compose([#transforms.Resize(256),transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))]) # change to [C, H, W]augment['1'] =transforms.Compose([# transforms.Resize(256),transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))]) # change to [C, H, W]val_dataset = DemoDataset(opt,augment=augment)val_dataloader = DataLoader(val_dataset,shuffle=False,drop_last=False,num_workers=opt.num_workers,batch_size = opt.batch_size_v,pin_memory=True)with torch.no_grad:for i, result inenumerate(val_dataloader):'warped cloth'warped_cloth =warped_image(gmm, result)if opt.warp_cloth:warped_cloth_name =result['warped_cloth_name']warped_cloth_path =os.path.join('dataset', 'warped_cloth', warped_cloth_name[0])if notos.path.exists(os.path.split(warped_cloth_path)[0]):os.makedirs(os.path.split(warped_cloth_path)[0])utils.save_image(warped_cloth * 0.5 + 0.5, warped_cloth_path)print('processing_%d'%i)continuesource_parse =result['source_parse'].float.cudatarget_pose_embedding =result['target_pose_embedding'].float.cudasource_image =result['source_image'].float.cudacloth_parse =result['cloth_parse'].cudacloth_image =result['cloth_image'].cudatarget_pose_img =result['target_pose_img'].float.cudacloth_parse =result['cloth_parse'].float.cudasource_parse_vis =result['source_parse_vis'].float.cuda"filter add clothinfomation"real_s =source_parseindex = [x for x inlist(range(20)) if x != 5 and x != 6 and x != 7]real_s_ =torch.index_select(real_s, 1, torch.tensor(index).cuda)input_parse =torch.cat((real_s_, target_pose_embedding, cloth_parse), 1).cuda'P'generate_parse =generator_parsing(input_parse) # tanhgenerate_parse =F.softmax(generate_parse, dim=1)generate_parse_argmax =torch.argmax(generate_parse, dim=1, keepdim=True).floatres = for index in range(20):res.append(generate_parse_argmax == index)generate_parse_argmax =torch.cat(res, dim=1).float"A"image_without_cloth =create_part(source_image, source_parse, 'image_without_cloth', False)input_app = torch.cat((image_without_cloth,warped_cloth, generate_parse), 1).cuda

源碼地址：

提取碼：qcj6