Social Icons

January 18, 2016

Fast YUV to RGB conversion in Python 3

This Python scripts implements fast matrix-based conversion from YUV to RGB colospace. Conversion from YV12 (YUV420) requires one extra line of code to perform up-sampling of chroma planes.

import numpy as np
from scipy import ndimage

def ConvertYUVtoRGB(yuv_planes)
    plane_y  = yuv_planes[0]
    plane_u  = yuv_planes[1]
    plane_v  = yuv_planes[2]
    
    height = plane_y.shape[0]
    width  = plane_y.shape[1]
    
    # upsample if YV12
    plane_u = ndimage.zoom(plane_u, 2, order=0)
    plane_v = ndimage.zoom(plane_v, 2, order=0)
    # alternativelly can perform upsampling with numpy.repeat()
    #plane_u = plane_u.repeat(2, axis=0).repeat(2, axis=1)
    #plane_v = plane_v.repeat(2, axis=0).repeat(2, axis=1)
    
    # reshape
    plane_y  = plane_y.reshape((plane_y.shape[0], plane_y.shape[1], 1))
    plane_u  = plane_u.reshape((plane_u.shape[0], plane_u.shape[1], 1))
    plane_v  = plane_v.reshape((plane_v.shape[0], plane_v.shape[1], 1))
    
    # make YUV of shape [height, width, color_plane]
    yuv = np.concatenate((plane_y, plane_u, plane_v), axis=2)
    
    # according to ITU-R BT.709
    yuv[:,:, 0] = yuv[:,:, 0].clip(16, 235).astype(yuv.dtype) - 16
    yuv[:,:,1:] = yuv[:,:,1:].clip(16, 240).astype(yuv.dtype) - 128
    
    A = np.array([[1.164,  0.000,  1.793],
                  [1.164, -0.213, -0.533],
                  [1.164,  2.112,  0.000]])
    
    # our result
    rgb = np.dot(yuv, A.T).clip(0, 255).astype('uint8')
    
    return rgb


Based on this discussion on stackoverflow.
See also: color conversion.

No comments:

Post a Comment