models/StyleCLIP/global_directions/GetCode.py

import os
import pickle
import numpy as np
from dnnlib import tflib  
import tensorflow as tf 

import argparse

def LoadModel(dataset_name):
    # Initialize TensorFlow.
    tflib.init_tf()
    model_path='./model/'
    model_name=dataset_name+'.pkl'
    
    tmp=os.path.join(model_path,model_name)
    with open(tmp, 'rb') as f:
        _, _, Gs = pickle.load(f)
    return Gs

def lerp(a,b,t):
     return a + (b - a) * t

#stylegan-ada
def SelectName(layer_name,suffix):
    if suffix==None:
        tmp1='add:0' in layer_name 
        tmp2='shape=(?,' in layer_name
        tmp4='G_synthesis_1' in layer_name
        tmp= tmp1 and tmp2 and tmp4  
    else:
        tmp1=('/Conv0_up'+suffix) in layer_name 
        tmp2=('/Conv1'+suffix) in layer_name 
        tmp3=('4x4/Conv'+suffix) in layer_name 
        tmp4='G_synthesis_1' in layer_name
        tmp5=('/ToRGB'+suffix) in layer_name
        tmp= (tmp1 or tmp2 or tmp3 or tmp5) and tmp4 
    return tmp


def GetSNames(suffix):
    #get style tensor name 
    with tf.Session() as sess:
        op = sess.graph.get_operations()
    layers=[m.values() for m in op]
    
    
    select_layers=[]
    for layer in layers:
        layer_name=str(layer)
        if SelectName(layer_name,suffix):
            select_layers.append(layer[0])
    return select_layers

def SelectName2(layer_name):
    tmp1='mod_bias' in layer_name 
    tmp2='mod_weight' in layer_name
    tmp3='ToRGB' in layer_name 
    
    tmp= (tmp1 or tmp2) and (not tmp3) 
    return tmp

def GetKName(Gs):
    
    layers=[var for name, var in Gs.components.synthesis.vars.items()]
    
    select_layers=[]
    for layer in layers:
        layer_name=str(layer)
        if SelectName2(layer_name):
            select_layers.append(layer)
    return select_layers

def GetCode(Gs,random_state,num_img,num_once,dataset_name):
    rnd = np.random.RandomState(random_state)  #5
    
    truncation_psi=0.7
    truncation_cutoff=8
    
    dlatent_avg=Gs.get_var('dlatent_avg')
    
    dlatents=np.zeros((num_img,512),dtype='float32')
    for i in range(int(num_img/num_once)):
        src_latents =  rnd.randn(num_once, Gs.input_shape[1])
        src_dlatents = Gs.components.mapping.run(src_latents, None) # [seed, layer, component]
        
        # Apply truncation trick.
        if truncation_psi is not None and truncation_cutoff is not None:
                layer_idx = np.arange(src_dlatents.shape[1])[np.newaxis, :, np.newaxis]
                ones = np.ones(layer_idx.shape, dtype=np.float32)
                coefs = np.where(layer_idx < truncation_cutoff, truncation_psi * ones, ones)
                src_dlatents_np=lerp(dlatent_avg, src_dlatents, coefs)
                src_dlatents=src_dlatents_np[:,0,:].astype('float32')
                dlatents[(i*num_once):((i+1)*num_once),:]=src_dlatents
    print('get all z and w')
    
    tmp='./npy/'+dataset_name+'/W'
    np.save(tmp,dlatents)

    
def GetImg(Gs,num_img,num_once,dataset_name,save_name='images'):
    print('Generate Image')
    tmp='./npy/'+dataset_name+'/W.npy'
    dlatents=np.load(tmp) 
    fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
    
    all_images=[]
    for i in range(int(num_img/num_once)):
        print(i)
        images=[]
        for k in range(num_once):
            tmp=dlatents[i*num_once+k]
            tmp=tmp[None,None,:]
            tmp=np.tile(tmp,(1,Gs.components.synthesis.input_shape[1],1))
            image2= Gs.components.synthesis.run(tmp, randomize_noise=False, output_transform=fmt)
            images.append(image2)
            
        images=np.concatenate(images)
        
        all_images.append(images)
        
    all_images=np.concatenate(all_images)
    
    tmp='./npy/'+dataset_name+'/'+save_name
    np.save(tmp,all_images)

def GetS(dataset_name,num_img):
    print('Generate S')
    tmp='./npy/'+dataset_name+'/W.npy'
    dlatents=np.load(tmp)[:num_img]
    
    with tf.Session() as sess:
        init = tf.global_variables_initializer()
        sess.run(init)
        
        Gs=LoadModel(dataset_name)
        Gs.print_layers()  #for ada
        select_layers1=GetSNames(suffix=None)  #None,'/mul_1:0','/mod_weight/read:0','/MatMul:0'
        dlatents=dlatents[:,None,:]
        dlatents=np.tile(dlatents,(1,Gs.components.synthesis.input_shape[1],1))
        
        all_s = sess.run(
            select_layers1,
            feed_dict={'G_synthesis_1/dlatents_in:0': dlatents})
    
    layer_names=[layer.name for layer in select_layers1]
    save_tmp=[layer_names,all_s]
    return save_tmp

    


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 GetCodeMS(dlatents):
        m=[]
        std=[]
        for i in range(len(dlatents)):
            tmp= dlatents[i] 
            tmp_mean=tmp.mean(axis=0)
            tmp_std=tmp.std(axis=0)
            m.append(tmp_mean)
            std.append(tmp_std)
        return m,std



#%%
if __name__ == "__main__":
    
    
    parser = argparse.ArgumentParser(description='Process some integers.')
    
    parser.add_argument('--dataset_name',type=str,default='ffhq',
                    help='name of dataset, for example, ffhq')
    parser.add_argument('--code_type',choices=['w','s','s_mean_std'],default='w')
    
    args = parser.parse_args()
    random_state=5
    num_img=100_000 
    num_once=1_000
    dataset_name=args.dataset_name
    
    if not os.path.isfile('./model/'+dataset_name+'.pkl'):
        url='https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/'
        name='stylegan2-'+dataset_name+'-config-f.pkl'
        os.system('wget ' +url+name + '  -P  ./model/')
        os.system('mv ./model/'+name+' ./model/'+dataset_name+'.pkl')
    
    if not os.path.isdir('./npy/'+dataset_name):
        os.system('mkdir ./npy/'+dataset_name)
    
    if args.code_type=='w':
        Gs=LoadModel(dataset_name=dataset_name)
        GetCode(Gs,random_state,num_img,num_once,dataset_name)
#        GetImg(Gs,num_img=num_img,num_once=num_once,dataset_name=dataset_name,save_name='images_100K') #no need 
    elif args.code_type=='s':
        save_name='S'
        save_tmp=GetS(dataset_name,num_img=2_000)
        tmp='./npy/'+dataset_name+'/'+save_name
        with open(tmp, "wb") as fp:
            pickle.dump(save_tmp, fp)
        
    elif args.code_type=='s_mean_std':
        save_tmp=GetS(dataset_name,num_img=num_img)
        dlatents=save_tmp[1]
        m,std=GetCodeMS(dlatents)
        save_tmp=[m,std]
        save_name='S_mean_std'
        tmp='./npy/'+dataset_name+'/'+save_name
        with open(tmp, "wb") as fp:
            pickle.dump(save_tmp, fp)