NV21和RGB互相转换的工具

nv21Tojpg

#!/usr/bin/env python3
import cv2
import numpy as np

import glob


def convert_fhq(h, w, msg):
    img_y=np.fromstring(msg[:h*w],dtype='uint8').reshape((h,w)).astype('int32')
    img_u=np.fromstring(msg[h*w:h*w+h*w//2:2],dtype='uint8').reshape((h//2,w//2)).astype('int32')
    img_v=np.fromstring(msg[h*w+1:h*w+h*w//2:2],dtype='uint8').reshape((h//2,w//2)).astype('int32')
    ruv=((359*(img_v-128))>>8)
    guv=-1*((88*(img_u-128)+183*(img_v-128))>>8)
    buv=((454*(img_u-128))>>8)
    ruv=np.repeat(np.repeat(ruv,2,axis=0),2,axis=1)
    guv=np.repeat(np.repeat(guv,2,axis=0),2,axis=1)
    buv=np.repeat(np.repeat(buv,2,axis=0),2,axis=1)
    img_r=(img_y+ruv).clip(0,255).astype('uint8')
    img_g=(img_y+guv).clip(0,255).astype('uint8')
    img_b=(img_y+buv).clip(0,255).astype('uint8')
    img=np.dstack([img_b[:,:,None],img_g[:,:,None],img_r[:,:,None]])
    img=img.transpose((1,0,2))[::-1].copy()
    #img = cv2.resize(img,(0,0),fx=0.25,fy=0.25)
    return img[:,:,::-1].copy()


h, w = 480, 640

for i in glob.glob('*.raw'):
    img = convert_fhq(h, w, open(i, 'rb').read())
    cv2.imwrite(i + '.png', img)

jpg2nv21

#!/usr/bin/env python3
import numpy as np
import glob
from PIL import Image
import cv2
def my_open_jpg(filename):
    "open a jpg file and return a file object"
    return Image.open(filename);

def my_jpg2nv21(im_jpg):
    "conver a jpg file to yuv file"
    #if(len(im_jpg.split())==4):
    #    r,g,b,a=im_jpg.split();
    #else:
    r,g,b=im_jpg.split();

    width,height=im_jpg.size;
    print(width)
    print(height)
    im_new=list(range(int(width*height*3/2)));
    r = list(r.getdata())
    g = list(g.getdata())
    b = list(b.getdata())
    Y,U,V = 0,0,0
    for xi in range(height):
        for xj in range(width):
            i = (xi//2)*2;
            j = (xj//2)*2;
            index1 = i*width+j
            index2 = i*width+j + 1
            index3 = (i + 1)*width + j;
            index4 = (i + 1)*width + j + 1
            R = (r[index1] + r[index2] + r[index3] + r[index4])//4
            G = (g[index1] + g[index2] + g[index3] + g[index4])//4
            B = (b[index1] + b[index2] + b[index3] + b[index4])//4
            #Y = k_r*R + k_g*G + k_b*B;
            Y,U,V = rgb_to_ycbcr(R, G, B)
            im_new[int(width*height+i//2*width+j)] = V;
            im_new[int(width*height+i//2*width+j+1)] = U;
            im_new[xi*width+xj]=Y;
    return im_new;
def my_save_yuv(filename,im_new):
    #use the numpy to write the data to file
    fp = open(filename,"wb");

    data=np.array(im_new,"B");
    data.tofile(fp);
    fp.close();
    print("save yuv file %s successfully",  filename)

def rgb_to_ycbcr(R, G, B):
    assert(255>=R)
    assert(255>=G)
    assert(255>=B)
    Y = 0.298*R + 0.612*G + 0.117*B;
    U = -0.168*R - 0.330*G + 0.498*B + 128;
    V = 0.449*R - 0.435*G - 0.083*B + 128;
    return int(Y), int(U), int(V)

def rgb2ycbcr(im):
    xform = np.array([[.299, .587, .114], [-.1687, -.3313, .5], [.5, -.4187, -.0813]])
    ycbcr = im.dot(xform.T)
    ycbcr[:,:,[1,2]] += 128
    return np.uint8(ycbcr)

def ycbcr2rgb(im):
    xform = np.array([[1, 0, 1.402], [1, -0.34414, -.71414], [1, 1.772, 0]])
    rgb = im.astype(np.float)
    rgb[:,:,[1,2]] -= 128
    rgb = rgb.dot(xform.T)
    np.putmask(rgb, rgb > 255, 255)
    np.putmask(rgb, rgb < 0, 0)
    return np.uint8(rgb)

if __name__ == "__main__":
    h, w = 480, 640
    for file in glob.glob('*.jpg'):
        img = my_open_jpg(file)
        print(img);
        yuv = my_jpg2nv21(img)
        yuvFileName = file + ".yuv"
        my_save_yuv(yuvFileName,yuv)

        #png file name to test
        pngFileName = file + "1.jpg"
        fin = open(yuvFileName,"rb")
        img_nv21 = np.fromstring(fin.read(h * w * 3 // 2), dtype=np.uint8)
        img_nv21 = img_nv21.reshape((h + h // 2, w))
        img = cv2.cvtColor(img_nv21, cv2.COLOR_YUV2BGR_NV21)
        cv2.imwrite(pngFileName, img)
    print("end");