#python #view-yuv Load YUV file and render it using Pygame

view-yuv.py

#!/usr/bin/env python

import argparse
import os
import pygame
from pygame.locals import *
import sys

# constants
supported_formats = ('yv12', 'iyuv', 'i420', 'nv12', 'nv21')

# TODO:
# create FourCC constants and map FourCC string to them
# map FourCC constants to bits per pixel

# http://www.fourcc.org/yuv.php
# yv12: 420 planar, Y, V, U
# iyuv, i420: 420 planar, Y, U, V
# nv12: 420 semi-planar, Y, UVUV
# nv21: 420 semi-planar, Y, VUVU
# yuv2, yuyv: 422 packed, YUYV
# uyvy: 422 packed, UYVY
# yuyv: 422 packed, YUYV


def testHardwareSupport():
    size = (32, 32)
    pygame.init()
    pygame.display.set_mode(size)
    pygame.display.set_caption('viewyuv: ' + file)
    # , pygame.YUV2_OVERLAY
    for pixformat in (pygame.YV12_OVERLAY, pygame.IYUV_OVERLAY, pygame.UYVY_OVERLAY, pygame.YVYU_OVERLAY):
        overlay = pygame.Overlay(pixformat, size)
        print("pixformat[%x]:    %s" % (pixformat, overlay.get_hardware()))


def getBitsPerPixel(pixformat):
    # all currently supported formats are 420
    return 12


def renderFrame(frame):
    # print('frame %d' % frame)
    fd.seek(frame*frame_size)
    if pixformat == 'yv12':
        y = fd.read(plane_size)
        u = fd.read(plane_size / 4)
        v = fd.read(plane_size / 4)
    elif pixformat == 'iyuv' or pixformat == 'i420':
        y = fd.read(plane_size)
        v = fd.read(plane_size / 4)
        u = fd.read(plane_size / 4)
    elif pixformat == 'nv12':
        y = fd.read(plane_size)
        uv = fd.read(plane_size/2)
        u = uv[0:][::2]  # even pixels
        v = uv[1:][::2]  # odd pixels
    elif pixformat == 'nv21':
        y = fd.read(plane_size)
        uv = fd.read(plane_size/2)
        v = uv[0:][::2]  # even pixels
        u = uv[1:][::2]  # odd pixels

    overlay.display((y, u, v))

parser = argparse.ArgumentParser(
    description='Load YUV file and render it using Pygame.')
parser.add_argument('file',
                    help='YUV file to render')
parser.add_argument('-s', '--size', nargs=2, type=int, metavar=('width', 'height'),
                    help='specify width and height of YUV file', default=(1280, 720))
parser.add_argument('-f', '--pix_fmt',
                    help='specify pixel format of the YUV file', default='yv12')
parser.add_argument('--fps', type=int,
                    help='specify the frame rate per second', default=30)

args = parser.parse_args()
# print(str(args))

# testHardwareSupport()

file = args.file
width = args.size[0]
height = args.size[1]
pixformat = args.pix_fmt
fps = args.fps

paused = False

if not os.path.exists(file):
    print("file not exists: \"%s\"" % file)
    sys.exit(-1)

pygame.init()
window = pygame.display.set_mode(args.size)
pygame.display.set_caption('viewyuv: ' + file)
overlay = pygame.Overlay(pygame.YV12_OVERLAY, args.size)
clock = pygame.time.Clock()

i = 0
plane_size = width * height
frame_size = plane_size * getBitsPerPixel(pixformat) / 8
file_size = os.stat(file).st_size
num_frames = file_size/frame_size
if not file_size % frame_size == 0:
    print("Warning: file should contain integer number of frames!!")
    print("file_size: %d, frame_size: %d" % (file_size, frame_size))


fd = open(file, 'rb')
while i < num_frames:
    for event in pygame.event.get():
        if event.type == QUIT:
            sys.exit(0)
        elif event.type == KEYDOWN:
            if event.key == K_SPACE:
                paused = not paused
            else:
                # these will all pause playback
                # and render one frame
                if event.key == K_LEFT:
                    if i > 0:
                        i = i-1
                if event.key == K_RIGHT:
                    if i < num_frames:
                        i = i+1
                if event.key == K_HOME:
                    i = 0
                if event.key == K_END:
                    i = num_frames-1

                paused = True
                renderFrame(i)

    if not paused:
        renderFrame(i)
        i += 1

    clock.tick(fps)

print("rendered %d frames" % num_frames)