Chroma subsampling on CbCr channels
| 71 | |
| 72 | |
| 73 | class ChromaSubsampling(nn.Module): |
| 74 | """ Chroma subsampling on CbCr channels |
| 75 | """ |
| 76 | |
| 77 | def __init__(self): |
| 78 | super(ChromaSubsampling, self).__init__() |
| 79 | |
| 80 | def forward(self, image): |
| 81 | """ |
| 82 | Args: |
| 83 | image(tensor): batch x height x width x 3 |
| 84 | |
| 85 | Returns: |
| 86 | y(tensor): batch x height x width |
| 87 | cb(tensor): batch x height/2 x width/2 |
| 88 | cr(tensor): batch x height/2 x width/2 |
| 89 | """ |
| 90 | image_2 = image.permute(0, 3, 1, 2).clone() |
| 91 | cb = F.avg_pool2d(image_2[:, 1, :, :].unsqueeze(1), kernel_size=2, stride=(2, 2), count_include_pad=False) |
| 92 | cr = F.avg_pool2d(image_2[:, 2, :, :].unsqueeze(1), kernel_size=2, stride=(2, 2), count_include_pad=False) |
| 93 | cb = cb.permute(0, 2, 3, 1) |
| 94 | cr = cr.permute(0, 2, 3, 1) |
| 95 | return image[:, :, :, 0], cb.squeeze(3), cr.squeeze(3) |
| 96 | |
| 97 | |
| 98 | class BlockSplitting(nn.Module): |