| 55 | |
| 56 | |
| 57 | class Upsample3D(torch.nn.Module): |
| 58 | def __init__( |
| 59 | self, |
| 60 | in_channels: int, |
| 61 | out_channels: int, |
| 62 | kernel_size: int = 3, |
| 63 | stride: int = 1, |
| 64 | padding: int = 1, |
| 65 | compress_time: bool = False, |
| 66 | ) -> None: |
| 67 | super().__init__() |
| 68 | self.conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding) |
| 69 | self.compress_time = compress_time |
| 70 | |
| 71 | def forward(self, inputs: torch.Tensor, xq: torch.Tensor) -> torch.Tensor: |
| 72 | if self.compress_time: |
| 73 | if inputs.shape[2] > 1 and inputs.shape[2] % 2 == 1: |
| 74 | # split first frame |
| 75 | x_first, x_rest = inputs[:, :, 0], inputs[:, :, 1:] |
| 76 | |
| 77 | x_first = torch.nn.functional.interpolate(x_first, scale_factor=2.0) |
| 78 | x_rest = torch.nn.functional.interpolate(x_rest, scale_factor=2.0) |
| 79 | x_first = x_first[:, :, None, :, :] |
| 80 | inputs = torch.cat([x_first, x_rest], dim=2) |
| 81 | elif inputs.shape[2] > 1: |
| 82 | inputs = torch.nn.functional.interpolate(inputs, scale_factor=2.0) |
| 83 | else: |
| 84 | inputs = inputs.squeeze(2) |
| 85 | inputs = torch.nn.functional.interpolate(inputs, scale_factor=2.0) |
| 86 | inputs = inputs[:, :, None, :, :] |
| 87 | else: |
| 88 | # only interpolate 2D |
| 89 | b, c, t, h, w = inputs.shape |
| 90 | inputs = inputs.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) |
| 91 | inputs = torch.nn.functional.interpolate(inputs, scale_factor=2.0) |
| 92 | inputs = inputs.reshape(b, t, c, *inputs.shape[2:]).permute(0, 2, 1, 3, 4) |
| 93 | |
| 94 | b, c, t, h, w = inputs.shape |
| 95 | inputs = inputs.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) |
| 96 | inputs = self.conv(inputs) |
| 97 | inputs = inputs.reshape(b, t, *inputs.shape[1:]).permute(0, 2, 1, 3, 4) |
| 98 | |
| 99 | return inputs |
| 100 | |
| 101 | |
| 102 | |