Full JPEG decompression algorithm Args: rounding(function): rounding function to use
| 399 | |
| 400 | |
| 401 | class DeCompressJpeg(nn.Module): |
| 402 | """Full JPEG decompression algorithm |
| 403 | |
| 404 | Args: |
| 405 | rounding(function): rounding function to use |
| 406 | """ |
| 407 | |
| 408 | def __init__(self, rounding=torch.round): |
| 409 | super(DeCompressJpeg, self).__init__() |
| 410 | self.c_dequantize = CDequantize() |
| 411 | self.y_dequantize = YDequantize() |
| 412 | self.idct = iDCT8x8() |
| 413 | self.merging = BlockMerging() |
| 414 | self.chroma = ChromaUpsampling() |
| 415 | self.colors = YCbCr2RGBJpeg() |
| 416 | |
| 417 | def forward(self, y, cb, cr, imgh, imgw, factor=1): |
| 418 | """ |
| 419 | Args: |
| 420 | compressed(dict(tensor)): batch x h*w/64 x 8 x 8 |
| 421 | imgh(int) |
| 422 | imgw(int) |
| 423 | factor(float) |
| 424 | |
| 425 | Returns: |
| 426 | Tensor: batch x 3 x height x width |
| 427 | """ |
| 428 | components = {'y': y, 'cb': cb, 'cr': cr} |
| 429 | for k in components.keys(): |
| 430 | if k in ('cb', 'cr'): |
| 431 | comp = self.c_dequantize(components[k], factor=factor) |
| 432 | height, width = int(imgh / 2), int(imgw / 2) |
| 433 | else: |
| 434 | comp = self.y_dequantize(components[k], factor=factor) |
| 435 | height, width = imgh, imgw |
| 436 | comp = self.idct(comp) |
| 437 | components[k] = self.merging(comp, height, width) |
| 438 | # |
| 439 | image = self.chroma(components['y'], components['cb'], components['cr']) |
| 440 | image = self.colors(image) |
| 441 | |
| 442 | image = torch.min(255 * torch.ones_like(image), torch.max(torch.zeros_like(image), image)) |
| 443 | return image / 255 |
| 444 | |
| 445 | |
| 446 | # ------------------------ main DiffJPEG ------------------------ # |