(self, opt)
| 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(): |
no test coverage detected