def _build(self):
# Graph construction
#im1_pl = tf.placeholder(tf.float32, [None, None, None, 3])
#im2_pl = tf.placeholder(tf.float32, [None, None, None, 3])
#im1_in = tf.image.resize_images(im1_pl, (384,512))
#im2_in = tf.image.resize_images(im2_pl, (384,512))
im1_in = im1_pl = tf.placeholder(tf.float32, [None, 384, 512, 3])
im2_in = im2_pl = tf.placeholder(tf.float32, [None, 384, 512, 3])
flownet2 = FlowNet2()
inputs = {'input_a': im1_in, 'input_b': im2_in}
flow_dict = flownet2.model(inputs, LONG_SCHEDULE, trainable=False)
pred_flow = flow_dict['flow']
self.im1_ = im1_pl
self.im2_ = im2_pl
self.flow_ = pred_flow
def get_network():
network = FlowNet2()
model_path = MODEL_PATH
pretrained_dict = torch.load(model_path)['state_dict']
model_dict = network.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
network.load_state_dict(model_dict)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
network = nn.DataParallel(network)
return network.to(DEVICE)
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
from FlowNet2_src import FlowNet2
from FlowNet2_src import flow_to_image
import matplotlib.pyplot as plt
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)
def calc_optical_flow(dataset):
of_root_dir = './optical_flow'
len_original_root_dir = len(dataset.dir.split('/')) - 1
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()
dataset_loader = DataLoader(dataset=dataset,
batch_size=1,
shuffle=False,
num_workers=1)
for idx, (batch, _) in enumerate(dataset_loader):
print('Calculating optical flow for {}-th frame'.format(idx + 1))
cur_img_addr = dataset.all_frame_addr[idx]
cur_img_name = cur_img_addr.split('/')[-1]
cur_img_name = cur_img_name.split('.')[0]
# parent path to store optical flow
of_path = of_root_dir
tmp_path_segment = cur_img_addr.split('/')[len_original_root_dir:-1]
for cur_seg in tmp_path_segment:
of_path = os.path.join(of_path, cur_seg)
if os.path.exists(of_path) is False:
os.makedirs(of_path, exist_ok=True)
# Calculate optical flow by FlowNet2
cur_imgs = np.transpose(batch[0].numpy(), [0, 2, 3, 1])
frameRange = dataset.context_range(idx)
old_size = (cur_imgs.shape[2], cur_imgs.shape[1])
if frameRange[1] == frameRange[0] or frameRange[1] == frameRange[2]:
if cur_imgs.shape[3] == 1:
im1 = cv2.resize(cur_imgs[0], (512, 384))[:, :, np.newaxis]
im2 = cv2.resize(cur_imgs[1], (512, 384))[:, :, np.newaxis]
im1 = np.concatenate([im1] * 3, axis=2)
im2 = np.concatenate([im2] * 3, axis=2)
else:
im1 = cv2.resize(cur_imgs[0], (512, 384))
im2 = cv2.resize(cur_imgs[1], (512, 384))
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)
pred_flow = flownet2(ims_v).cpu().data
pred_flow = pred_flow[0].numpy().transpose((1, 2, 0))
new_inputs = cv2.resize(pred_flow, old_size)
else:
if cur_imgs.shape[3] == 1:
im1 = cv2.resize(cur_imgs[1], (512, 384))[:, :, np.newaxis]
im2 = cv2.resize(cur_imgs[2], (512, 384))[:, :, np.newaxis]
im1 = np.concatenate([im1] * 3, axis=2)
im2 = np.concatenate([im2] * 3, axis=2)
else:
im1 = cv2.resize(cur_imgs[1], (512, 384))
im2 = cv2.resize(cur_imgs[2], (512, 384))
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)
pred_flow = flownet2(ims_v).cpu().data
pred_flow = pred_flow[0].numpy().transpose((1, 2, 0))
# Visualization
# cv2.imshow('of', flow_to_image(pred_flow))
# cv2.waitKey(0)
new_inputs = cv2.resize(pred_flow, old_size)
# Save optical flow
np.save(os.path.join(of_path, cur_img_name + '.npy'), new_inputs)