Full JPEG compression algorithm Args: rounding(function): rounding function to use
| 206 | |
| 207 | |
| 208 | class CompressJpeg(nn.Module): |
| 209 | """Full JPEG compression algorithm |
| 210 | |
| 211 | Args: |
| 212 | rounding(function): rounding function to use |
| 213 | """ |
| 214 | |
| 215 | def __init__(self, rounding=torch.round): |
| 216 | super(CompressJpeg, self).__init__() |
| 217 | self.l1 = nn.Sequential(RGB2YCbCrJpeg(), ChromaSubsampling()) |
| 218 | self.l2 = nn.Sequential(BlockSplitting(), DCT8x8()) |
| 219 | self.c_quantize = CQuantize(rounding=rounding) |
| 220 | self.y_quantize = YQuantize(rounding=rounding) |
| 221 | |
| 222 | def forward(self, image, factor=1): |
| 223 | """ |
| 224 | Args: |
| 225 | image(tensor): batch x 3 x height x width |
| 226 | |
| 227 | Returns: |
| 228 | dict(tensor): Compressed tensor with batch x h*w/64 x 8 x 8. |
| 229 | """ |
| 230 | y, cb, cr = self.l1(image * 255) |
| 231 | components = {'y': y, 'cb': cb, 'cr': cr} |
| 232 | for k in components.keys(): |
| 233 | comp = self.l2(components[k]) |
| 234 | if k in ('cb', 'cr'): |
| 235 | comp = self.c_quantize(comp, factor=factor) |
| 236 | else: |
| 237 | comp = self.y_quantize(comp, factor=factor) |
| 238 | |
| 239 | components[k] = comp |
| 240 | |
| 241 | return components['y'], components['cb'], components['cr'] |
| 242 | |
| 243 | |
| 244 | # ------------------------ decompression ------------------------# |