fix colors (channel order)

0319d141 · Sebastian · a48b5c69 · 0319d141 · 0319d141
Commit 0319d141 authored 5 years ago by Sebastian
--- a/python/ftl/ftlstream.py
+++ b/python/ftl/ftlstream.py
 import msgpack
+
+import numpy as np
+
+from enum import IntEnum
 from collections import namedtuple
 from . libde265 import Decoder

+try:
+    import cv2 as cv
+    def ycbcr2rgb(img):
+        raise NotImplementedError("TODO")
+    
+except ImportError:
+    import skimage.color
+    def ycbcr2rgb(img):
+        res = skimage.color.ycbcr2rgb(img.astype(np.float))
+
+        # clip
+        res[res > 1.0] = 1.0
+        res[res < 0.0] = 0.0
+
+        # skimage ycbcr2rgb() returns dtype float64, convert to uint8
+        return (res * 255).astype(np.uint8)
+
+# FTL definitions
+
 _packet = namedtuple("Packet", ["codec", "definition", "block_total", "block_number", "flags", "data"])
 _stream_packet = namedtuple("StreamPacket", ["timestamp", "streamID", "chanel_count", "channel"])

@@ -17,6 +40,88 @@ _definition_t = {
    8 : ()
 }

+class NALType(IntEnum):
+    CODED_SLICE_TRAIL_N = 0
+    CODED_SLICE_TRAIL_R = 1
+
+    CODED_SLICE_TSA_N = 2
+    CODED_SLICE_TSA_R = 3
+
+    CODED_SLICE_STSA_N = 4
+    CODED_SLICE_STSA_R = 5
+
+    CODED_SLICE_RADL_N = 6
+    CODED_SLICE_RADL_R = 7
+
+    CODED_SLICE_RASL_N = 8
+    CODED_SLICE_RASL_R = 9
+
+    RESERVED_VCL_N10 = 10
+    RESERVED_VCL_R11 = 11
+    RESERVED_VCL_N12 = 12
+    RESERVED_VCL_R13 = 13
+    RESERVED_VCL_N14 = 14
+    RESERVED_VCL_R15 = 15
+
+    CODED_SLICE_BLA_W_LP = 16
+    CODED_SLICE_BLA_W_RADL = 17
+    CODED_SLICE_BLA_N_LP = 18
+    CODED_SLICE_IDR_W_RADL = 19
+    CODED_SLICE_IDR_N_LP = 20
+    CODED_SLICE_CRA = 21
+    RESERVED_IRAP_VCL22 = 22
+    RESERVED_IRAP_VCL23 = 23
+
+    RESERVED_VCL24 = 24
+    RESERVED_VCL25 = 25
+    RESERVED_VCL26 = 26
+    RESERVED_VCL27 = 27
+    RESERVED_VCL28 = 28
+    RESERVED_VCL29 = 29
+    RESERVED_VCL30 = 30
+    RESERVED_VCL31 = 31
+
+    VPS = 32
+    SPS = 33
+    PPS = 34
+    ACCESS_UNIT_DELIMITER = 35
+    EOS = 36
+    EOB = 37
+    FILLER_DATA = 38
+    PREFIX_SEI = 39
+    SUFFIX_SEI = 40
+
+    RESERVED_NVCL41 = 41
+    RESERVED_NVCL42 = 42
+    RESERVED_NVCL43 = 43
+    RESERVED_NVCL44 = 44
+    RESERVED_NVCL45 = 45
+    RESERVED_NVCL46 = 46
+    RESERVED_NVCL47 = 47
+    UNSPECIFIED_48 = 48
+    UNSPECIFIED_49 = 49
+    UNSPECIFIED_50 = 50
+    UNSPECIFIED_51 = 51
+    UNSPECIFIED_52 = 52
+    UNSPECIFIED_53 = 53
+    UNSPECIFIED_54 = 54
+    UNSPECIFIED_55 = 55
+    UNSPECIFIED_56 = 56
+    UNSPECIFIED_57 = 57
+    UNSPECIFIED_58 = 58
+    UNSPECIFIED_59 = 59
+    UNSPECIFIED_60 = 60
+    UNSPECIFIED_61 = 61
+    UNSPECIFIED_62 = 62
+    UNSPECIFIED_63 = 63
+    INVALID = 64
+
+def get_NAL_type(data):
+    if not isinstance(data, bytes):
+        raise ValueError("expected bytes")
+    
+    return NALType((data[4] >> 1) & 0x3f)
+
 class FTLStream:
    def __init__(self, file):
        self._file = open(file, "br")
