(self, num_input_frame, num_feat, load_path)
| 278 | """ |
| 279 | |
| 280 | def __init__(self, num_input_frame, num_feat, load_path): |
| 281 | |
| 282 | super(EDVRFeatureExtractor, self).__init__() |
| 283 | |
| 284 | self.center_frame_idx = num_input_frame // 2 |
| 285 | |
| 286 | # extract pyramid features |
| 287 | self.conv_first = nn.Conv2d(3, num_feat, 3, 1, 1) |
| 288 | self.feature_extraction = make_layer(ResidualBlockNoBN, 5, num_feat=num_feat) |
| 289 | self.conv_l2_1 = nn.Conv2d(num_feat, num_feat, 3, 2, 1) |
| 290 | self.conv_l2_2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) |
| 291 | self.conv_l3_1 = nn.Conv2d(num_feat, num_feat, 3, 2, 1) |
| 292 | self.conv_l3_2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) |
| 293 | |
| 294 | # pcd and tsa module |
| 295 | self.pcd_align = PCDAlignment(num_feat=num_feat, deformable_groups=8) |
| 296 | self.fusion = TSAFusion(num_feat=num_feat, num_frame=num_input_frame, center_frame_idx=self.center_frame_idx) |
| 297 | |
| 298 | # activation function |
| 299 | self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True) |
| 300 | |
| 301 | if load_path: |
| 302 | self.load_state_dict(torch.load(load_path, map_location=lambda storage, loc: storage)['params']) |
| 303 | |
| 304 | def forward(self, x): |
| 305 | b, n, c, h, w = x.size() |
nothing calls this directly
no test coverage detected