def __init__(self, model_dir, sess,
use_mu=False, ret_c_centers=False, ret_c5=False,
checkpoint=None):
assert not (ret_c_centers and ret_c5)
self.sess = sess
self.ret_c5 = ret_c5
self.ret_c_centers = ret_c_centers
self.images = tf.placeholder(tf.float32, [None, vh, vw, 3])
ret = vss(self.images, False, True,
ret_mu=use_mu, ret_c_centers=ret_c_centers,
ret_c5=ret_c5)
if ret_c5:
self.descriptor = ret[0]
self.c5 = ret[1]
elif ret_c_centers:
self.descriptor = ret[0]
self.cc = ret[1]
else:
self.descriptor = ret if not use_mu else tf.reduce_mean(ret, axis=0, keepdims=True)
saver = tf.train.Saver()
if checkpoint is None:
ckpt = tf.train.get_checkpoint_state(model_dir)
cpath = ckpt.model_checkpoint_path
else:
cpath = checkpoint
print("loading model: ", cpath)
saver.restore(self.sess, cpath)
def show_example(image_fl, model_dir):
im = cv2.cvtColor(cv2.resize(cv2.imread(image_fl),
(vw, vh)), cv2.COLOR_BGR2RGB)[np.newaxis, ...] / 255.0
_im = tf.placeholder_with_default(im.astype(np.float32), im.shape)
_, _, rec, seg, _, _, _ = vss(_im, False)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
rec, seg = sess.run([rec, seg])
rec = (255*np.squeeze(rec)).astype(np.uint8)
seg = np.argmax(np.squeeze(seg), axis=-1)
rgb_seg = np.zeros(rec.shape).astype(np.uint8)
np.random.seed(0)
for i in range(N_CLASSES):
c = np.random.rand(3)
case = seg==i
rgb_seg[case, :] = c
if not os.path.isdir('plots'):
os.mkdir('plots')
cv2.imwrite('plots/seg.jpg', rgb_seg)
cv2.imwrite('plots/rec.jpg', rec)