| 385 | return flownets1_flow |
| 386 | |
| 387 | class FlowNet2CSS(nn.Module): |
| 388 | |
| 389 | def __init__(self, args, batchNorm=False, div_flow = 20.): |
| 390 | super(FlowNet2CSS,self).__init__() |
| 391 | self.batchNorm = batchNorm |
| 392 | self.div_flow = div_flow |
| 393 | self.rgb_max = args.rgb_max |
| 394 | self.args = args |
| 395 | |
| 396 | self.channelnorm = ChannelNorm() |
| 397 | |
| 398 | # First Block (FlowNetC) |
| 399 | self.flownetc = FlowNetC.FlowNetC(args, batchNorm=self.batchNorm) |
| 400 | self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear') |
| 401 | self.resample1 = Resample2d() if not args.fp16 else fp16_resample2d() |
| 402 | |
| 403 | # Block (FlowNetS1) |
| 404 | self.flownets_1 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm) |
| 405 | self.upsample2 = nn.Upsample(scale_factor=4, mode='bilinear') |
| 406 | self.resample2 = Resample2d() if not args.fp16 else fp16_resample2d() |
| 407 | |
| 408 | # Block (FlowNetS2) |
| 409 | self.flownets_2 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm) |
| 410 | self.upsample3 = nn.Upsample(scale_factor=4, mode='nearest') |
| 411 | |
| 412 | for m in self.modules(): |
| 413 | if isinstance(m, nn.Conv2d): |
| 414 | if m.bias is not None: |
| 415 | init.uniform(m.bias) |
| 416 | init.xavier_uniform(m.weight) |
| 417 | |
| 418 | if isinstance(m, nn.ConvTranspose2d): |
| 419 | if m.bias is not None: |
| 420 | init.uniform(m.bias) |
| 421 | init.xavier_uniform(m.weight) |
| 422 | # init_deconv_bilinear(m.weight) |
| 423 | |
| 424 | def forward(self, inputs): |
| 425 | rgb_mean = inputs.contiguous().view(inputs.size()[:2]+(-1,)).mean(dim=-1).view(inputs.size()[:2] + (1,1,1,)) |
| 426 | |
| 427 | x = (inputs - rgb_mean) / self.rgb_max |
| 428 | x1 = x[:,:,0,:,:] |
| 429 | x2 = x[:,:,1,:,:] |
| 430 | x = torch.cat((x1,x2), dim = 1) |
| 431 | |
| 432 | # flownetc |
| 433 | flownetc_flow2 = self.flownetc(x)[0] |
| 434 | flownetc_flow = self.upsample1(flownetc_flow2*self.div_flow) |
| 435 | |
| 436 | # warp img1 to img0; magnitude of diff between img0 and and warped_img1, |
| 437 | resampled_img1 = self.resample1(x[:,3:,:,:], flownetc_flow) |
| 438 | diff_img0 = x[:,:3,:,:] - resampled_img1 |
| 439 | norm_diff_img0 = self.channelnorm(diff_img0) |
| 440 | |
| 441 | # concat img0, img1, img1->img0, flow, diff-mag ; |
| 442 | concat1 = torch.cat((x, resampled_img1, flownetc_flow/self.div_flow, norm_diff_img0), dim=1) |
| 443 | |
| 444 | # flownets1 |
nothing calls this directly
no outgoing calls
no test coverage detected