def flow_completion(args): data_list_dir = os.path.join(args.dataset_root, 'data') if not os.path.exists(data_list_dir): os.makedirs(data_list_dir) initial_data_list = os.path.join(data_list_dir, 'initial_test_list.txt') print('Generate datalist for initial step') from dataset.data_list import gen_flow_initial_test_mask_list gen_flow_initial_test_mask_list(flow_root=args.DATA_ROOT, output_txt_path=initial_data_list) args.EVAL_LIST = os.path.join(data_list_dir, 'initial_test_list.txt') from tools.test_scripts import test_initial_stage args.output_root = os.path.join(args.dataset_root, 'Flow_res', 'initial_res') args.PRETRAINED_MODEL = args.PRETRAINED_MODEL_1 if args.img_size is not None: args.IMAGE_SHAPE = [args.img_size[0] // 2, args.img_size[1] // 2] args.RES_SHAPE = args.IMAGE_SHAPE print('Flow Completion in First Step') test_initial_stage(args) args.flow_root = args.output_root if args.MS: args.ResNet101 = False from tools.test_scripts import test_refine_stage args.PRETRAINED_MODEL = args.PRETRAINED_MODEL_2 args.IMAGE_SHAPE = [320, 600] args.RES_SHAPE = [320, 600] args.DATA_ROOT = args.output_root args.output_root = os.path.join(args.dataset_root, 'Flow_res', 'stage2_res') stage2_data_list = os.path.join(data_list_dir, 'stage2_test_list.txt') from dataset.data_list import gen_flow_refine_test_mask_list gen_flow_refine_test_mask_list(flow_root=args.DATA_ROOT, output_txt_path=stage2_data_list) args.EVAL_LIST = stage2_data_list test_refine_stage(args) args.PRETRAINED_MODEL = args.PRETRAINED_MODEL_3 args.IMAGE_SHAPE = [480, 840] args.RES_SHAPE = [480, 840] args.DATA_ROOT = args.output_root args.output_root = os.path.join(args.dataset_root, 'Flow_res', 'stage3_res') stage3_data_list = os.path.join(data_list_dir, 'stage3_test_list.txt') from dataset.data_list import gen_flow_refine_test_mask_list gen_flow_refine_test_mask_list(flow_root=args.DATA_ROOT, output_txt_path=stage3_data_list) args.EVAL_LIST = stage3_data_list test_refine_stage(args) args.flow_root = args.output_root
def flow_completion(self): if self.i == -1: data_list_dir = os.path.join(self.args.dataset_root, 'data') os.makedirs(data_list_dir, exist_ok=True) initial_data_list = os.path.join(data_list_dir, 'initial_test_list.txt') print('Generate datalist for initial step') data_list.gen_flow_initial_test_mask_list( flow_root=self.args.DATA_ROOT, output_txt_path=initial_data_list) self.args.EVAL_LIST = os.path.join(data_list_dir, 'initial_test_list.txt') self.args.output_root = os.path.join(self.args.dataset_root, 'Flow_res', 'initial_res') self.args.PRETRAINED_MODEL = self.args.PRETRAINED_MODEL_1 if self.args.img_size is not None: self.args.IMAGE_SHAPE = [ self.args.img_size[0] // 2, self.args.img_size[1] // 2 ] self.args.RES_SHAPE = self.args.IMAGE_SHAPE print('Flow Completion in First Step') self.args.MASK_ROOT = self.args.mask_root eval_dataset = FlowInitial.FlowSeq(self.args, isTest=True) self.flow_refinement_dataloader = iter( DataLoader(eval_dataset, batch_size=self.settings.batch_size, shuffle=False, drop_last=False, num_workers=self.args.n_threads)) if self.args.ResNet101: dfc_resnet101 = resnet_models.Flow_Branch(33, 2) self.dfc_resnet = nn.DataParallel(dfc_resnet101).to( self.args.device) else: dfc_resnet50 = resnet_models.Flow_Branch_Multi( input_chanels=33, NoLabels=2) self.dfc_resnet = nn.DataParallel(dfc_resnet50).to( self.args.device) self.dfc_resnet.eval() io.load_ckpt(self.args.PRETRAINED_MODEL, [('model', self.dfc_resnet)], strict=True) print('Load Pretrained Model from', self.args.PRETRAINED_MODEL) self.i += 1 complete = False with torch.no_grad(): try: item = next(self.flow_refinement_dataloader) input_x = item[0].to(self.args.device) flow_masked = item[1].to(self.args.device) mask = item[3].to(self.args.device) output_dir = item[4][0] res_flow = self.dfc_resnet(input_x) res_complete = res_flow * mask[:, 10: 11, :, :] + flow_masked[:, 10:12, :, :] * ( 1. - mask[:, 10:11, :, :]) output_dir_split = output_dir.split(',') output_file = os.path.join(self.args.output_root, output_dir_split[0]) output_basedir = os.path.dirname(output_file) if not os.path.exists(output_basedir): os.makedirs(output_basedir) res_save = res_complete[0].permute( 1, 2, 0).contiguous().cpu().data.numpy() cvb.write_flow(res_save, output_file) except StopIteration: complete = True if self.i == len(self.flow_refinement_dataloader) - 1 or complete: self.args.flow_root = self.args.output_root del self.flow_refinement_dataloader, self.dfc_resnet self.i = -1 self.state += 1