Args: compressed(dict(tensor)): batch x h*w/64 x 8 x 8 imgh(int) imgw(int) factor(float) Returns: Tensor: batch x 3 x height x width
(self, y, cb, cr, imgh, imgw, factor=1)
| 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 ------------------------ # |