def test_444_to_420():
x = torch.rand(1, 3, 32, 32)
y, u, v = yuv_444_to_420(x)
assert u.size(0) == v.size(0) == y.size(0) == x.size(0)
assert u.size(1) == v.size(1) == y.size(1) == 1
assert y.size(2) == x.size(2) and y.size(3) == x.size(3)
assert u.size(2) == v.size(2) == (y.size(2) // 2)
assert u.size(3) == v.size(3) == (y.size(2) // 2)
assert (x[:, [0]] == y).all()
with pytest.raises(ValueError):
y, u, v = yuv_444_to_420(x, mode="toto")
y, u, v = yuv_444_to_420(x.chunk(3, 1))
def test_444_to_420():
x = torch.rand(1, 3, 32, 32)
y, u, v = yuv_444_to_420(x)
assert u.size(0) == v.size(0) == y.size(0) == x.size(0)
assert u.size(1) == v.size(1) == y.size(1) == 1
assert y.size(2) == x.size(2) and y.size(3) == x.size(3)
assert u.size(2) == v.size(2) == (y.size(2) // 2)
assert u.size(3) == v.size(3) == (y.size(2) // 2)
assert (x[:, [0]] == y).all()
def convert_rgb_yuv420(
frame: Tensor) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
# yuv420 [0, 2**bitdepth-1] to rgb 444 [0, 1] only for now
return yuv_444_to_420(rgb2ycbcr(frame), mode="avg_pool")