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(),