Source code for hello.transforms.transforms

import cv2 as cv
import numpy as np




def gen_bgr(h, w, input_dtype="uint8"):
    bgr = np.zeros((h, w, 3), dtype=input_dtype)
    return bgr





def gen_nv12(h, w, input_dtype="uint8"):
    nv12 = np.random.randint(256, size=(int(h*1.5), w)).astype(input_dtype)
    return nv12.flatten()





def bgr_to_yuv444(bgr):
    img_h, img_w = bgr.shape[:2]
    y_size = img_h * img_w

    uv_h, uv_w = img_h // 2, img_w // 2
    uv_size = uv_h * uv_w

    # bgr -> yuv420p(YU12)
    yuv420p = cv.cvtColor(bgr, cv.COLOR_BGR2YUV_I420)
    yuv420p = yuv420p.flatten()

    # yuv420p(YU12) -> yuv444
    img_y = yuv420p[:y_size].reshape((img_h, img_w, 1))

    img_u = yuv420p[y_size:y_size + uv_size]
    img_u = img_u.reshape(uv_h, uv_w, 1)

    img_u = np.repeat(img_u, 2, axis=0)
    img_u = np.repeat(img_u, 2, axis=1)

    img_v = yuv420p[y_size + uv_size:]
    img_v = img_v.reshape(uv_h, uv_w, 1)

    img_v = np.repeat(img_v, 2, axis=0)
    img_v = np.repeat(img_v, 2, axis=1)

    yuv444 = np.concatenate((img_y, img_u, img_v), axis=2)
    return yuv444





def bgr_to_nv12(bgr):
    # NV12 shape as (h*1.5,w)
    # YUV_NV12: Y(y:h*w)UV(u,v:h*w/4)
    # YUV_I420: Y(y:h*w)U(u:h*w/4)V(v:h*w/4)
    h, w = bgr.shape[:2]
    area = h * w

    yuv420p = cv.cvtColor(bgr, cv.COLOR_BGR2YUV_I420)
    yuv420p = yuv420p.reshape((area * 3 // 2,))

    y = yuv420p[:area]
    uv_planar = yuv420p[area:].reshape((2, area // 4))
    uv_packed = uv_planar.transpose((1, 0)).reshape((area // 2,))

    nv12 = np.zeros_like(yuv420p)
    nv12[:area] = y
    nv12[area:] = uv_packed
    return nv12.reshape((-1, w))





def bgr_to_nv21(bgr):
    # NV21 shape as (h*1.5,w)
    # YUV_NV21: Y(y:h*w)VU(v,u:h*w/4)
    # YUV_YV12: Y(y:h*w)V(v:h*w/4)U(u:h*w/4)
    h, w = bgr.shape[:2]
    area = h * w

    yuv420p = cv.cvtColor(bgr, cv.COLOR_BGR2YUV_YV12)
    yuv420p = yuv420p.reshape((area * 3 // 2,))

    y = yuv420p[:area]
    vu_planar = yuv420p[area:].reshape((2, area // 4))
    vu_packed = vu_planar.transpose((1, 0)).reshape((area // 2,))

    nv21 = np.zeros_like(yuv420p)
    nv21[:area] = y
    nv21[area:] = vu_packed
    return nv21.reshape((-1, w))





def nv12_to_bgr(nv12, width=None):
    if nv12.ndim == 1:
        assert width is not None
        nv12 = nv12.reshape((-1, width))

    assert nv12.ndim == 2, "shape as (height*1.5, width)"

    bgr = cv.cvtColor(nv12, cv.COLOR_YUV2BGR_NV12)
    return bgr





def nv21_to_bgr(nv21, width=None):
    if nv21.ndim == 1:
        assert width is not None
        nv21 = nv21.reshape((-1, width))

    assert nv21.ndim == 2, "shape as (height*1.5, width)"

    bgr = cv.cvtColor(nv21, cv.COLOR_YUV2BGR_NV21)
    return bgr