Dev (#4)

* utils branch

* creating tests for generate_frame

* changes to setup

* ignoreing local, added version

* constants, vidio etc

- add constants
- move tpg to vidio
- init modules in __init__
- compare in tests

* adding ConvertAXIStream CS

added colorconv back
restructured code
added constants

* update for readme

* added 422 and 420

renamed GenRGAXIStream to GenAXIStream,
changed some constants

* added 422 and 420

renamed GenRGAXIStream to GenAXIStream,
changed some constants

* adding test_video

* readme change

* todo 420

* GenAXIStream also returns the matrix now

version change

.gitignore

@@ -1,3 +1,4 @@
 *.egg-info
 dist
-*local
+*local
+*venv*

README.md

@@ -1,14 +1,15 @@
 # vidio (video io)
 
+This generates the AXIStream frame for video subsystems to be used with cocotb.  
 This generates the AXIStream frame for video subsystems to be used with cocotb.  
 
 Functions:
-- GenRGBAXIStream: Generates an AXIS frame
+- GenAXIStream: Generates an AXIS frame
 - ConvertAXIStreamCS: Converts color space RGB to YUV [In development] 
 
 Roadmap:
-- AXIS to matrix
-- Matrix to AXIS
+- AXIS to matrix [Add 422, 420]
+- Matrix to AXIS 
 - Generate color bars AXIS/matrix, supporting ITU-R BT.601-7, ITU-R BT.709-5
 - Color convertion
 - Chroma convertion

cocotbext/vidio/__init__.py

@@ -1,2 +1,2 @@
 from .version import __version__
-from .vidio import GenRGBAXIStream, ConvertAXIStreamCS
+from .vidio import GenAXIStream, ConvertAXIStreamCS

cocotbext/vidio/constants.py

@@ -1,13 +1,10 @@
 import enum
 
-class VideoType(enum.IntEnum):
+class ColorFormat(enum.IntEnum):
     RGB = 0
-    YUV = 1
-
-class ChromaType(enum.IntEnum):
-    YUV444 = 0
-    YUV422 = 1
-    YUV420 = 2
+    YUV444 = 1
+    YUV422 = 2
+    YUV420 = 3
 
 class Standard(enum.IntEnum):
     BT601 = 0
@@ -18,5 +15,3 @@ class Pattern(enum.IntEnum):
     p_incr = 0
     rand = 1
     h_incr = 2
-
-

cocotbext/vidio/utils.py

@@ -1,8 +1,9 @@
 """Utilities"""
 import numpy as np
 from cocotbext.axi import AxiStreamFrame
+from .constants import ColorFormat
 
-def mat2axis(matrix, resH, pixelPerClock, pixelQuant):
+def mat2axis(matrix, resH, pixelPerClock, pixelQuant, outFormat):
     """ 
     converts matrix to axis
     - matrix: numpy array of shape mx3, where m = rows*columns 3 = color channels
@@ -19,22 +20,60 @@ def mat2axis(matrix, resH, pixelPerClock, pixelQuant):
     for j in range(v):
         packed_list = []
         for _ in range(pack_range):
-            packed_value = (frame[k+1, 2] << q*5) | (frame[k+1, 1] << q*4) | (frame[k+1, 0] << q*3) | (frame[k, 2] << q*2) | (frame[k, 1] << q) | frame[k, 0]
-            packed_list.append(packed_value.tobytes())
+            if outFormat == ColorFormat.RGB or outFormat == ColorFormat.YUV444:
+                if p==2 :
+                    packed_value = (frame[k+1, 2] << q*5) | (frame[k+1, 1] << q*4) | (frame[k+1, 0] << q*3) | (frame[k, 2] << q*2) | (frame[k, 1] << q) | frame[k, 0]
+                    packed_list.append(packed_value.tobytes())
+                elif p==4 :
+                    packed_value = (frame[k+2, 0] << q*6) |  (frame[k+1, 2] << q*5) | (frame[k+1, 1] << q*4) | (frame[k+1, 0] << q*3) | (frame[k, 2] << q*2) | (frame[k, 1] << q) | frame[k, 0]
+                    packed_list.append(packed_value.tobytes())
+                    packed_value = (frame[k+3, 2] << q*5-4) | (frame[k+3, 1] << q*4-4) | (frame[k+3, 0] << q*3-4) | (frame[k+2, 2] << q*2-4) | (frame[k+2, 1] << q-4) | (frame[k+2, 0]  >> 4)
+                    packed_list.append(packed_value.tobytes())
+            elif outFormat == ColorFormat.YUV422: 
+                if p==2 :
+                    packed_value = (frame[k, 2] << q*3) | (frame[k+1, 0] << q*2) | (frame[k, 1] << q) | frame[k, 0]
+                    packed_list.append(packed_value.tobytes())
+                elif p==4 :
+                    packed_value = (frame[k+3, 0] << q*6) | (frame[k+2, 1] << q*5) | (frame[k+2, 0] << q*4) | (frame[k, 2] << q*3) | (frame[k+1, 0] << q*2) | (frame[k, 1] << q) | frame[k, 0]
+                    packed_list.append(packed_value.tobytes())
+                    packed_value = (frame[k+2, 2] << q-4) | (frame[k+3, 0] >> 4)
+                    packed_list.append(packed_value.tobytes())
+            elif outFormat == ColorFormat.YUV420:
+                if (j%2) :
+                    if p==2 :
+                        packed_value = (0 << q*3) | (frame[k+1, 0] << q*2) | (0 << q) | frame[k, 0]
+                        packed_list.append(packed_value.tobytes())
+                    elif p==4 :
+                        packed_value = (frame[k+3, 0] << q*6) | (0 << q*5) | (frame[k+2, 0] << q*4) | (0 << q*3) | (frame[k+1, 0] << q*2) | (0 << q) | frame[k, 0]
+                        packed_list.append(packed_value.tobytes())
+                        packed_value = (0 << q-4) | (frame[k+3, 0] >> 4)
+                        packed_list.append(packed_value.tobytes())
+                else :
+                    if p==2 :
+                        packed_value = (frame[k, 2] << q*3) | (frame[k+1, 0] << q*2) | (frame[k, 1] << q) | frame[k, 0]
+                        packed_list.append(packed_value.tobytes())
+                    elif p==4 : # 64-bit intergers so ned to seperate this
+                        packed_value = (frame[k+3, 0] << q*6) | (frame[k+2, 1] << q*5) | (frame[k+2, 0] << q*4) | (frame[k, 2] << q*3) | (frame[k+1, 0] << q*2) | (frame[k, 1] << q) | frame[k, 0]
+                        packed_list.append(packed_value.tobytes())
+                        packed_value = (frame[k+2, 2] << q-4) | (frame[k+3, 0] >> 4)
+                        packed_list.append(packed_value.tobytes())
             k = k + p
         bytearray_ = bytearray()
         for byte_array in packed_list:
             bytearray_.extend(byte_array)
-
+        tuser = []
+        if p == 2: 
+            tuser=[1]*8 + [0] * (len(bytearray_)-8) if j == 0 else [0]*len(bytearray_)
+        elif p == 4:
+            tuser=[1]*16 + [0] * (len(bytearray_)-16) if j == 0 else [0]*len(bytearray_)
         single_axi_frame = AxiStreamFrame(
                 tdata=bytearray_,
-                tuser=[1]*8 + [0] * (len(bytearray_)-8) if j == 0 else [0]*len(bytearray_)
+                tuser = tuser
             )
         full_axi_frame.append(single_axi_frame)
-
     return full_axi_frame
 
-def axis2mat(axisFrame, resH, pixelPerClock, pixelQuant):
+def axis2mat(axisFrame, resH, pixelPerClock, pixelQuant, fmt):
     """ 
     converts axis to matrix
     - axisFrame: cocotbext.axi.AxiStreamFrame
@@ -46,15 +85,30 @@ def axis2mat(axisFrame, resH, pixelPerClock, pixelQuant):
     pack_range = int(c/p)
     q = pixelQuant
     mat = np.zeros((r, c, 3), dtype=np.int16)
-    mask = (1 << q) - 1
-    for i in range(r):
-        for j in range(pack_range):
-            d_val = int.from_bytes(axisFrame[i].tdata[j*8:8*(j+1)], byteorder='little', signed=False)
-            mat[i][2*j][0] = d_val >> 0*10 & mask
-            mat[i][2*j][1] = d_val >> 1*10 & mask
-            mat[i][2*j][2] = d_val >> 2*10 & mask
-            mat[i][2*j+1][0] = d_val >> 3*10 & mask
-            mat[i][2*j+1][1] = d_val >> 4*10 & mask
-            mat[i][2*j+1][2] = d_val >> 5*10 & mask
+    if (fmt == ColorFormat.RGB or fmt == ColorFormat.YUV444):
+        mask = (1 << q) - 1
+        for i in range(r):
+            for j in range(pack_range):
+                d_val = int.from_bytes(axisFrame[i].tdata[j*8:8*(j+1)], byteorder='little', signed=False)
+                mat[i][2*j][0] = d_val >> 0*10 & mask
+                mat[i][2*j][1] = d_val >> 1*10 & mask
+                mat[i][2*j][2] = d_val >> 2*10 & mask
+                mat[i][2*j+1][0] = d_val >> 3*10 & mask
+                mat[i][2*j+1][1] = d_val >> 4*10 & mask
+                mat[i][2*j+1][2] = d_val >> 5*10 & mask
+    elif (fmt == ColorFormat.YUV422):
+        pass
+    else:
+        pass
 
-    return np.reshape(mat, newshape=(r*c,3))
+    return np.reshape(mat, newshape=(r*c,3))
+
+def rgbtoy444(matrix):
+    pass
+
+def csy444toy422(y444):
+    pass
+
+def csy422toy420():
+    pass
+

cocotbext/vidio/version.py

@@ -1 +1 @@
-__version__ = "0.0.3"
+__version__ = "0.0.5"

cocotbext/vidio/vidio.py

@@ -1,55 +1,77 @@
 import numpy as np
 
-from .utils import mat2axis, axis2mat
-from .constants import VideoType, Standard, Pattern
+from .utils import mat2axis, axis2mat, csy444toy422
+from .constants import Standard, Pattern, ColorFormat
 from .csc import rgb2yuv
 
-def GenRGBAXIStream(resH, resV, pixelPerClock, pixelQuant, pattern):
-    """ 
+def GenAXIStream(resH, resV, pixelPerClock, pixelQuant, pattern, format):
+    """
     - returns: AXIS transaction for one frame
     - resh: Horizontal resoltuion, >1
     - resv: Vertixal resolution, >1
     - pixelperclock: Pixel per clock in AXI transaction, [2]
     - pixelquant: Quantization, 10
-    - pattern: 
+    - pattern:
         - p_incr = counter increment on every pixel
         - rand = random
         - h_incr = horizontal pixel increment
+    - format: RGB, YUV444, YUV422, YUV420
     """
     c ,r, q = resH, resV, pixelQuant
-    mat = np.zeros((r, c, 3), dtype=np.uint16)
+    rgb = np.zeros((r, c, 3), dtype=np.uint16)
 
     if pattern == Pattern.p_incr:
         for i in range(r):
             for j in range(c):
                 for k in range(3):
-                    mat[i, j, k] = 512 + (i * r * 3 + j * 3 + k) % (2**q)
+                    rgb[i, j, k] = (i * r * 3 + j * 3 + k) % (2**q)
     elif pattern == Pattern.rand:
-        mat = np.random.rand(r, c, 3)
-        mat = (mat*1023).astype(np.uint16)
+        rgb = np.random.rand(r, c, 3)
+        rgb = (rgb*((2**q)-1)).astype(np.uint16)
     elif pattern == Pattern.h_incr:
         for i in range(r):
             for j in range(c):
-                mat[i, j, :] = j
+                rgb[i, j, :] = j % (2**q)
     else:
         raise ValueError("Unknown pattern")
 
-    mat = np.reshape(mat, newshape=(r*c,3))
+    mask = np.ones((r*c, 3), dtype=np.uint16)
+    vcount, hcount = 0, 0
+    if (format == ColorFormat.YUV422):
+        for i in range(r*c):
+            if (i%2) :
+                mask[i][1:] = [0, 0]
+    elif (format == ColorFormat.YUV420):
+        for i in range(r*c):
+            if (i%2) :
+                mask[i][1:] = [0, 0]
+            if vcount%2 :
+                mask[i][1:] = [0, 0]
+            if hcount == resH-1:
+                vcount = vcount + 1
+                hcount = 0
+            else:
+                hcount = hcount + 1
+
+    rgbreshape = np.reshape(rgb, newshape=(r*c,3))
+    rgbreshape = rgbreshape*mask
+    axis = mat2axis(rgbreshape, resH, pixelPerClock, 10, format)
+
+    return axis, rgbreshape
 
-    return mat2axis(mat, resH, pixelPerClock, pixelQuant)
-
-def ConvertAXIStreamCS(resH, pixelPerClock, pixelQuant, outputVideoType, axisFrame):
-    """ 
-    - returns: Color converted AXIS 
+
+def ConvertAXIStreamCS(resH, pixelPerClock, pixelQuant, outputFormat, axisFrame):
+    """
+    - returns: Color converted AXIS
     - resh: Horizontal resoltuion, >1
     - pixelperclock: Pixel per clock in AXI transaction, [2]
     - pixelquant: Quantization, 10
-    - outputVideoType: VideoType
+    - outputFormat: RGB, YUV444, YUV422, YUV420
     - axisFrame: input frame to color convert
     """
     mat = axis2mat(axisFrame, resH, pixelPerClock, pixelQuant)
-    if outputVideoType == VideoType.YUV:
+    if outputFormat == ColorFormat.YUV444:
         conv = rgb2yuv(mat, Standard.BT2020)
-        return mat2axis(conv.T, resH, pixelPerClock, pixelQuant)
+        return mat2axis(conv.T, resH, pixelPerClock, pixelQuant, outputFormat)
     else:
-        raise ValueError("Unknown outputVideoType")
+        raise ValueError("Unknown outputFormat")

tests/test_tpg.py → tests/test_vidio.py

@@ -1,14 +1,19 @@
 import unittest
-from cocotbext.vidio import GenRGBAXIStream
+
+from cocotbext.vidio import GenAXIStream
+from cocotbext.vidio.utils import axis2mat
 
 class Test_tpg(unittest.TestCase):
 	def test_GenRGBAXIStream(self):
-		axi_frame = GenRGBAXIStream(8, 8, 2, 10, "p_incr")
+		axi_frame = GenAXIStream(8, 8, 2, 10, "p_incr")
 		tuser = axi_frame[0].tuser
 		self.assertEqual(tuser, [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
 
 class Test_utils(unittest.TestCase):
 	pass
 
 
-unittest.main()
+# unittest.main()
+if __name__=="__main__":
+	axis = GenAXIStream(16,4,4,10,0,2)
+	print((axis[0]))