def __init__(self, args): super(DummyModel, self).__init__() self.args = args self.flownet = FlowNet2S(args) checkpoint_file = "./checkpoints/FlowNet2-S_checkpoint.pth.tar" checkpoint = torch.load(checkpoint_file) self.flownet.load_state_dict(checkpoint['state_dict']) # self.flownet.training = False # for param in self.flownet.parameters(): # param.requires_grad = False if args.fp16: self.resample1 = nn.Sequential(tofp32(), Resample2d(), tofp16()) else: self.resample1 = Resample2d() self.rgb_max = 255
def __init__(self, args, batchNorm=False, div_flow=20.): super(FlowNet2CSS, self).__init__() self.batchNorm = batchNorm self.div_flow = div_flow self.rgb_max = args.rgb_max self.args = args self.channelnorm = ChannelNorm() # First Block (FlowNetC) self.flownetc = FlowNetC.FlowNetC(args, batchNorm=self.batchNorm) self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear', align_corners=True) if args.fp16: self.resample1 = nn.Sequential(tofp32(), Resample2d(), tofp16()) else: self.resample1 = Resample2d() # Block (FlowNetS1) self.flownets_1 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm) self.upsample2 = nn.Upsample(scale_factor=4, mode='bilinear', align_corners=True) if args.fp16: self.resample2 = nn.Sequential(tofp32(), Resample2d(), tofp16()) else: self.resample2 = Resample2d() # Block (FlowNetS2) self.flownets_2 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm) self.upsample3 = nn.Upsample(scale_factor=4, mode='nearest') for m in self.modules(): if isinstance(m, nn.Conv2d): if m.bias is not None: init.uniform(m.bias) init.xavier_uniform(m.weight) if isinstance(m, nn.ConvTranspose2d): if m.bias is not None: init.uniform(m.bias) init.xavier_uniform(m.weight)
def detect_occlusion(fw_flow, bw_flow): ## fw-flow: img1 => img2 ## bw-flow: img2 => img1 with torch.no_grad(): ## convert to tensor fw_flow_t = img2tensor(fw_flow).cuda() bw_flow_t = img2tensor(bw_flow).cuda() ## warp fw-flow to img2 flow_warping = Resample2d().cuda() fw_flow_w = flow_warping(fw_flow_t, bw_flow_t) ## convert to numpy array fw_flow_w = tensor2img(fw_flow_w) ## occlusion fb_flow_sum = fw_flow_w + bw_flow fb_flow_mag = compute_flow_magnitude(fb_flow_sum) fw_flow_w_mag = compute_flow_magnitude(fw_flow_w) bw_flow_mag = compute_flow_magnitude(bw_flow) mask1 = fb_flow_mag > 0.01 * (fw_flow_w_mag + bw_flow_mag) + 0.5 ## motion boundary fx_du, fx_dv, fy_du, fy_dv = compute_flow_gradients(bw_flow) fx_mag = fx_du ** 2 + fx_dv ** 2 fy_mag = fy_du ** 2 + fy_dv ** 2 mask2 = (fx_mag + fy_mag) > 0.01 * bw_flow_mag + 0.002 ## combine mask mask = np.logical_or(mask1, mask2) occlusion = np.zeros((fw_flow.shape[0], fw_flow.shape[1])) occlusion[mask == 1] = 1 return occlusion
output_dir = os.path.join(opts.data_dir, opts.phase, opts.method, opts.task, opts.dataset) ## print average if result already exists metric_filename = os.path.join(output_dir, "WarpError.txt") if os.path.exists(metric_filename) and not opts.redo: print("Output %s exists, skip..." % metric_filename) cmd = 'tail -n1 %s' % metric_filename utils.run_cmd(cmd) sys.exit() ## flow warping layer device = torch.device("cuda" if opts.cuda else "cpu") flow_warping = Resample2d().to(device) ### load video list list_filename = os.path.join(opts.list_dir, "%s_%s.txt" % (opts.dataset, opts.phase)) with open(list_filename) as f: video_list = [line.rstrip() for line in f.readlines()] ### start evaluation err_all = np.zeros(len(video_list)) for v in range(len(video_list)): video = video_list[v] frame_dir = os.path.join(opts.data_dir, opts.phase, opts.method,
def __init__(self, args, batchNorm=False, div_flow = 20.): super(FlowNet2,self).__init__() self.batchNorm = batchNorm self.div_flow = div_flow self.rgb_max = args.rgb_max self.args = args self.channelnorm = ChannelNorm() # First Block (FlowNetC) self.flownetc = FlowNetC.FlowNetC(args, batchNorm=self.batchNorm) self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear') # import ipdb; ipdb.set_trace() if args.fp16: self.resample1 = nn.Sequential( tofp32(), Resample2d(), tofp16()) else: self.resample1 = Resample2d() # Block (FlowNetS1) self.flownets_1 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm) self.upsample2 = nn.Upsample(scale_factor=4, mode='bilinear') if args.fp16: self.resample2 = nn.Sequential( tofp32(), Resample2d(), tofp16()) else: self.resample2 = Resample2d() # Block (FlowNetS2) self.flownets_2 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm) # Block (FlowNetSD) self.flownets_d = FlowNetSD.FlowNetSD(args, batchNorm=self.batchNorm) self.upsample3 = nn.Upsample(scale_factor=4, mode='nearest') self.upsample4 = nn.Upsample(scale_factor=4, mode='nearest') if args.fp16: self.resample3 = nn.Sequential( tofp32(), Resample2d(), tofp16()) else: self.resample3 = Resample2d() if args.fp16: self.resample4 = nn.Sequential( tofp32(), Resample2d(), tofp16()) else: self.resample4 = Resample2d() # Block (FLowNetFusion) self.flownetfusion = FlowNetFusion.FlowNetFusion(args, batchNorm=self.batchNorm) for m in self.modules(): if isinstance(m, nn.Conv2d): if m.bias is not None: init.uniform_(m.bias) init.xavier_uniform_(m.weight) if isinstance(m, nn.ConvTranspose2d): if m.bias is not None: init.uniform_(m.bias) init.xavier_uniform_(m.weight)
from networks.resample2d_package.resample2d import Resample2d import networks class Object(object): pass """ Flownet """ args = Object() args.rgb_max = 1.0 args.fp16 = False FlowNet = networks.FlowNet2(args, requires_grad=False) model_filename = os.path.join("pretrained_models", "FlowNet2_checkpoint.pth.tar") checkpoint = torch.load(model_filename) FlowNet.load_state_dict(checkpoint['state_dict']) FlowNet = FlowNet.cuda() """ Submodules """ flow_warping = Resample2d().cuda() downsampler = nn.AvgPool2d((2, 2), stride=2).cuda() def norm(t): return torch.sum(t*t, dim=1, keepdim=True) def repackage_hidden(h): """Wraps hidden states in new Variables, to detach them from their history.""" if isinstance(h, torch.Tensor): return h.detach() else: return tuple(repackage_hidden(v) for v in h) def train_lstm_epoch(epoch, data_loader, model, criterion_L1, criterion_ssim, optimizer, opt): opt.w_ST, opt.w_LT, opt.w_Flow = 1.0, 1.0, 10.0
def __init__(self, args): super(InterpolNetOld, self).__init__() self.args = args # flownet predictor self.flownet = FlowNet2(args) checkpoint = torch.load("./checkpoints/FlowNet2_checkpoint.pth.tar") self.flownet.load_state_dict(checkpoint['state_dict']) self.flownet.training = False for param in self.flownet.parameters(): param.requires_grad = False if args.fp16: self.resample1 = nn.Sequential(tofp32(), Resample2d(), tofp16()) else: self.resample1 = Resample2d() res_kernel_number = 128 self.convResIn_img1 = nn.Conv2d(3, res_kernel_number, kernel_size=3, stride=1, padding=1) self.resBlock_img1 = BasicResBlock(res_kernel_number) self.convResOut_img1 = nn.Conv2d(res_kernel_number, 3, kernel_size=3, stride=1, padding=1) self.convResIn_img2 = nn.Conv2d(3, res_kernel_number, kernel_size=3, stride=1, padding=1) self.resBlock_img2 = BasicResBlock(res_kernel_number) self.convResOut_img2 = nn.Conv2d(res_kernel_number, 3, kernel_size=3, stride=1, padding=1) self.convResIn_mid = nn.Conv2d(3, res_kernel_number, kernel_size=3, stride=1, padding=1) self.resBlock_mid = BasicResBlock(res_kernel_number) self.convResOut_mid = nn.Conv2d(res_kernel_number, 3, kernel_size=3, stride=1, padding=1) self.convResIn_final = nn.Conv2d(9, res_kernel_number, kernel_size=3, stride=1, padding=1) self.resBlock_final = BasicResBlock(res_kernel_number) self.convResOut_final = nn.Conv2d(res_kernel_number, 3, kernel_size=3, stride=1, padding=1) self.relu = nn.ReLU() self.rgb_max = 255