from src.flowlib import flow_to_image
from src.flownet2.flownet2 import FlowNet2
from src.training_schedules import LONG_SCHEDULE
img_dir = [
y for x in os.walk('./test') for y in glob(os.path.join(x[0], '*.png'))
]
img_dir.sort()
a = cv2.imread(img_dir[0]) / 255.0
b = cv2.imread(img_dir[1]) / 255.0
inputs = {
'input_a': tf.expand_dims(tf.constant(a, dtype=tf.float32), 0),
'input_b': tf.expand_dims(tf.constant(b, dtype=tf.float32), 0)
}
model = FlowNet2(mode=Mode.TEST, debug=False)
predictions = model(inputs, LONG_SCHEDULE)
pred_flow = predictions['flow']
saver = tf.train.Saver()
with tf.Session() as sess:
#saver = tf.train.import_meta_graph('./checkpoints/FlowNet2/flownet-2.ckpt-0.meta')
saver.restore(sess, './checkpoints/FlowNet2/flownet-2.ckpt-0')
pred = sess.run(pred_flow)
print(pred.keys())
pred_flow = pred['flow']
fimg = flow_to_image(pred_flow)
cv2.imwrite('aaa.png', fimg)
def main():
"""
.Input("image_a: float32")
.Input("image_b: float32")
.Attr("crop: list(int) >= 2")
.Attr("params_a_name: list(string)")
.Attr("params_a_rand_type: list(string)")
.Attr("params_a_exp: list(bool)")
.Attr("params_a_mean: list(float32)")
.Attr("params_a_spread: list(float32)")
.Attr("params_a_prob: list(float32)")
.Attr("params_b_name: list(string)")
.Attr("params_b_rand_type: list(string)")
.Attr("params_b_exp: list(bool)")
.Attr("params_b_mean: list(float32)")
.Attr("params_b_spread: list(float32)")
.Attr("params_b_prob: list(float32)")
.Output("aug_image_a: float32")
.Output("aug_image_b: float32")
.Output("spatial_transform_a: float32")
.Output("inv_spatial_transform_b: float32")
"""
crop = [364, 492]
params_a_name = ['translate_x', 'translate_y']
params_a_rand_type = ['uniform_bernoulli', 'uniform_bernoulli']
params_a_exp = [False, False]
params_a_mean = [0.0, 0.0]
params_a_spread = [0.4, 0.4]
params_a_prob = [1.0, 1.0]
params_b_name = []
params_b_rand_type = []
params_b_exp = []
params_b_mean = []
params_b_spread = []
params_b_prob = []
with tf.Session() as sess:
with tf.device('/gpu:0'):
image_a = imread('./img0.ppm') / 255.0
image_b = imread('./img1.ppm') / 255.0
flow = read_flow('./flow.flo')
image_a_tf = tf.expand_dims(
tf.to_float(tf.constant(image_a, dtype=tf.float64)), 0)
image_b_tf = tf.expand_dims(
tf.to_float(tf.constant(image_b, dtype=tf.float64)), 0)
preprocess = _preprocessing_ops.data_augmentation(
image_a_tf, image_b_tf, crop, params_a_name,
params_a_rand_type, params_a_exp, params_a_mean,
params_a_spread, params_a_prob, params_b_name,
params_b_rand_type, params_b_exp, params_b_mean,
params_b_spread, params_b_prob)
out = sess.run(preprocess)
trans = out.spatial_transform_a
inv_trans = out.inv_spatial_transform_b
print(trans.shape)
print(inv_trans.shape)
flow_tf = tf.expand_dims(tf.to_float(tf.constant(flow)), 0)
aug_flow_tf = _preprocessing_ops.flow_augmentation(
flow_tf, trans, inv_trans, crop)
aug_flow = sess.run(aug_flow_tf)[0, :, :, :]
# Plot img0, img0aug
plt.subplot(321)
plt.imshow(image_a)
plt.subplot(322)
plt.imshow(out.aug_image_a[0, :, :, :])
# Plot img1, img1aug
plt.subplot(323)
plt.imshow(image_b)
plt.subplot(324)
plt.imshow(out.aug_image_b[0, :, :, :])
# Plot flow, flowaug
plt.subplot(325)
plt.imshow(flow_to_image(flow))
plt.subplot(326)
plt.imshow(flow_to_image(aug_flow))
plt.show()
# image_b_aug = sess.run(image_b_tf)
#
# display(np.expand_dims(image_a, 0), 0)
# display(np.expand_dims(image_b, 0), 1)
# display(image_a_aug, 2)
# display(image_b_aug, 3)
# plt.show()
# o = _preprocessing_ops.flow_augmentation(flow, trans, inv_t, [4, 8])
# print n[:, :, :]
# print n[0, 0, 1], n[0, 0, 0]
# print n[1, 0, 1], n[1, 0, 0]
# print n[2, 0, 1], n[2, 0, 0]
# print '---'
# print sess.run(o)
"""# Goes along width first!!
def main():
"""
.Input("image_a: float32")
.Input("image_b: float32")
.Attr("crop: list(int) >= 2")
.Attr("params_a_name: list(string)")
.Attr("params_a_rand_type: list(string)")
.Attr("params_a_exp: list(bool)")
.Attr("params_a_mean: list(float32)")
.Attr("params_a_spread: list(float32)")
.Attr("params_a_prob: list(float32)")
.Attr("params_b_name: list(string)")
.Attr("params_b_rand_type: list(string)")
.Attr("params_b_exp: list(bool)")
.Attr("params_b_mean: list(float32)")
.Attr("params_b_spread: list(float32)")
.Attr("params_b_prob: list(float32)")
.Output("aug_image_a: float32")
.Output("aug_image_b: float32")
.Output("spatial_transform_a: float32")
.Output("inv_spatial_transform_b: float32")
"""
crop = [364, 492]
params_a_name = ['translate_x', 'translate_y']
params_a_rand_type = ['uniform_bernoulli', 'uniform_bernoulli']
params_a_exp = [False, False]
params_a_mean = [0.0, 0.0]
params_a_spread = [0.4, 0.4]
params_a_prob = [1.0, 1.0]
params_b_name = []
params_b_rand_type = []
params_b_exp = []
params_b_mean = []
params_b_spread = []
params_b_prob = []
with tf.Session() as sess:
with tf.device('/gpu:0'):
image_a = imread('./img0.ppm') / 255.0
image_b = imread('./img1.ppm') / 255.0
flow = read_flow('./flow.flo')
image_a_tf = tf.expand_dims(tf.to_float(tf.constant(image_a, dtype=tf.float64)), 0)
image_b_tf = tf.expand_dims(tf.to_float(tf.constant(image_b, dtype=tf.float64)), 0)
preprocess = _preprocessing_ops.data_augmentation(image_a_tf,
image_b_tf,
crop,
params_a_name,
params_a_rand_type,
params_a_exp,
params_a_mean,
params_a_spread,
params_a_prob,
params_b_name,
params_b_rand_type,
params_b_exp,
params_b_mean,
params_b_spread,
params_b_prob)
out = sess.run(preprocess)
trans = out.spatial_transform_a
inv_trans = out.inv_spatial_transform_b
print trans.shape
print inv_trans.shape
flow_tf = tf.expand_dims(tf.to_float(tf.constant(flow)), 0)
aug_flow_tf = _preprocessing_ops.flow_augmentation(flow_tf, trans, inv_trans, crop)
aug_flow = sess.run(aug_flow_tf)[0, :, :, :]
# Plot img0, img0aug
plt.subplot(321)
plt.imshow(image_a)
plt.subplot(322)
plt.imshow(out.aug_image_a[0, :, :, :])
# Plot img1, img1aug
plt.subplot(323)
plt.imshow(image_b)
plt.subplot(324)
plt.imshow(out.aug_image_b[0, :, :, :])
# Plot flow, flowaug
plt.subplot(325)
plt.imshow(flow_to_image(flow))
plt.subplot(326)
plt.imshow(flow_to_image(aug_flow))
plt.show()
# image_b_aug = sess.run(image_b_tf)
#
# display(np.expand_dims(image_a, 0), 0)
# display(np.expand_dims(image_b, 0), 1)
# display(image_a_aug, 2)
# display(image_b_aug, 3)
# plt.show()
# o = _preprocessing_ops.flow_augmentation(flow, trans, inv_t, [4, 8])
# print n[:, :, :]
# print n[0, 0, 1], n[0, 0, 0]
# print n[1, 0, 1], n[1, 0, 0]
# print n[2, 0, 1], n[2, 0, 0]
# print '---'
# print sess.run(o)
"""# Goes along width first!!