| 3 | |
| 4 | |
| 5 | class SimpleAdapter(nn.Module): |
| 6 | def __init__(self, in_dim, out_dim, kernel_size, stride, downscale_factor=8, num_residual_blocks=1): |
| 7 | super(SimpleAdapter, self).__init__() |
| 8 | |
| 9 | # Pixel Unshuffle: reduce spatial dimensions by a factor of 8 |
| 10 | self.pixel_unshuffle = nn.PixelUnshuffle(downscale_factor=downscale_factor) |
| 11 | |
| 12 | # Convolution: reduce spatial dimensions by a factor |
| 13 | # of 2 (without overlap) |
| 14 | self.conv = nn.Conv2d(in_dim * downscale_factor * downscale_factor, out_dim, kernel_size=kernel_size, stride=stride, padding=0) |
| 15 | |
| 16 | # Residual blocks for feature extraction |
| 17 | self.residual_blocks = nn.Sequential( |
| 18 | *[ResidualBlock(out_dim) for _ in range(num_residual_blocks)] |
| 19 | ) |
| 20 | |
| 21 | def forward(self, x): |
| 22 | # Reshape to merge the frame dimension into batch |
| 23 | bs, c, f, h, w = x.size() |
| 24 | x = x.permute(0, 2, 1, 3, 4).contiguous().view(bs * f, c, h, w) |
| 25 | |
| 26 | # Pixel Unshuffle operation |
| 27 | x_unshuffled = self.pixel_unshuffle(x) |
| 28 | |
| 29 | # Convolution operation |
| 30 | x_conv = self.conv(x_unshuffled) |
| 31 | |
| 32 | # Feature extraction with residual blocks |
| 33 | out = self.residual_blocks(x_conv) |
| 34 | |
| 35 | # Reshape to restore original bf dimension |
| 36 | out = out.view(bs, f, out.size(1), out.size(2), out.size(3)) |
| 37 | |
| 38 | # Permute dimensions to reorder (if needed), e.g., swap channels and feature frames |
| 39 | out = out.permute(0, 2, 1, 3, 4) |
| 40 | |
| 41 | return out |
| 42 | |
| 43 | |
| 44 | class ResidualBlock(nn.Module): |