From 01ae276b034129bb7c3df95863e9d511fb6063f7 Mon Sep 17 00:00:00 2001
From: Sebastian Hahta <joseha@utu.fi>
Date: Wed, 15 Jan 2020 19:41:29 +0200
Subject: [PATCH] encoding and StreamWriter
---
python/ftl/codecs.py | 175 +++++++++++++++++++++++++---------
python/ftl/ftlstreamwriter.py | 75 +++++++++++++++
2 files changed, 205 insertions(+), 45 deletions(-)
create mode 100644 python/ftl/ftlstreamwriter.py
diff --git a/python/ftl/codecs.py b/python/ftl/codecs.py
index a4c2d2ac1..bf4e21e68 100644
--- a/python/ftl/codecs.py
+++ b/python/ftl/codecs.py
@@ -10,46 +10,21 @@ from . misc import Calibration
from enum import IntEnum
-class FTLDecoder:
- def decode(self, packet):
- raise NotImplementedError()
-
-def _int_to_float(im):
- return im.astype(float) / np.iinfo(im.dtype).max
-
-################################################################################
-# OpenCV (optional)
-################################################################################
-
+_has_opencv = False
try:
import cv2 as cv
-
- def decode_codec_opencv(packet):
- if packet.block_total != 1 or packet.block_number != 0:
- raise Exception("Unsupported block format (todo)")
-
- return _int_to_float(cv.imdecode(np.frombuffer(packet.data, dtype=np.uint8),
- cv.IMREAD_UNCHANGED))
-
- def decode_codec_opencv_float(packet):
- if packet.block_total != 1 or packet.block_number != 0:
- raise Exception("Unsupported block format (todo)")
-
- return cv.imdecode(np.frombuffer(packet.data, dtype=np.uint8),
- cv.IMREAD_UNCHANGED).astype(np.float) / 1000.0
-
- def _ycrcb2rgb(img):
- return _int_to_float(cv.cvtColor(img, cv.COLOR_YCrCb2RGB))
-
+ _has_opencv = True
except ImportError:
warn("OpenCV not available. OpenCV required for full functionality.")
- def decode_codec_opencv(packet):
- raise Exception("OpenCV required for OpenCV (png/jpeg) decoding")
+def _int_to_float(im):
+ return im.astype(float) / np.iinfo(im.dtype).max
- def decode_codec_opencv_float(packet):
- raise Exception("OpenCV required for OpenCV (png/jpeg) decoding")
+if _has_opencv:
+ def _ycrcb2rgb(img):
+ return _int_to_float(cv.cvtColor(img, cv.COLOR_YCrCb2RGB))
+else:
def _ycrcb2rgb(img):
""" YCrCb to RGB, based on OpenCV documentation definition.
@@ -84,13 +59,39 @@ def _ycbcr2rgb(img):
return rgb / 255
+################################################################################
+# Decoding
+################################################################################
+
+class FTLDecoder:
+ def decode(self, packet):
+ raise NotImplementedError()
+
+################################################################################
+# OpenCV (optional)
+################################################################################
+
+def decode_codec_opencv(packet):
+ if packet.block_total != 1 or packet.block_number != 0:
+ raise Exception("Unsupported block format (todo)")
+
+ return _int_to_float(cv.imdecode(np.frombuffer(packet.data, dtype=np.uint8),
+ cv.IMREAD_UNCHANGED))
+
+def decode_codec_opencv_float(packet):
+ if packet.block_total != 1 or packet.block_number != 0:
+ raise Exception("Unsupported block format (todo)")
+
+ return cv.imdecode(np.frombuffer(packet.data, dtype=np.uint8),
+ cv.IMREAD_UNCHANGED).astype(np.float) / 1000.0
+
################################################################################
# HEVC
################################################################################
# components/codecs/include/ftl/codecs/hevc.hpp
-class NALType(IntEnum):
+class _NALType(IntEnum):
CODED_SLICE_TRAIL_N = 0
CODED_SLICE_TRAIL_R = 1
@@ -166,14 +167,14 @@ class NALType(IntEnum):
UNSPECIFIED_63 = 63
INVALID = 64
-def get_NAL_type(data):
+def _get_NAL_type(data):
if not isinstance(data, bytes):
raise ValueError("expected bytes")
- return NALType((data[4] >> 1) & 0x3f)
+ return _NALType((data[4] >> 1) & 0x3f)
-def is_iframe(data):
- return get_NAL_type(data) == NALType.VPS
+def _is_iframe(data):
+ return _get_NAL_type(data) == _NALType.VPS
class FTLDecoder_HEVC:
def __init__(self):
@@ -182,7 +183,7 @@ class FTLDecoder_HEVC:
def decode(self, packet):
if not self._seen_iframe:
- if not is_iframe(packet.data):
+ if not _is_iframe(packet.data):
# can't decode before first I-frame has been received
warn("received P-frame before I-frame")
return
@@ -284,12 +285,11 @@ def decode_codec_pose(packet):
################################################################################
def create_decoder(codec, channel, version=3):
- """ Create decoder for given channel, codec and ftlf version.
-
- @param codec Codec id
- @param channel Channel id
- @param version FTL file version
- @returns callable, which takes packet as argument
+ """ @brief Create decoder for given channel, codec and ftlf version.
+ @param codec Codec id
+ @param channel Channel id
+ @param version FTL file version
+ @returns callable which takes packet as argument
"""
if codec == ftltype.codec_t.HEVC:
@@ -302,12 +302,18 @@ def create_decoder(codec, channel, version=3):
return FTLDecoder_HEVC_YCbCr().decode
elif codec == ftltype.codec_t.PNG:
+ if not _has_opencv:
+ raise Exception("OpenCV required for OpenCV (png/jpeg) decoding")
+
if ftltype.is_float_channel(channel):
return decode_codec_opencv_float
else:
return decode_codec_opencv
elif codec == ftltype.codec_t.JPG:
+ if not _has_opencv:
+ raise Exception("OpenCV required for OpenCV (png/jpeg) decoding")
+
return decode_codec_opencv
elif codec == ftltype.codec_t.MSGPACK:
@@ -326,3 +332,82 @@ def create_decoder(codec, channel, version=3):
else:
raise ValueError("Unknown codec %i" % codec)
+
+################################################################################
+# ENCODING
+################################################################################
+
+def create_packet(codec, definition, flags, data):
+ return ftltype.Packet._make((codec, definition, 1, 0, flags, data))
+
+# TODO exception types?
+
+def encode_codec_opencv_jpg(data, **kwargs):
+ params = []
+ retval, encoded = cv.imencode(".jpg", data, params)
+ if retval:
+ return create_packet(ftltype.codec_t.JPG,
+ ftltype.get_definition(data),
+ 0,
+ encoded)
+ else:
+ # todo
+ raise Exception("encoding error")
+
+def encode_codec_opencv_png(data, **kwargs):
+ params = [cv.IMWRITE_PNG_COMPRESSION, 9]
+ retval, encoded = cv.imencode(".png", data, params)
+ if retval:
+ return create_packet(ftltype.codec_t.PNG,
+ ftltype.get_definition(data),
+ 0,
+ encoded)
+ else:
+ # todo
+ raise Exception("encoding error")
+
+def encode_codec_opencv_png_float(data, compression=9):
+ data = (data * 1000).astype(np.uint16)
+ params = [cv.IMWRITE_PNG_COMPRESSION, compression]
+ retval, encoded = cv.imencode(".png", data, params)
+ if retval:
+ return create_packet(ftltype.codec_t.PNG,
+ ftltype.get_definition(data),
+ 0,
+ encoded)
+ else:
+ # todo
+ raise Exception("encoding error")
+
+def create_encoder(codec, channel, **options):
+ """ @brief Create encoder
+ @param codec codec id
+ @param channel channel id
+ @param **options options passed to codec constructor
+ @returns callable which takes unencoded data and optional parameters
+ """
+
+ if codec == ftltype.codec_t.JPG:
+ if not ftltype.is_float_channel(channel):
+ return encode_codec_opencv_jpg
+ else:
+ raise Exception("JPG not supported for float channels")
+
+ elif codec == ftltype.codec_t.PNG:
+ if ftltype.is_float_channel(channel):
+ return encode_codec_opencv_png_float
+ else:
+ return encode_codec_opencv_png
+
+ elif codec == ftltype.codec_t.MSGPACK:
+ if channel == ftltype.Channel.Pose:
+ raise NotImplementedError("todo")
+
+ elif channel == ftltype.Channel.Calibration:
+ raise NotImplementedError("todo")
+
+ else:
+ raise Exception("msgpack only available for pose/calibration")
+
+ else:
+ raise Exception("unsupported/unknown codec")
diff --git a/python/ftl/ftlstreamwriter.py b/python/ftl/ftlstreamwriter.py
new file mode 100644
index 000000000..3dfc8903b
--- /dev/null
+++ b/python/ftl/ftlstreamwriter.py
@@ -0,0 +1,75 @@
+import msgpack
+import struct
+
+from . import ftltype
+
+from codecs import create_encoder
+
+class FTLStreamWriter:
+ def __init__(self, file, version=3):
+ self._file = open(file, "wb")
+ self._file.write(bytes(ord(c) for c in "FTLF")) # magic
+ self._file.write(bytes([version])) # version
+ self._file.write(bytes([0]*64)) # reserved
+ self._packer = msgpack.Packer(strict_types=False, use_bin_type=True)
+
+ self._encoders = {}
+ self._channel_count = 0
+
+ def __del__(self):
+ self.close()
+
+ def close(self):
+ self._file.close()
+
+ def add_raw(self, sp, p):
+ if len(sp) != len(ftltype.StreamPacket._fields):
+ raise ValueError("invalid StreamPacket")
+
+ if len(p) != len(ftltype.Packet._fields):
+ raise ValueError("invalid Packet")
+
+ self._file.write(self._packer.pack((sp, p)))
+ self._file.flush()
+
+ def create_encoder(self, source, codec, channel, **kwargs):
+ if channel not in ftltype.Channel:
+ raise ValueError("unknown channel")
+
+ if not isinstance(source, int):
+ raise ValueError("source id must be int")
+
+ if source < 0:
+ raise ValueError("source id must be positive")
+
+ encoder = create_encoder(codec, channel, **kwargs)
+ self._encoders[(int(source), int(channel))] = encoder
+ self._channel_count += 1
+
+ def encode(self, source, timestamp, channel, data):
+ if not isinstance(source, int):
+ raise ValueError("source id must be int")
+
+ if source < 0:
+ raise ValueError("source id must be positive")
+
+ if timestamp < 0:
+ raise ValueError("timestamp must be positive")
+
+ if channel not in ftltype.Channel:
+ raise ValueError("unknown channel")
+
+ try:
+ p = self._encoders[(int(source), int(channel))](data)
+ except KeyError:
+ raise Exception("no encoder found, create_encoder() has to be" +
+ "called for every source and channel")
+ except Exception as ex:
+ raise Exception("Encoding error:" + str(ex))
+
+ sp = ftltype.StreamPacket._make((timestamp,
+ int(source),
+ int(channel),
+ self._channel_count))
+
+ self.add_raw(sp, p)
--
GitLab