def __init__(
self,
args_flow,
model_name="FlowNet2C",
restore_path="flow_inference/models/FlowNet2-C_checkpoint.pth.tar",
train_flow=False,
resume=None):
super(Inference_flow, self).__init__()
args_flow.rgb_max = 255.0
args_flow.fp16 = False
args_flow.grads = {}
if model_name == 'FlowNet2':
self.model = models.FlowNet2(args_flow).cuda()
elif model_name == 'FlowNet2C':
self.model = models.FlowNet2C(args_flow).cuda()
elif model_name == 'FlowNet2S':
self.model = models.FlowNet2S(args_flow).cuda()
elif model_name == 'FlowNet2SD':
self.model = models.FlowNet2SD(args_flow).cuda()
elif model_name == 'FlowNet2CS':
self.model = models.FlowNet2CS(args_flow).cuda()
elif model_name == 'FlowNet2CSS':
self.model = models.FlowNet2CSS(args_flow).cuda()
else:
assert False, "No such model %s" % (model_name)
print("loading %s pretrained model..." % (model_name))
if train_flow:
self.model.train()
else:
self.model.eval()
if resume is not None:
self.model.load_state_dict(torch.load(resume)['flow'])
else:
self.model.load_state_dict(torch.load(restore_path)['state_dict'])
prototxt = tempfile.NamedTemporaryFile(mode='w', delete=True)
prototxt.write(template)
prototxt.flush()
net = caffe.Net(prototxt.name, args.caffe_model, caffe.TEST)
weights = {}
biases = {}
for k, v in list(net.params.items()):
weights[k] = np.array(v[0].data).reshape(v[0].data.shape)
biases[k] = np.array(v[1].data).reshape(v[1].data.shape)
print((k, weights[k].shape, biases[k].shape))
if 'FlowNet2/' in args.caffe_model:
model = models.FlowNet2(args)
parse_flownetc(model.flownetc.modules(), weights, biases)
parse_flownets(model.flownets_1.modules(),
weights,
biases,
param_prefix='net2_')
parse_flownets(model.flownets_2.modules(),
weights,
biases,
param_prefix='net3_')
parse_flownetsd(model.flownets_d.modules(),
weights,
biases,
param_prefix='netsd_')
parse_flownetfusion(model.flownetfusion.modules(),
def main():
global args
args = parser.parse_args()
test_list = make_dataset(args.data)
# test_list = make_real_dataset(args.data)
if args.arch == 'pwc':
model = models.pwc_dc_net('models/pwc_net_ft.pth.tar').cuda()
elif args.arch == 'spynet':
model = models.spynet(nlevels=5, strmodel='F').cuda()
elif args.arch == 'flownet2':
model = models.FlowNet2().cuda()
print("=> using pre-trained weights for FlowNet2")
weights = torch.load('models/FlowNet2_checkpoint.pth.tar')
model.load_state_dict(weights['state_dict'])
if args.pretrained is not None:
network_data = torch.load(args.pretrained)
args.arch = network_data['arch']
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](data=network_data).cuda()
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
model.eval()
flow_epe = AverageMeter()
avg_mot_err = AverageMeter()
avg_parts_epe = {}
for bk in BODY_MAP.keys():
avg_parts_epe[bk] = AverageMeter()
if args.no_norm:
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255])
])
else:
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0], std=[args.div_flow, args.div_flow])
])
for i, (img_paths, flow_path, seg_path) in enumerate(tqdm(test_list)):
raw_im1 = flow_transforms.ArrayToTensor()(imread(img_paths[0],
mode='RGB')[:, :, :3])
raw_im2 = flow_transforms.ArrayToTensor()(imread(img_paths[1],
mode='RGB')[:, :, :3])
img1 = input_transform(imread(img_paths[0], mode='RGB')[:, :, :3])
img2 = input_transform(imread(img_paths[1], mode='RGB')[:, :, :3])
if flow_path is None:
_, h, w = img1.size()
new_h = int(np.floor(h / 256) * 256)
new_w = int(np.floor(w / 448) * 448)
# if i>744:
# import ipdb; ipdb.set_trace()
img1 = F.upsample(img1.unsqueeze(0), (new_h, new_w),
mode='bilinear').squeeze()
img2 = F.upsample(img2.unsqueeze(0), (new_h, new_w),
mode='bilinear').squeeze()
if flow_path is not None:
gtflow = target_transform(load_flo(flow_path))
segmask = flow_transforms.ArrayToTensor()(cv2.imread(seg_path))
input_var = torch.cat([img1, img2]).unsqueeze(0)
if flow_path is not None:
gtflow_var = gtflow.unsqueeze(0)
segmask_var = segmask.unsqueeze(0)
input_var = input_var.to(device)
if flow_path is not None:
gtflow_var = gtflow_var.to(device)
segmask_var = segmask_var.to(device)
# compute output
output = model(input_var)
if flow_path is not None:
epe = args.div_flow * realEPE(
output,
gtflow_var,
sparse=True if 'KITTI' in args.dataset else False)
epe_parts = partsEPE(output, gtflow_var, segmask_var)
epe_parts.update(
(x, args.div_flow * y) for x, y in epe_parts.items())
# record EPE
flow_epe.update(epe.item(), gtflow_var.size(0))
for bk in avg_parts_epe:
if epe_parts[bk].item() > 0:
avg_parts_epe[bk].update(epe_parts[bk].item(),
gtflow_var.size(0))
# record motion warping error
raw_im1 = raw_im1.cuda().unsqueeze(0)
raw_im2 = raw_im2.cuda().unsqueeze(0)
mot_err = motion_warping_error(raw_im1, raw_im2,
args.div_flow * output)
avg_mot_err.update(mot_err.item(), raw_im1.size(0))
if args.output_dir is not None:
if flow_path is not None:
_, h, w = gtflow.size()
output_path = flow_path.replace(args.data, args.output_dir)
output_path = output_path.replace('/test/', '/')
os.system('mkdir -p ' + output_path[:-15])
else:
output_path = img_paths[0].replace(args.data, args.output_dir)
os.system('mkdir -p ' + output_path[:-10])
output_path = output_path.replace('.png', '.flo')
output_path = output_path.replace('/flow/', '/')
upsampled_output = F.interpolate(output, (h, w),
mode='bilinear',
align_corners=False)
flow_write(output_path,
upsampled_output.cpu()[0].data.numpy()[0],
upsampled_output.cpu()[0].data.numpy()[1])
if args.save_name is not None:
epe_dict = {}
for bk in BODY_MAP.keys():
epe_dict[bk] = avg_parts_epe[bk].avg
epe_dict['full_epe'] = flow_epe.avg
np.save(os.path.join('results', args.save_name), epe_dict)
print("Averge EPE", flow_epe.avg)
print("Motion warping error", avg_mot_err.avg)
def load_image(path: str) -> np.ndarray:
img = cv2.imread(path, cv2.IMREAD_COLOR)
# OpenCV stores images as BGR, but the model expects RGB.
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# Convert to float and add batch dimension.
return np.expand_dims(img.astype(np.float32) / 255., axis=0)
if __name__ == "__main__":
# Graph construction.
image1_placeholder = tf.placeholder(tf.float32, [1, 384, 512, 3])
image2_placeholder = tf.placeholder(tf.float32, [1, 384, 512, 3])
flownet2 = models.FlowNet2()
inputs = {"input_a": image1_placeholder, "input_b": image2_placeholder}
predicted_flow = flownet2.run(inputs)["flow"]
# Load inputs.
image1 = load_image("example/0img0.ppm")
image2 = load_image("example/0img1.ppm")
ckpt_file = "checkpoints/FlowNet2/flownet-2.ckpt-0"
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, ckpt_file)
feed_dict = {image1_placeholder: image1, image2_placeholder: image2}
predicted_flow_np = sess.run(predicted_flow, feed_dict=feed_dict)
######################################################################################
######################################################################################
args = parser.parse_args()
args.model_class = tools.module_to_dict(models)[args.model]
args.cuda = not args.no_cuda and torch.cuda.is_available()
assert args.flow_output_type in ('jpg', 'flo')
######################################################################################
######################################################################################
# Intialize the model
logging.info('Initialized FlowNet2 model.')
model = models.FlowNet2(args, batchNorm=False, div_flow=20.)
if args.cuda:
model = model.cuda()
model.eval()
######################################################################################
######################################################################################
# Restore weights from checkpoint
if os.path.isfile(args.checkpoint):
checkpoint = torch.load(args.checkpoint)
model.load_state_dict(checkpoint['state_dict'])
logging.info('Restored checkpoint file from: {}'.format(
args.checkpoint))
else:
raise FileNotFoundError('no existing checkpoint found...')
net = caffe.Net(prototxt.name, opt.caffe_model, caffe.TEST)
weights = {}
biases = {}
for k, v in net.params.items():
weights[k] = np.array(v[0].data).reshape(v[0].data.shape)
if len(v) > 1:
# the last upsampling layer in PWC-net does not have bias term
biases[k] = np.array(v[1].data).reshape(v[1].data.shape)
print(k, weights[k].shape, biases[k].shape)
else:
print(k, weights[k].shape)
if 'FlowNet2/' in opt.caffe_model:
model = models.FlowNet2(opt)
parse_flownetc(model.flownetc.modules(), weights, biases)
parse_flownets(model.flownets_1.modules(),
weights,
biases,
param_prefix='net2_')
parse_flownets(model.flownets_2.modules(),
weights,
biases,
param_prefix='net3_')
parse_flownetsd(model.flownets_d.modules(),
weights,
biases,
param_prefix='netsd_')
parse_flownetfusion(model.flownetfusion.modules(),
