Feature/python

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: