import os import os.path import pickle import numpy as np import tensorflow as tf from dnnlib import tflib from global_directions.utils.visualizer import HtmlPageVisualizer def Vis(bname,suffix,out,rownames=None,colnames=None): num_images=out.shape[0] step=out.shape[1] if colnames is None: colnames=[f'Step {i:02d}' for i in range(1, step + 1)] if rownames is None: rownames=[str(i) for i in range(num_images)] visualizer = HtmlPageVisualizer( num_rows=num_images, num_cols=step + 1, viz_size=256) visualizer.set_headers( ['Name'] +colnames) for i in range(num_images): visualizer.set_cell(i, 0, text=rownames[i]) for i in range(num_images): for k in range(step): image=out[i,k,:,:,:] visualizer.set_cell(i, 1+k, image=image) # Save results. visualizer.save(f'./html/'+bname+'_'+suffix+'.html') def LoadData(img_path): tmp=img_path+'S' with open(tmp, "rb") as fp: #Pickling s_names,all_s=pickle.load( fp) dlatents=all_s pindexs=[] mindexs=[] for i in range(len(s_names)): name=s_names[i] if not('ToRGB' in name): mindexs.append(i) else: pindexs.append(i) tmp=img_path+'S_mean_std' with open(tmp, "rb") as fp: #Pickling m,std=pickle.load( fp) return dlatents,s_names,mindexs,pindexs,m,std def LoadModel(model_path,model_name): # Initialize TensorFlow. tflib.init_tf() tmp=os.path.join(model_path,model_name) with open(tmp, 'rb') as f: _, _, Gs = pickle.load(f) Gs.print_layers() return Gs def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False): """Convert a minibatch of images from float32 to uint8 with configurable dynamic range. Can be used as an output transformation for Network.run(). """ if nchw_to_nhwc: images = np.transpose(images, [0, 2, 3, 1]) scale = 255 / (drange[1] - drange[0]) images = images * scale + (0.5 - drange[0] * scale) np.clip(images, 0, 255, out=images) images=images.astype('uint8') return images def convert_images_from_uint8(images, drange=[-1,1], nhwc_to_nchw=False): """Convert a minibatch of images from uint8 to float32 with configurable dynamic range. Can be used as an input transformation for Network.run(). """ if nhwc_to_nchw: images=np.rollaxis(images, 3, 1) return images/ 255 *(drange[1] - drange[0])+ drange[0] class Manipulator(): def __init__(self,dataset_name='ffhq'): self.file_path='./' self.img_path=self.file_path+'npy/'+dataset_name+'/' self.model_path=self.file_path+'model/' self.dataset_name=dataset_name self.model_name=dataset_name+'.pkl' self.alpha=[0] #manipulation strength self.num_images=10 self.img_index=0 #which image to start self.viz_size=256 self.manipulate_layers=None #which layer to manipulate, list self.dlatents,self.s_names,self.mindexs,self.pindexs,self.code_mean,self.code_std=LoadData(self.img_path) self.sess=tf.InteractiveSession() init = tf.global_variables_initializer() self.sess.run(init) self.Gs=LoadModel(self.model_path,self.model_name) self.num_layers=len(self.dlatents) self.Vis=Vis self.noise_constant={} for i in range(len(self.s_names)): tmp1=self.s_names[i].split('/') if not 'ToRGB' in tmp1: tmp1[-1]='random_normal:0' size=int(tmp1[1].split('x')[0]) tmp1='/'.join(tmp1) tmp=(1,1,size,size) self.noise_constant[tmp1]=np.random.random(tmp) tmp=self.Gs.components.synthesis.input_shape[1] d={} d['G_synthesis_1/dlatents_in:0']=np.zeros([1,tmp,512]) names=list(self.noise_constant.keys()) tmp=tflib.run(names,d) for i in range(len(names)): self.noise_constant[names[i]]=tmp[i] self.fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) self.img_size=self.Gs.output_shape[-1] def GenerateImg(self,codes): num_images,step=codes[0].shape[:2] out=np.zeros((num_images,step,self.img_size,self.img_size,3),dtype='uint8') for i in range(num_images): for k in range(step): d={} for m in range(len(self.s_names)): d[self.s_names[m]]=codes[m][i,k][None,:] #need to change d['G_synthesis_1/4x4/Const/Shape:0']=np.array([1,18, 512], dtype=np.int32) d.update(self.noise_constant) img=tflib.run('G_synthesis_1/images_out:0', d) image=convert_images_to_uint8(img, nchw_to_nhwc=True) out[i,k,:,:,:]=image[0] return out def MSCode(self,dlatent_tmp,boundary_tmp): step=len(self.alpha) dlatent_tmp1=[tmp.reshape((self.num_images,-1)) for tmp in dlatent_tmp] dlatent_tmp2=[np.tile(tmp[:,None],(1,step,1)) for tmp in dlatent_tmp1] # (10, 7, 512) l=np.array(self.alpha) l=l.reshape( [step if axis == 1 else 1 for axis in range(dlatent_tmp2[0].ndim)]) if type(self.manipulate_layers)==int: tmp=[self.manipulate_layers] elif type(self.manipulate_layers)==list: tmp=self.manipulate_layers elif self.manipulate_layers is None: tmp=np.arange(len(boundary_tmp)) else: raise ValueError('manipulate_layers is wrong') for i in tmp: dlatent_tmp2[i]+=l*boundary_tmp[i] codes=[] for i in range(len(dlatent_tmp2)): tmp=list(dlatent_tmp[i].shape) tmp.insert(1,step) codes.append(dlatent_tmp2[i].reshape(tmp)) return codes def EditOne(self,bname,dlatent_tmp=None): if dlatent_tmp==None: dlatent_tmp=[tmp[self.img_index:(self.img_index+self.num_images)] for tmp in self.dlatents] boundary_tmp=[] for i in range(len(self.boundary)): tmp=self.boundary[i] if len(tmp)<=bname: boundary_tmp.append([]) else: boundary_tmp.append(tmp[bname]) codes=self.MSCode(dlatent_tmp,boundary_tmp) out=self.GenerateImg(codes) return codes,out def EditOneC(self,cindex,dlatent_tmp=None): if dlatent_tmp==None: dlatent_tmp=[tmp[self.img_index:(self.img_index+self.num_images)] for tmp in self.dlatents] boundary_tmp=[[] for i in range(len(self.dlatents))] #'only manipulate 1 layer and one channel' assert len(self.manipulate_layers)==1 ml=self.manipulate_layers[0] tmp=dlatent_tmp[ml].shape[1] #ada tmp1=np.zeros(tmp) tmp1[cindex]=self.code_std[ml][cindex] #1 boundary_tmp[ml]=tmp1 codes=self.MSCode(dlatent_tmp,boundary_tmp) out=self.GenerateImg(codes) return codes,out def W2S(self,dlatent_tmp): all_s = self.sess.run( self.s_names, feed_dict={'G_synthesis_1/dlatents_in:0': dlatent_tmp}) return all_s #%% if __name__ == "__main__": M=Manipulator(dataset_name='ffhq') #%% M.alpha=[-5,0,5] M.num_images=20 lindex,cindex=6,501 M.manipulate_layers=[lindex] codes,out=M.EditOneC(cindex) #dlatent_tmp tmp=str(M.manipulate_layers)+'_'+str(cindex) M.Vis(tmp,'c',out)