93 lines
3.2 KiB
Python
93 lines
3.2 KiB
Python
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()
|