@@ -60,12 +165,18 @@ class FTLStream:
            self._decoders[k] = Decoder(_definition_t[p.definition])
        
        decoder = self._decoders[k]
+        
        decoder.push_data(p.data)
        decoder.decode()
+        
        img = decoder.get_next_picture()
        
        if img is not None:
-            self._frames[k] = img
+            self._frames[k] = ycbcr2rgb(img[:,:,(0,2,1)]) # note: channels BGR
+    
+    def _flush_decoders(self):
+        for decoder in self._decoders.values():
+            decoder.flush_data()
    
    def read(self):
        '''
@@ -99,6 +210,7 @@ class FTLStream:
                    raise ValueError("unkowno codec %i" % p.codec)
            
            except Exception as e:
+                # TODO: Multiple exceptions possible. Re-design read()?
                ex = e
            
            try:

--- a/python/ftl/libde265.py
+++ b/python/ftl/libde265.py
@@ -4,62 +4,78 @@ Python wrapper for libde265. Only decoding is (partly) implemented.
 Requirements:
 * libde265 library (libde265.so.0)
 * numpy
- * skimage (rescaling)
+ * opencv or skimage

 '''

+try:
+    import cv2 as cv
+    
+    # TODO: test
+    def resize(img, size):
+        return cv.resize(img, dsize=reversed(size), interpolation=cv.INTER_CUBIC)
+    
+except ImportError:
+    from skimage.transform import resize as resize_skimage
+    
+    def resize(img, size):
+        # skimage resize() return dtype float64, convert back to uint8
+        # order: 0 nn, 1 bilinear, 3 bicubic
+        return (resize_skimage(img, size, order=3, mode="constant", cval=0) * 255).astype(np.uint8)
+
 import ctypes
+from enum import IntEnum

 import numpy as np
-from skimage.transform import rescale

 # error codes copied from header (de265.h)

