def __init__(self, args, batchNorm=False, div_flow = 20.):
super(FlowNet2CS,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')
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')
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
### start training
while model.epoch < opts.epoch_max:
model.epoch += 1
### re-generate train data loader for every epoch
data_loader = utils.create_data_loader(train_dataset, opts, "train")
### update learning rate
current_lr = utils.learning_rate_decay(opts, model.epoch)
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
## submodule
flow_warping = Resample2d().to(device)
downsampler = nn.AvgPool2d((2, 2), stride=2).to(device)
### criterion and loss recorder
if opts.loss == 'L2':
criterion = nn.MSELoss(size_average=True)
elif opts.loss == 'L1':
criterion = nn.L1Loss(size_average=True)
else:
raise Exception("Unsupported criterion %s" %opts.loss)
### start epoch
ts = datetime.now()
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 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
model.train()
ts = opt.t_stride
### start epoch
for i, (inputs, masks, _) in enumerate(data_loader):
# inputs: BxCxTxHxW
def __init__(self, args, batchNorm=False, div_flow = 20., requires_grad=False):
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 = networks.FlowNetC.FlowNetC(args, batchNorm=self.batchNorm)
self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear')
if args.fp16:
self.resample1 = nn.Sequential(
tofp32(),
Resample2d(),
tofp16())
else:
self.resample1 = Resample2d()
# Block (FlowNetS1)
self.flownets_1 = networks.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 = networks.FlowNetS.FlowNetS(args, batchNorm=self.batchNorm)
# Block (FlowNetSD)
self.flownets_d = networks.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 = networks.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)
# init_deconv_bilinear(m.weight)
if not requires_grad:
for param in self.parameters():
param.requires_grad = False
def __init__(self):
super(FlowWarping, self).__init__()
self.channelnorm = ChannelNorm()
self.resample1 = Resample2d()
