1.1.0版本

UniTAP/common/dsc_compression_info.py

@@ -0,0 +1,96 @@
+from enum import IntEnum
+
+
+class DscCompressionInfo:
+    """
+
+    Class contains information about DSC compression used on frame.
+
+    """
+    class DscColorFormat(IntEnum):
+        """
+        Contains values of possible color format.
+        """
+        CF_NONE = -1
+        CF_RGB = 0
+        CF_YCbCr_422 = 1
+        CF_YCbCr_444 = 2
+        CF_YCbCr_420 = 3
+        CF_Simple_422 = 4
+
+    def __init__(self):
+        self.color_format = DscCompressionInfo.DscColorFormat.CF_NONE
+        self.bpp = 0
+        self.is_block_prediction_enabled = False
+        self.h_slice_size = 0
+        self.v_slice_size = 0
+        self.buffer_bit_depth = 0
+        self.version = (0, 0)
+        self.is_simple_as_444 = False
+
+    def is_valid(self) -> bool:
+        """
+        Return state of the video frame and check color_format, bpp, h and v slice_size and DSC version.
+        If everything ok, return True, otherwise - False.
+
+        Returns:
+            object of `bool` type.
+        """
+        return self.color_format is not None and\
+               self.bpp > 0 and\
+               self.h_slice_size > 0 and\
+               self.v_slice_size > 0 and\
+               self.version in [(1, 2), (1, 1)]
+
+
+def create_from_pps(pps_bytearray: bytearray) -> DscCompressionInfo:
+    """
+    Fill structure 'DscCompressionInfo' from PPS header of the DSC image.
+
+    Returns:
+        object of `DscCompressionInfo` type.
+    """
+    if b'DSCF' not in pps_bytearray:
+        if len(pps_bytearray) < 128:
+            raise ValueError("Incorrect PPS size!")
+        pps = pps_bytearray
+    else:
+        if len(pps_bytearray) < 132:
+            raise ValueError("Incorrect PPS size!")
+        pps = pps_bytearray[4:]
+
+    is_yuv = ((pps[4] >> 4) & 1) == 0
+    is_simple422 = ((pps[4] >> 3) & 1) == 1
+    is_native422 = (pps[88] & 1) == 1
+    is_native420 = ((pps[88] >> 1) & 1) == 1
+    width = (pps[8] << 8) | pps[9]
+    height = (pps[6] << 8) | pps[7]
+    slice_height = (pps[10] << 8) | pps[11]
+    slice_width = (pps[12] << 8) | pps[13]
+
+    info = DscCompressionInfo()
+
+    info.version = (pps[0] >> 4 & 0xf, pps[0] & 0xf)
+    info.buffer_bit_depth = pps[3] & 0xf
+    info.is_block_prediction_enabled = bool(pps[4] >> 5 & 0x1)
+    info.h_slice_size = int(width / slice_width)
+    info.v_slice_size = int(height / slice_height)
+
+    if is_native422 or is_native420:
+        info.bpp = int(pps[4] & 0x3 << 8 | pps[5]) / 2
+    else:
+        info.bpp = int(pps[4] & 0x3 << 8 | pps[5])
+
+    if is_yuv:
+        if is_simple422:
+            info.color_format = DscCompressionInfo.DscColorFormat.CF_Simple_422
+        elif is_native422:
+            info.color_format = DscCompressionInfo.DscColorFormat.CF_YCbCr_422
+        elif is_native420:
+            info.color_format = DscCompressionInfo.DscColorFormat.CF_YCbCr_420
+        else:
+            info.color_format = DscCompressionInfo.DscColorFormat.CF_YCbCr_444
+    else:
+        info.color_format = DscCompressionInfo.DscColorFormat.CF_RGB
+
+    return info