Cleanup and optimize yuv420p_to_png.py

video/yuv420p_to_png.py

@@ -1,61 +1,60 @@
+#!/usr/bin/env python3
+
 # Author : sherin.s@gmail.com (Sherin Sasidhan)
 # Date   : 05-Oct-2010
- 
-import Image
+# URL    : https://shrex999.wordpress.com/2013/07/31/yuv-to-rgb-python-imaging-library/
+#
+# Updated:  2019-06-15 - Doncho Gunchev <dgunchev@gmail.com>
+
+from PIL import Image
 import sys
 from struct import *
-import array
- 
+from array import array
+from os.path import splitext
+
 if len(sys.argv) != 4:
-        print "***** Usage syntax Error!!!! *****\n"
-        print "Usage:"
-        print "python <script> <.yuv file yuv420p> <width&gt> <height> "
-        sys.exit(1) # exit
-else:
-        pass
- 
-image_name = sys.argv[1]
-width = int(sys.argv[2])
-height = int(sys.argv[3])
- 
-y = array.array('B')
-u = array.array('B')
-v = array.array('B')
- 
-f_y = open(image_name, "rb")
-f_uv = open(image_name, "rb")
-f_uv.seek(width*height, 1)
- 
-image_out = Image.new("RGB", (width, height))
-pix = image_out.load()
- 
-print "width=", width, "height=", height
- 
-for i in range(0, height/2):
-    for j in range(0, width/2):
-        u.append(ord(f_uv.read(1)));
- 
-for i in range(0, height/2):
-    for j in range(0, width/2):
-        v.append(ord(f_uv.read(1)));
-for i in range(0,height):
-    for j in range(0, width):
-        y.append(ord(f_y.read(1)));
-        #print "i=", i, "j=", j , (i*width), ((i*width) +j)
-        #pix[j, i] = y[(i*width) +j], y[(i*width) +j], y[(i*width) +j]
-        Y_val = y[(i*width)+j]
-        U_val = u[((i/2)*(width/2))+(j/2)]
-        V_val = v[((i/2)*(width/2))+(j/2)]
-        B = 1.164 * (Y_val-16) + 2.018 * (U_val - 128)
-        G = 1.164 * (Y_val-16) - 0.813 * (V_val - 128) - 0.391 * (U_val - 128)
-        R = 1.164 * (Y_val-16) + 1.596*(V_val - 128)
-        pix[j, i] = int(R), int(G), int(B)
- 
-######################################################
-# B = 1.164(Y - 16)                   + 2.018(U - 128)
-# G = 1.164(Y - 16) - 0.813(V - 128) - 0.391(U - 128)
-# R = 1.164(Y - 16) + 1.596(V - 128)
-######################################################
- 
-#image_out.save("out.bmp")
-image_out.show()
+    sys.exit("""***** Usage syntax Error!!!! *****
+
+Usage:
+    %r yuv420p_FILE.yuv WIDTH HEIGHT
+""" % sys.argv[0])
+
+def load_yuv420p(yuv_file_path, width, height):
+    """Read YUV420p file and return PIL Image"""
+    yuv_file = open(yuv_file_path, "rb")
+
+    y_data = array('B')
+    y_data.fromfile(yuv_file, height * width)
+    u_data = array('B')
+    u_data.fromfile(yuv_file, (height // 2) * (width // 2))
+    v_data = array('B')
+    v_data.fromfile(yuv_file, (height // 2) * (width // 2))
+
+    result = Image.new("RGB", (width, height))
+    pixels = result.load()
+    for i in range(height):
+        for j in range(width):
+            y_value = (y_data[(i * width) + j] - 16) * 1.164
+            offset = ((i // 2) * (width // 2)) + (j // 2)
+            u_value = u_data[offset] - 128
+            v_value = v_data[offset] - 128
+            # B = 1.164(Y - 16)                  + 2.018(U - 128)
+            # G = 1.164(Y - 16) - 0.813(V - 128) - 0.391(U - 128)
+            # R = 1.164(Y - 16) + 1.596(V - 128)
+            b_value = y_value + 2.018 * u_value
+            g_value = y_value - 0.813 * v_value - 0.391 * u_value
+            r_value = y_value + 1.596 * v_value
+            pixels[j, i] = (int(r_value), int(g_value), int(b_value))
+
+    return result
+
+def main():
+    yuv_file_path = sys.argv[1]
+    width, height = int(sys.argv[2]), int(sys.argv[3])
+    picture = load_yuv420p(yuv_file_path, width, height)
+    picture.save(splitext(yuv_file_path)[0] + '.png')
+    picture.show()
+
+
+if __name__ == "__main__":
+    main()