models/StyleCLIP/global_directions/manipulate.py

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)