| 4 | from .base_model import BaseModel |
| 5 | |
| 6 | class FlowNet(BaseModel): |
| 7 | def name(self): |
| 8 | return 'FlowNet' |
| 9 | |
| 10 | def initialize(self, opt): |
| 11 | BaseModel.initialize(self, opt) |
| 12 | |
| 13 | # flownet 2 |
| 14 | from .flownet2_pytorch import models as flownet2_models |
| 15 | from .flownet2_pytorch.utils import tools as flownet2_tools |
| 16 | from .flownet2_pytorch.networks.resample2d_package.resample2d import Resample2d |
| 17 | |
| 18 | self.flowNet = flownet2_tools.module_to_dict(flownet2_models)['FlowNet2'](fp16=opt.fp16).cuda(self.gpu_ids[0]) |
| 19 | checkpoint = torch.load('models/flownet2_pytorch/FlowNet2_checkpoint.pth.tar') |
| 20 | self.flowNet.load_state_dict(checkpoint['state_dict']) |
| 21 | self.flowNet.eval() |
| 22 | self.resample = Resample2d() |
| 23 | self.downsample = torch.nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False) |
| 24 | |
| 25 | def forward(self, input_A, input_B, dummy_bs=0): |
| 26 | with torch.no_grad(): |
| 27 | if input_A.get_device() == self.gpu_ids[0]: |
| 28 | input_A, input_B = input_A[dummy_bs:], input_B[dummy_bs:] |
| 29 | if input_A.size(0) == 0: |
| 30 | b, n, c, h, w = input_A.size() |
| 31 | return self.Tensor(1, n, 2, h, w), self.Tensor(1, n, 1, h, w) |
| 32 | size = input_A.size() |
| 33 | assert(len(size) == 4 or len(size) == 5) |
| 34 | if len(size) == 5: |
| 35 | b, n, c, h, w = size |
| 36 | input_A = input_A.contiguous().view(-1, c, h, w) |
| 37 | input_B = input_B.contiguous().view(-1, c, h, w) |
| 38 | flow, conf = self.compute_flow_and_conf(input_A, input_B) |
| 39 | return flow.view(b, n, 2, h, w), conf.view(b, n, 1, h, w) |
| 40 | else: |
| 41 | return self.compute_flow_and_conf(input_A, input_B) |
| 42 | |
| 43 | def compute_flow_and_conf(self, im1, im2): |
| 44 | assert(im1.size()[1] == 3) |
| 45 | assert(im1.size() == im2.size()) |
| 46 | old_h, old_w = im1.size()[2], im1.size()[3] |
| 47 | new_h, new_w = old_h//64*64, old_w//64*64 |
| 48 | if old_h != new_h: |
| 49 | downsample = torch.nn.Upsample(size=(new_h, new_w), mode='bilinear') |
| 50 | upsample = torch.nn.Upsample(size=(old_h, old_w), mode='bilinear') |
| 51 | im1 = downsample(im1) |
| 52 | im2 = downsample(im2) |
| 53 | data1 = torch.cat([im1.unsqueeze(2), im2.unsqueeze(2)], dim=2) |
| 54 | flow1 = self.flowNet(data1) |
| 55 | conf = (self.norm(im1 - self.resample(im2, flow1)) < 0.02).float() |
| 56 | if old_h != new_h: |
| 57 | flow1 = upsample(flow1) * old_h / new_h |
| 58 | conf = upsample(conf) |
| 59 | return flow1.detach(), conf.detach() |
| 60 | |
| 61 | def norm(self, t): |
| 62 | return torch.sum(t*t, dim=1, keepdim=True) |