Spaces:
Runtime error
Runtime error
File size: 2,332 Bytes
2fec875 |
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 |
from manipulate import Manipulator
import tensorflow as tf
import numpy as np
import torch
import clip
from MapTS import GetBoundary,GetDt
class StyleCLIP():
def __init__(self,dataset_name='ffhq'):
print('load clip')
device = "cuda" if torch.cuda.is_available() else "cpu"
self.model, preprocess = clip.load("ViT-B/32", device=device)
self.LoadData(dataset_name)
def LoadData(self, dataset_name):
tf.keras.backend.clear_session()
M=Manipulator(dataset_name=dataset_name)
np.set_printoptions(suppress=True)
fs3=np.load('./npy/'+dataset_name+'/fs3.npy')
self.M=M
self.fs3=fs3
w_plus=np.load('./data/'+dataset_name+'/w_plus.npy')
self.M.dlatents=M.W2S(w_plus)
if dataset_name=='ffhq':
self.c_threshold=20
else:
self.c_threshold=100
self.SetInitP()
def SetInitP(self):
self.M.alpha=[3]
self.M.num_images=1
self.target=''
self.neutral=''
self.GetDt2()
img_index=0
self.M.dlatent_tmp=[tmp[img_index:(img_index+1)] for tmp in self.M.dlatents]
def GetDt2(self):
classnames=[self.target,self.neutral]
dt=GetDt(classnames,self.model)
self.dt=dt
num_cs=[]
betas=np.arange(0.1,0.3,0.01)
for i in range(len(betas)):
boundary_tmp2,num_c=GetBoundary(self.fs3,self.dt,self.M,threshold=betas[i])
print(betas[i])
num_cs.append(num_c)
num_cs=np.array(num_cs)
select=num_cs>self.c_threshold
if sum(select)==0:
self.beta=0.1
else:
self.beta=betas[select][-1]
def GetCode(self):
boundary_tmp2,num_c=GetBoundary(self.fs3,self.dt,self.M,threshold=self.beta)
codes=self.M.MSCode(self.M.dlatent_tmp,boundary_tmp2)
return codes
def GetImg(self):
codes=self.GetCode()
out=self.M.GenerateImg(codes)
img=out[0,0]
return img
#%%
if __name__ == "__main__":
style_clip=StyleCLIP()
self=style_clip
|