JPEG Quantization for Y channel Args: rounding(function): rounding function to use
| 146 | |
| 147 | |
| 148 | class YQuantize(nn.Module): |
| 149 | """ JPEG Quantization for Y channel |
| 150 | |
| 151 | Args: |
| 152 | rounding(function): rounding function to use |
| 153 | """ |
| 154 | |
| 155 | def __init__(self, rounding): |
| 156 | super(YQuantize, self).__init__() |
| 157 | self.rounding = rounding |
| 158 | self.y_table = y_table |
| 159 | |
| 160 | def forward(self, image, factor=1): |
| 161 | """ |
| 162 | Args: |
| 163 | image(tensor): batch x height x width |
| 164 | |
| 165 | Returns: |
| 166 | Tensor: batch x height x width |
| 167 | """ |
| 168 | if isinstance(factor, (int, float)): |
| 169 | image = image.float() / (self.y_table * factor) |
| 170 | else: |
| 171 | b = factor.size(0) |
| 172 | table = self.y_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1) |
| 173 | image = image.float() / table |
| 174 | image = self.rounding(image) |
| 175 | return image |
| 176 | |
| 177 | |
| 178 | class CQuantize(nn.Module): |