-DE265_OK = 0
-DE265_ERROR_NO_SUCH_FILE=1
-DE265_ERROR_COEFFICIENT_OUT_OF_IMAGE_BOUNDS=4
-DE265_ERROR_CHECKSUM_MISMATCH=5
-DE265_ERROR_CTB_OUTSIDE_IMAGE_AREA=6
-DE265_ERROR_OUT_OF_MEMORY=7
-DE265_ERROR_CODED_PARAMETER_OUT_OF_RANGE=8
-DE265_ERROR_IMAGE_BUFFER_FULL=9
-DE265_ERROR_CANNOT_START_THREADPOOL=10
-DE265_ERROR_LIBRARY_INITIALIZATION_FAILED=11
-DE265_ERROR_LIBRARY_NOT_INITIALIZED=12
-DE265_ERROR_WAITING_FOR_INPUT_DATA=13
-DE265_ERROR_CANNOT_PROCESS_SEI=14
-DE265_ERROR_PARAMETER_PARSING=15
-DE265_ERROR_NO_INITIAL_SLICE_HEADER=16
-DE265_ERROR_PREMATURE_END_OF_SLICE=17
-DE265_ERROR_UNSPECIFIED_DECODING_ERROR=18
-DE265_ERROR_NOT_IMPLEMENTED_YET = 502
-DE265_WARNING_NO_WPP_CANNOT_USE_MULTITHREADING = 1000
-DE265_WARNING_WARNING_BUFFER_FULL=1001
-DE265_WARNING_PREMATURE_END_OF_SLICE_SEGMENT=1002
-DE265_WARNING_INCORRECT_ENTRY_POINT_OFFSET=1003
-DE265_WARNING_CTB_OUTSIDE_IMAGE_AREA=1004
-DE265_WARNING_SPS_HEADER_INVALID=1005
-DE265_WARNING_PPS_HEADER_INVALID=1006
-DE265_WARNING_SLICEHEADER_INVALID=1007
-DE265_WARNING_INCORRECT_MOTION_VECTOR_SCALING=1008
-DE265_WARNING_NONEXISTING_PPS_REFERENCED=1009
-DE265_WARNING_NONEXISTING_SPS_REFERENCED=1010
-DE265_WARNING_BOTH_PREDFLAGS_ZERO=1011
-DE265_WARNING_NONEXISTING_REFERENCE_PICTURE_ACCESSED=1012
-DE265_WARNING_NUMMVP_NOT_EQUAL_TO_NUMMVQ=1013
-DE265_WARNING_NUMBER_OF_SHORT_TERM_REF_PIC_SETS_OUT_OF_RANGE=1014
-DE265_WARNING_SHORT_TERM_REF_PIC_SET_OUT_OF_RANGE=1015
-DE265_WARNING_FAULTY_REFERENCE_PICTURE_LIST=1016
-DE265_WARNING_EOSS_BIT_NOT_SET=1017
-DE265_WARNING_MAX_NUM_REF_PICS_EXCEEDED=1018
-DE265_WARNING_INVALID_CHROMA_FORMAT=1019
-DE265_WARNING_SLICE_SEGMENT_ADDRESS_INVALID=1020
-DE265_WARNING_DEPENDENT_SLICE_WITH_ADDRESS_ZERO=1021
-DE265_WARNING_NUMBER_OF_THREADS_LIMITED_TO_MAXIMUM=1022
-DE265_NON_EXISTING_LT_REFERENCE_CANDIDATE_IN_SLICE_HEADER=1023
-DE265_WARNING_CANNOT_APPLY_SAO_OUT_OF_MEMORY=1024
-DE265_WARNING_SPS_MISSING_CANNOT_DECODE_SEI=1025
-DE265_WARNING_COLLOCATED_MOTION_VECTOR_OUTSIDE_IMAGE_AREA=1026
+class libde265error(IntEnum):
+    DE265_OK = 0
+    DE265_ERROR_NO_SUCH_FILE=1
+    DE265_ERROR_COEFFICIENT_OUT_OF_IMAGE_BOUNDS=4
+    DE265_ERROR_CHECKSUM_MISMATCH=5
+    DE265_ERROR_CTB_OUTSIDE_IMAGE_AREA=6
+    DE265_ERROR_OUT_OF_MEMORY=7
+    DE265_ERROR_CODED_PARAMETER_OUT_OF_RANGE=8
+    DE265_ERROR_IMAGE_BUFFER_FULL=9
+    DE265_ERROR_CANNOT_START_THREADPOOL=10
+    DE265_ERROR_LIBRARY_INITIALIZATION_FAILED=11
+    DE265_ERROR_LIBRARY_NOT_INITIALIZED=12
+    DE265_ERROR_WAITING_FOR_INPUT_DATA=13
+    DE265_ERROR_CANNOT_PROCESS_SEI=14
+    DE265_ERROR_PARAMETER_PARSING=15
+    DE265_ERROR_NO_INITIAL_SLICE_HEADER=16
+    DE265_ERROR_PREMATURE_END_OF_SLICE=17
+    DE265_ERROR_UNSPECIFIED_DECODING_ERROR=18
+    DE265_ERROR_NOT_IMPLEMENTED_YET = 502
+    DE265_WARNING_NO_WPP_CANNOT_USE_MULTITHREADING = 1000
+    DE265_WARNING_WARNING_BUFFER_FULL=1001
+    DE265_WARNING_PREMATURE_END_OF_SLICE_SEGMENT=1002
+    DE265_WARNING_INCORRECT_ENTRY_POINT_OFFSET=1003
+    DE265_WARNING_CTB_OUTSIDE_IMAGE_AREA=1004
+    DE265_WARNING_SPS_HEADER_INVALID=1005
+    DE265_WARNING_PPS_HEADER_INVALID=1006
+    DE265_WARNING_SLICEHEADER_INVALID=1007
+    DE265_WARNING_INCORRECT_MOTION_VECTOR_SCALING=1008
+    DE265_WARNING_NONEXISTING_PPS_REFERENCED=1009
+    DE265_WARNING_NONEXISTING_SPS_REFERENCED=1010
+    DE265_WARNING_BOTH_PREDFLAGS_ZERO=1011
+    DE265_WARNING_NONEXISTING_REFERENCE_PICTURE_ACCESSED=1012
+    DE265_WARNING_NUMMVP_NOT_EQUAL_TO_NUMMVQ=1013
+    DE265_WARNING_NUMBER_OF_SHORT_TERM_REF_PIC_SETS_OUT_OF_RANGE=1014
+    DE265_WARNING_SHORT_TERM_REF_PIC_SET_OUT_OF_RANGE=1015
+    DE265_WARNING_FAULTY_REFERENCE_PICTURE_LIST=1016
+    DE265_WARNING_EOSS_BIT_NOT_SET=1017
+    DE265_WARNING_MAX_NUM_REF_PICS_EXCEEDED=1018
+    DE265_WARNING_INVALID_CHROMA_FORMAT=1019
+    DE265_WARNING_SLICE_SEGMENT_ADDRESS_INVALID=1020
+    DE265_WARNING_DEPENDENT_SLICE_WITH_ADDRESS_ZERO=1021
+    DE265_WARNING_NUMBER_OF_THREADS_LIMITED_TO_MAXIMUM=1022
+    DE265_NON_EXISTING_LT_REFERENCE_CANDIDATE_IN_SLICE_HEADER=1023
+    DE265_WARNING_CANNOT_APPLY_SAO_OUT_OF_MEMORY=1024
+    DE265_WARNING_SPS_MISSING_CANNOT_DECODE_SEI=1025
+    DE265_WARNING_COLLOCATED_MOTION_VECTOR_OUTSIDE_IMAGE_AREA=1026

 libde265 = ctypes.cdll.LoadLibrary("libde265.so.0")

