This JPEG algorithm result is slightly different from cv2. DiffJPEG supports batch processing. Args: differentiable(bool): If True, uses custom differentiable rounding function, if False, uses standard torch.round
| 447 | |
| 448 | |
| 449 | class DiffJPEG(nn.Module): |
| 450 | """This JPEG algorithm result is slightly different from cv2. |
| 451 | DiffJPEG supports batch processing. |
| 452 | |
| 453 | Args: |
| 454 | differentiable(bool): If True, uses custom differentiable rounding function, if False, uses standard torch.round |
| 455 | """ |
| 456 | |
| 457 | def __init__(self, differentiable=True): |
| 458 | super(DiffJPEG, self).__init__() |
| 459 | if differentiable: |
| 460 | rounding = diff_round |
| 461 | else: |
| 462 | rounding = torch.round |
| 463 | |
| 464 | self.compress = CompressJpeg(rounding=rounding) |
| 465 | self.decompress = DeCompressJpeg(rounding=rounding) |
| 466 | |
| 467 | def forward(self, x, quality): |
| 468 | """ |
| 469 | Args: |
| 470 | x (Tensor): Input image, bchw, rgb, [0, 1] |
| 471 | quality(float): Quality factor for jpeg compression scheme. |
| 472 | """ |
| 473 | factor = quality |
| 474 | if isinstance(factor, (int, float)): |
| 475 | factor = quality_to_factor(factor) |
| 476 | else: |
| 477 | for i in range(factor.size(0)): |
| 478 | factor[i] = quality_to_factor(factor[i]) |
| 479 | h, w = x.size()[-2:] |
| 480 | h_pad, w_pad = 0, 0 |
| 481 | # why should use 16 |
| 482 | if h % 16 != 0: |
| 483 | h_pad = 16 - h % 16 |
| 484 | if w % 16 != 0: |
| 485 | w_pad = 16 - w % 16 |
| 486 | x = F.pad(x, (0, w_pad, 0, h_pad), mode='constant', value=0) |
| 487 | |
| 488 | y, cb, cr = self.compress(x, factor=factor) |
| 489 | recovered = self.decompress(y, cb, cr, (h + h_pad), (w + w_pad), factor=factor) |
| 490 | recovered = recovered[:, :, 0:h, 0:w] |
| 491 | return recovered |
| 492 | |
| 493 | |
| 494 | if __name__ == '__main__': |