File size: 2,623 Bytes
bb0f5a9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109



import numpy as np 
import torch
import clip
from PIL import Image
import copy
from manipulate import Manipulator
import argparse

def GetImgF(out,model,preprocess):
    imgs=out
    imgs1=imgs.reshape([-1]+list(imgs.shape[2:]))
    
    tmp=[]
    for i in range(len(imgs1)):
        
        img=Image.fromarray(imgs1[i])
        image = preprocess(img).unsqueeze(0).to(device)
        tmp.append(image)
    
    image=torch.cat(tmp)
    with torch.no_grad():
        image_features = model.encode_image(image)
    
    image_features1=image_features.cpu().numpy()
    image_features1=image_features1.reshape(list(imgs.shape[:2])+[512])
    
    return image_features1

def GetFs(fs):
    tmp=np.linalg.norm(fs,axis=-1)
    fs1=fs/tmp[:,:,:,None]
    fs2=fs1[:,:,1,:]-fs1[:,:,0,:]  # 5*sigma - (-5)* sigma
    fs3=fs2/np.linalg.norm(fs2,axis=-1)[:,:,None]
    fs3=fs3.mean(axis=1)
    fs3=fs3/np.linalg.norm(fs3,axis=-1)[:,None]
    return fs3

#%%
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Process some integers.')
    
    parser.add_argument('--dataset_name',type=str,default='cat',
                    help='name of dataset, for example, ffhq')
    args = parser.parse_args()
    dataset_name=args.dataset_name
    
    #%%
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, preprocess = clip.load("ViT-B/32", device=device)
    #%%
    M=Manipulator(dataset_name=dataset_name)
    np.set_printoptions(suppress=True)
    print(M.dataset_name)
    #%%
    img_sindex=0
    num_images=100
    dlatents_o=[]
    tmp=img_sindex*num_images
    for i in range(len(M.dlatents)):
        tmp1=M.dlatents[i][tmp:(tmp+num_images)]
        dlatents_o.append(tmp1)
    #%%
    
    all_f=[]
    M.alpha=[-5,5] #ffhq 5
    M.step=2
    M.num_images=num_images
    select=np.array(M.mindexs)<=16 #below or equal to 128 resolution 
    mindexs2=np.array(M.mindexs)[select]
    for lindex in mindexs2: #ignore ToRGB layers
        print(lindex)
        num_c=M.dlatents[lindex].shape[1]
        for cindex in range(num_c):
            
            M.dlatents=copy.copy(dlatents_o)
            M.dlatents[lindex][:,cindex]=M.code_mean[lindex][cindex]
            
            M.manipulate_layers=[lindex]
            codes,out=M.EditOneC(cindex) 
            image_features1=GetImgF(out,model,preprocess)
            all_f.append(image_features1)
    
    all_f=np.array(all_f)
    
    fs3=GetFs(all_f)
    
    #%%
    file_path='./npy/'+M.dataset_name+'/'
    np.save(file_path+'fs3',fs3)