Args: image(tensor): batch x 3 x height x width Returns: dict(tensor): Compressed tensor with batch x h*w/64 x 8 x 8.
(self, image, factor=1)
| 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 ------------------------# |