@@ -117,7 +133,7 @@ class Decoder:
        
    def __del__(self):
        libde265.de265_free_decoder(self.ctx_)
-        
+    
    def get_error_str(self, code):
        return libde265.de265_get_error_text(code).decode("ascii")
    
@@ -148,21 +164,21 @@ class Decoder:
    def decode(self):
        err = libde265.de265_decode(self.ctx_, self.more_)
        
-        if err and err != DE265_ERROR_WAITING_FOR_INPUT_DATA:
+        if err and err != libde265error.DE265_ERROR_WAITING_FOR_INPUT_DATA:
            raise Exception(self.get_error_str(err))
-            
+        
        return self.more_.value != 0
    
    def flush_data(self):
-        libde265.flush_data(self.ctx_)
+        err = libde265.de265_flush_data(self.ctx_)
+        
+        if err:
+            raise Exception(self.get_error_str(err))
        
    def get_next_picture(self):
        '''
-        Returns next decoded frame. Image in YCrCb format. If no frame available
+        Returns next decoded frame. Image in YCbCr format. If no frame available
        returns None.
-
-        TODO: Skimage's ycrcb2rgb() does not look correct. Is Cr/Cb channel
-              created correctly?
        '''
        img = libde265.de265_get_next_picture(self.ctx_)
        
@@ -172,20 +188,19 @@ class Decoder:
        res = np.zeros((self.size_[0], self.size_[1], 3), dtype=np.uint8)
        
        for c in range(0, 3):
-            size = (libde265.de265_get_image_height(img, c), libde265.de265_get_image_width(img, c))
-            bpp = libde265.de265_get_bits_per_pixel(img, c)
-            img_ptr = libde265.de265_get_image_plane(img, c, self.out_stride_)
-            scale_y = self.size_[0] / size[0]
-            scale_x = self.size_[1] / size[1]
+            size = (libde265.de265_get_image_height(img, c),
+                    libde265.de265_get_image_width(img, c))
            
+            bpp = libde265.de265_get_bits_per_pixel(img, c)
+
            if bpp != 8:
                raise NotImplementedError("unsupported bits per pixel %i" % bpp)
            
-            if scale_x != scale_y:
-                raise NotImplementedError("unsupported subsampling")
+            img_ptr = libde265.de265_get_image_plane(img, c, self.out_stride_)
+            
+            ch = np.frombuffer(img_ptr[:size[0] * size[1]], dtype=np.uint8)
+            ch.shape = size
            
-            channel = np.asarray(bytearray(img_ptr[:size[0] * size[1]]), dtype=np.uint8)
-            channel.shape = size
-            res[:,:,c] = (rescale(channel, scale_x) * 255).astype(np.uint8)
+            res[:,:,c] = resize(ch, self.size_)

-        return res
\ No newline at end of file
+        return res