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.

1 comment:

  1. Nice work. I'm not sure what the one line of code you are referring to for upsampling is. Would it be one of the two lines you have commented out?

    ReplyDelete