Discrete Cosine Transformation
| 119 | |
| 120 | |
| 121 | class DCT8x8(nn.Module): |
| 122 | """ Discrete Cosine Transformation |
| 123 | """ |
| 124 | |
| 125 | def __init__(self): |
| 126 | super(DCT8x8, self).__init__() |
| 127 | tensor = np.zeros((8, 8, 8, 8), dtype=np.float32) |
| 128 | for x, y, u, v in itertools.product(range(8), repeat=4): |
| 129 | tensor[x, y, u, v] = np.cos((2 * x + 1) * u * np.pi / 16) * np.cos((2 * y + 1) * v * np.pi / 16) |
| 130 | alpha = np.array([1. / np.sqrt(2)] + [1] * 7) |
| 131 | self.tensor = nn.Parameter(torch.from_numpy(tensor).float()) |
| 132 | self.scale = nn.Parameter(torch.from_numpy(np.outer(alpha, alpha) * 0.25).float()) |
| 133 | |
| 134 | def forward(self, image): |
| 135 | """ |
| 136 | Args: |
| 137 | image(tensor): batch x height x width |
| 138 | |
| 139 | Returns: |
| 140 | Tensor: batch x height x width |
| 141 | """ |
| 142 | image = image - 128 |
| 143 | result = self.scale * torch.tensordot(image, self.tensor, dims=2) |
| 144 | result.view(image.shape) |
| 145 | return result |
| 146 | |
| 147 | |
| 148 | class YQuantize(nn.Module): |