Args: y(tensor): y channel image cb(tensor): cb channel cr(tensor): cr channel Returns: Tensor: batch x height x width x 3
(self, y, cb, cr)
| 353 | super(ChromaUpsampling, self).__init__() |
| 354 | |
| 355 | def forward(self, y, cb, cr): |
| 356 | """ |
| 357 | Args: |
| 358 | y(tensor): y channel image |
| 359 | cb(tensor): cb channel |
| 360 | cr(tensor): cr channel |
| 361 | |
| 362 | Returns: |
| 363 | Tensor: batch x height x width x 3 |
| 364 | """ |
| 365 | |
| 366 | def repeat(x, k=2): |
| 367 | height, width = x.shape[1:3] |
| 368 | x = x.unsqueeze(-1) |
| 369 | x = x.repeat(1, 1, k, k) |
| 370 | x = x.view(-1, height * k, width * k) |
| 371 | return x |
| 372 | |
| 373 | cb = repeat(cb) |
| 374 | cr = repeat(cr) |
| 375 | return torch.cat([y.unsqueeze(3), cb.unsqueeze(3), cr.unsqueeze(3)], dim=3) |
| 376 | |
| 377 | |
| 378 | class YCbCr2RGBJpeg(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected