Read videos from HIM project files

data/xhdfs.py

@@ -0,0 +1,92 @@
+from typing import NamedTuple, Generator, Any
+
+import numpy as np
+
+CAMERA_NODE_TYPES = ('unknown', 'depth', 'ir', 'rgb', 'rgb_luminance', 'rgb_mapped')
+
+HEADER_BYTES = 32
+BYTES_PER_SHORT = 2
+TIMESTAMP_BYTES = 8
+RAW_UNRELIABLE_PIXEL_VALUE = -11
+
+ENDIANNESS = 'little'
+
+# Pixels are stored as signed 16-bit integers with little endianness.
+# However, the only valid negative value is `-11` (unreliable pixel), which should just be converted into the
+# corresponding unsigned value, without any loss of information
+SOURCE_NUMPY_DATATYPE = np.dtype('<i2')
+TARGET_NUMPY_DATATYPE = np.dtype('=u2')  # TODO: is this cross-platform?
+MIN_PIXEL_VALUE = np.iinfo(TARGET_NUMPY_DATATYPE).min
+MAX_PIXEL_VALUE = np.iinfo(TARGET_NUMPY_DATATYPE).max
+UNRELIABLE_PIXEL_VALUE = np.cast[TARGET_NUMPY_DATATYPE](RAW_UNRELIABLE_PIXEL_VALUE).min()
+
+XHDFSHeader = NamedTuple(typename='XHDFSHeader',
+                         fields=[('n_frames', int),
+                                 ('width', int),
+                                 ('height', int),
+                                 ('version_number', int),
+                                 ('camera_node_type', str)])
+
+
+def read_unsigned_integer(fileobj, n_bytes: int = 4, signed=False) -> int:
+    return int.from_bytes(fileobj.read(n_bytes), ENDIANNESS, signed=signed)
+
+
+def bytes_to_image(frame_bytes: bytes, width: int, height: int) -> np.ndarray:
+    frame_bytes = frame_bytes[8:]  # Cut off the timestamp
+
+    # Height & width in that order match what OpenCV expects
+    frame_image = np.frombuffer(frame_bytes, SOURCE_NUMPY_DATATYPE).reshape(height, width)
+    frame_image = frame_image.astype(TARGET_NUMPY_DATATYPE, casting='unsafe', copy=False)
+    frame_image[frame_image == UNRELIABLE_PIXEL_VALUE] = 0
+
+    return frame_image
+
+
+class XHDFS:
+    def __init__(self, fileobj):
+        self._fileobj = fileobj
+
+        self._header = None
+        self._pixels_per_frame = -1
+
+    def __enter__(self):
+        header = self._read_header()
+
+        self._header = header
+        self._pixels_per_frame = header.width * header.height * BYTES_PER_SHORT + TIMESTAMP_BYTES
+
+        return self
+
+    def _read_header(self) -> XHDFSHeader:
+        n_frames = read_unsigned_integer(self._fileobj)
+        width = read_unsigned_integer(self._fileobj)
+        height = read_unsigned_integer(self._fileobj)
+
+        # Skip `CameraNodeId` & `WithBodies`
+        _ = self._fileobj.read(4 + 1)
+
+        version = read_unsigned_integer(self._fileobj)
+        node_type = read_unsigned_integer(self._fileobj, n_bytes=2)
+        node_type = CAMERA_NODE_TYPES[node_type]
+
+        # Skip `CameraNodeSerial` & 3 "spare" bytes
+        _ = self._fileobj.read(6 + 3)
+
+        header = XHDFSHeader(n_frames, width, height, version, node_type)
+        return header
+
+    def _read_frame_bytes(self) -> bytes:
+        return self._fileobj.read(self._pixels_per_frame)
+
+    def frame_sequence(self) -> Generator[np.ndarray, None, None]:
+        w, h = self._header.width, self._header.height
+        for _ in range(self._header.n_frames):
+            frame_bytes = self._read_frame_bytes()
+            frame = bytes_to_image(frame_bytes, w, h)
+            del frame_bytes
+            yield frame
+        pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._fileobj.close()