def demo(prev, after):
# Prepare img pair
#im1 = imread('FlowNet2_src/example/0img0.ppm')
#im2 = imread('FlowNet2_src/example/0img1.ppm')
im1 = cv2.resize(imread(prev), (480, 360))
im2 = cv2.resize(imread(after), (480, 360))
# B x 3(RGB) x 2(pair) x H x W
ims = np.array([[im1, im2]]).transpose((0, 4, 1, 2, 3)).astype(np.float32)
ims = torch.from_numpy(ims)
ims_v = Variable(ims.cuda(), requires_grad=False)
# Build model
flownet2 = FlowNet2()
path = '/home/mike/workspace/segnet/pytorch_flownet2/FlowNet2_src/pretrained/FlowNet2_checkpoint.pth.tar'
pretrained_dict = torch.load(path)['state_dict']
model_dict = flownet2.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
flownet2.load_state_dict(model_dict)
flownet2.cuda()
t1 = time.time()
pred_flow = flownet2(ims_v).cpu().data
t2 = time.time()
print t2 - t1
pred_flow = pred_flow[0].numpy().transpose((1, 2, 0))
flow_im = flow_to_image(pred_flow)
# Visualization
plt.imshow(flow_im)
plt.show()
#pdb.set_trace()
return pred_flow
im1_pl = tf.placeholder(tf.float32, [1, 384, 512, 3])
im2_pl = tf.placeholder(tf.float32, [1, 384, 512, 3])
flownet2 = FlowNetSD()
inputs = {'input_a': im1_pl, 'input_b': im2_pl}
flow_dict = flownet2.model(inputs, LONG_SCHEDULE, trainable=False)
pred_flow = flow_dict['flow']
# Feed forward
im1 = imread('FlowNet2_src/example/0img0.ppm')/255.
im2 = imread('FlowNet2_src/example/0img1.ppm')/255.
im1 = np.array([im1]).astype(np.float32)
im2 = np.array([im2]).astype(np.float32)
ckpt_file = 'FlowNet2_src/checkpoints/FlowNet2/flownet-2.ckpt-0'
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, ckpt_file)
# Double check loading is correct
#for var in tf.all_variables():
# print(var.name, var.eval(session=sess).mean())
feed_dict = {im1_pl: im1, im2_pl: im2}
pred_flow_val = sess.run(pred_flow, feed_dict=feed_dict)
# Visualization
import matplotlib.pyplot as plt
flow_im = flow_to_image(pred_flow_val[0])
plt.imshow(flow_im)
plt.show()
if __name__ == '__main__':
# Prepare img pair
im1 = imread('FlowNet2_src/example/0img0.ppm')
im2 = imread('FlowNet2_src/example/0img1.ppm')
# B x 3(RGB) x 2(pair) x H x W
ims = np.array([[im1, im2]]).transpose((0, 4, 1, 2, 3)).astype(np.float32)
ims = torch.from_numpy(ims)
ims_v = Variable(ims.cuda(), requires_grad=False)
# Build model
flownet2 = FlowNet2()
path = 'FlowNet2_src/pretrained/FlowNet2_checkpoint.pth.tar'
pretrained_dict = torch.load(path)['state_dict']
model_dict = flownet2.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
flownet2.load_state_dict(model_dict)
flownet2.cuda()
pred_flow = flownet2(ims_v).cpu().data
pred_flow = pred_flow[0].numpy().transpose((1, 2, 0))
flow_im = flow_to_image(pred_flow)
# Visualization
plt.imshow(flow_im)
plt.show()
im1 = imread('FlowNet2_src/2_crop/3970.jpg') / 255. #(256,256)
im2 = imread('FlowNet2_src/2_crop/3971.jpg') / 255.
im1 = np.array([im1]).astype(np.float32) # （1,256,256）
print(im1.shape)
im1 = np.stack([im1, im1, im1], axis=-1) # 增加一个维度 （1,256,256,3）
print(im2.shape)
im2 = np.array([im2]).astype(np.float32)
im2 = np.stack([im2, im2, im2], axis=-1)
print(im1.shape)
print(im2.shape)
ckpt_file = 'FlowNet2_src/checkpoints/FlowNet2/flownet-2.ckpt-0'
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, ckpt_file)
# Double check loading is correct
#for var in tf.all_variables():
# print(var.name, var.eval(session=sess).mean())
feed_dict = {im1_pl: im1, im2_pl: im2}
pred_flow_val = sess.run(pred_flow, feed_dict=feed_dict)
print('preflow[0]', pred_flow_val[0].shape)
# Visualization
import matplotlib.pyplot as plt
flow_im = flow_to_image(pred_flow_val[0]) # 三通道的光流图
print('flow_to_image', flow_im.shape)
print(flow_im)
plt.imshow(flow_im)
plt.show()
flownet2 = FlowNet2()
inputs = {'input_a': im1_pl, 'input_b': im2_pl}
flow_dict = flownet2.model(inputs, LONG_SCHEDULE, trainable=False)
pred_flow = flow_dict['flow']
cpkt_file = 'FlowNet2_src/checkpoints/FlowNet2/flownet-2.ckpt-0'
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, cpkt_file)
for n_frame in range(0, len(imagelist) - 1):
frame_pre = imagelist[n_frame] # 上一帧图像
frame_cur = imagelist[n_frame + 1] # 当前帧图像
im1, im2 = process_img(frame_pre, frame_cur) # 处理图像为矩阵格式
feed_dict = {im1_pl: im1, im2_pl: im2}
pred_flow_val = sess.run(pred_flow, feed_dict=feed_dict)
flow_im = flow_to_image(pred_flow_val[0]) # 将光流数据转为图像
cv2.imwrite(os.path.join(output, os.path.basename(frame_cur)),
flow_im)
# 保存图像到指定路径，output的路径加上当前帧的最后边的名字。
# for filename in os.listdir('zth_src/4-6'):
# print(filename)
# def process_img(frame_1, frame_2):
# # Feed forward
# im1 = frame_1/255.
# im2 = frame_2/255.
# print('im1', im1.shape)
# im1 = np.array([im1]).astype(np.float32)
# print('npim1',im1.shape)
# #im1 = np.stack([im1,im1,im1], axis=-1)
def flow_to_im(flow):
return flow_to_image(flow.numpy().transpose((1, 2, 0)))
