-
Notifications
You must be signed in to change notification settings - Fork 12
/
vis.py
63 lines (51 loc) · 2.8 KB
/
vis.py
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
# -*- coding:utf-8 -*-
# Created Time: 2017/11/04 21:19:32
# Author: Taihong Xiao <[email protected]>
import tensorflow as tf
import numpy as np
import os
from scipy import misc
from model import Model
from dataset import Dataset
import argparse
def get_representation(img, model, model_dir, out_path):
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
ckpt = tf.train.get_checkpoint_state(model_dir)
# print(ckpt)
# print(ckpt.model_checkpoint_path)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
Ax = tf.placeholder(tf.float32, [model.batch_size,model.height,model.width,model.channel],name='Ax')
enc_Ax = model.splitter('encoder', Ax)
grad_att_1 =[tf.gradients(enc_Ax[0][:,:,:,i], Ax)[0] for i in range(128)]
grad_att_2 =[tf.gradients(enc_Ax[1][:,:,:,i], Ax)[0] for i in range(128)]
grad_att_3 =[tf.gradients(enc_Ax[2][:,:,:,i], Ax)[0] for i in range(256)]
# from IPython import embed;embed();exit()
grad_1 = sess.run(grad_att_1, feed_dict={Ax: img})
grad_2 = sess.run(grad_att_2, feed_dict={Ax: img})
grad_3 = sess.run(grad_att_3, feed_dict={Ax: img})
for i in range(128):
misc.imsave(os.path.join(out_path, '0_{:03d}.jpg'.format(i)), grad_1[i][0])
misc.imsave(os.path.join(out_path, '1_{:03d}.jpg'.format(i)), grad_2[i][0])
np.save(os.path.join(out_path, '0_{:03d}.npy'.format(i)), grad_2[i][0])
np.save(os.path.join(out_path, '1_{:03d}.npy'.format(i)), grad_2[i][0])
for i in range(256):
misc.imsave(os.path.join(out_path, '2_{:03d}.jpg'.format(i)), grad_3[i][0])
np.save(os.path.join(out_path, '2_{:03d}.npy'.format(i)), grad_3[i][0])
def main():
parser = argparse.ArgumentParser(description='test', formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument('-i', '--input', type=str, help='image path')
parser.add_argument('-a', '--attributes', type=str, nargs='+', help='attribute list')
parser.add_argument('--model_dir', type=str, default='train_log/model/', help='path to the model')
parser.add_argument('--latent_path', type=str, default='latent', help='path to the model')
parser.add_argument('-g', '--gpu', type=str, default='', help='gpu ids')
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
if not os.path.exists(args.latent_path): os.makedirs(args.latent_path)
DNA_GAN = Model(feature_list=args.attributes, is_train=False, nhwc=[1,64,64,3])
img = np.expand_dims(misc.imresize(misc.imread(args.input), (DNA_GAN.height, DNA_GAN.width)), axis=0)
get_representation(img, DNA_GAN, args.model_dir, args.latent_path)
if __name__ == "__main__":
main()