MAI/udp_dump_capture.py

import socket
from pathlib import Path

import cv2
import numpy as np


class MikPacketAssembler:
    HEADER_SIZE = 8
    FLAG_END = 0x01
    FLAG_START = 0x02
    MAX_ARRAY_SIZE = 256 * 1024 * 1024

    def __init__(self):
        self.current = None
        self.expected_packet = 0
        self.expected_offset = 0
        self.dropped_arrays = 0

    def push(self, payload):
        if len(payload) < self.HEADER_SIZE:
            raise ValueError("UDP payload is shorter than the MIK packet header")
        flags = payload[1]
        sequence = payload[2]
        packet_number = payload[3]
        value = int.from_bytes(payload[4:8], "little")
        packet_data = payload[8:]

        if flags & self.FLAG_START:
            if self.current is not None:
                self.dropped_arrays += 1
            if value <= 0 or value > self.MAX_ARRAY_SIZE:
                self.current = None
                return None
            self.current = {
                "sequence": sequence,
                "data": bytearray(value),
                "size": value,
            }
            self.expected_packet = (packet_number + 1) & 0xFF
            self.expected_offset = min(len(packet_data), value)
            self.current["data"][:self.expected_offset] = packet_data[:self.expected_offset]
            if flags & self.FLAG_END and self.expected_offset == value:
                result = bytes(self.current["data"])
                self.current = None
                return result
            return None

        if self.current is None or sequence != self.current["sequence"]:
            return None
        if packet_number != self.expected_packet or value != self.expected_offset:
            self.dropped_arrays += 1
            self.current = None
            return None

        self.expected_packet = (self.expected_packet + 1) & 0xFF
        end = self.expected_offset + len(packet_data)
        if end > self.current["size"]:
            self.dropped_arrays += 1
            self.current = None
            return None
        self.current["data"][self.expected_offset:end] = packet_data
        self.expected_offset = end

        if flags & self.FLAG_END:
            if self.expected_offset == self.current["size"]:
                result = bytes(self.current["data"])
                self.current = None
                return result
            self.dropped_arrays += 1
            self.current = None
        return None


class UdpDumpCapture:
    """VideoCapture-compatible reader for framed MIK UDP packet logs."""

    PORT = 59004
    PACKET_HEADER_SIZE = 8
    LABEL_SIZE = 40
    FLAG_END = 0x01
    FLAG_START = 0x02
    PIXEL_GRAY8 = 0x01
    PIXEL_GRAY16 = 0x02
    PIXEL_RGB888 = 0x03
    PIXEL_YCBCR422 = 0x0A
    PIXEL_INT16 = 0x12
    MAX_ARRAY_SIZE = 256 * 1024 * 1024

    def __init__(self, path, fps=30.0, port=None):
        self.path = Path(path)
        self.port = int(port) if port is not None else None
        self._init_decoder(fps)
        self._file = None
        self._next_frame = None
        self._assembler = MikPacketAssembler()
        try:
            self._file = self.path.open("rb", buffering=8 * 1024 * 1024)
            envelope = self._file.read(4)
            if len(envelope) != 4:
                raise ValueError("not a framed UDP packet log")
            observed_port = int.from_bytes(envelope[:2], "little")
            if observed_port <= 0 or (self.port is not None and observed_port != self.port):
                raise ValueError(f"unexpected UDP port: {observed_port}")
            self.port = observed_port
            self._file.seek(0)
            self._next_frame = self._read_frame()
            if self._next_frame is None:
                raise ValueError(self.last_error or "no complete video frame in UDP log")
        except (OSError, ValueError) as exc:
            self.last_error = str(exc)
            self.release()

    def _init_decoder(self, fps, width=0, height=0):
        self.fps = max(1.0, float(fps))
        self.width = max(0, int(width))
        self.height = max(0, int(height))
        self.pixel_id = 0
        self.row_padding = 0
        self.frames_read = 0
        self.dropped_arrays = 0
        self.last_labels = []
        self.last_error = ""
        self._contrast = None

    def isOpened(self):
        return self._file is not None

    def _packet(self):
        envelope = self._file.read(4)
        if not envelope:
            return None
        if len(envelope) != 4:
            raise ValueError("truncated UDP log envelope")
        port = int.from_bytes(envelope[:2], "little")
        size = int.from_bytes(envelope[2:4], "little")
        if self.port is None:
            self.port = port
        elif port != self.port:
            raise ValueError(f"unexpected UDP port: {port}")
        if size < self.PACKET_HEADER_SIZE:
            raise ValueError(f"invalid UDP payload size: {size}")
        payload = self._file.read(size)
        if len(payload) != size:
            raise ValueError("truncated UDP packet")
        return payload

    def _array(self):
        while True:
            payload = self._packet()
            if payload is None:
                return None
            dropped_before = self._assembler.dropped_arrays
            data = self._assembler.push(payload)
            self.dropped_arrays += self._assembler.dropped_arrays - dropped_before
            if data is not None:
                return data

    def _mono_to_bgr(self, raw, invalid_value=None):
        valid = np.ones(raw.shape, dtype=bool) if invalid_value is None else raw != invalid_value
        sample = raw[::4, ::4][valid[::4, ::4]]
        if sample.size < 16:
            sample = raw[valid]
        low, high = np.percentile(sample, (1.0, 99.0)) if sample.size else (0.0, 1.0)
        if high <= low:
            high = low + 1.0
        if self._contrast is None:
            self._contrast = (float(low), float(high))
        else:
            old_low, old_high = self._contrast
            self._contrast = (0.9 * old_low + 0.1 * low, 0.9 * old_high + 0.1 * high)
        low, high = self._contrast
        gray = np.clip((raw.astype(np.float32) - low) * (255.0 / (high - low)), 0, 255).astype(np.uint8)
        if not valid.all():
            median = cv2.medianBlur(gray, 3)
            gray[~valid] = median[~valid]
        return cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)

    def _decode_array(self, data):
        if len(data) < 12:
            raise ValueError("UDP data array is too short")
        label_count = int.from_bytes(data[:4], "little")
        video_offset = 4 + label_count * self.LABEL_SIZE
        if video_offset + 8 > len(data):
            raise ValueError("invalid label array size")
        self.last_labels = [
            data[4 + index * self.LABEL_SIZE:4 + (index + 1) * self.LABEL_SIZE]
            for index in range(label_count)
        ]

        header = data[video_offset:video_offset + 8]
        width = int.from_bytes(header[0:2], "little")
        height = int.from_bytes(header[2:4], "little")
        pixel_id = header[4]
        padding = header[6]
        bytes_per_pixel = {
            self.PIXEL_GRAY8: 1,
            self.PIXEL_GRAY16: 2,
            self.PIXEL_RGB888: 3,
            self.PIXEL_YCBCR422: 2,
            self.PIXEL_INT16: 2,
        }.get(pixel_id)
        if width <= 0 or height <= 0 or bytes_per_pixel is None:
            raise ValueError(f"unsupported video format: {width}x{height}, pixel_id=0x{pixel_id:02x}")

        row_bytes = width * bytes_per_pixel
        stride = row_bytes + padding
        pixels_offset = video_offset + 8
        pixels_end = pixels_offset + stride * height
        if pixels_end > len(data):
            raise ValueError("truncated video frame")
        rows = np.frombuffer(data[pixels_offset:pixels_end], dtype=np.uint8).reshape(height, stride)
        pixels = rows[:, :row_bytes].copy()

        if pixel_id == self.PIXEL_GRAY8:
            frame = self._mono_to_bgr(pixels.reshape(height, width))
        elif pixel_id == self.PIXEL_GRAY16:
            frame = self._mono_to_bgr(pixels.view("<u2").reshape(height, width), 0xFFFF)
        elif pixel_id == self.PIXEL_INT16:
            frame = self._mono_to_bgr(pixels.view("<i2").reshape(height, width), -1)
        elif pixel_id == self.PIXEL_RGB888:
            frame = cv2.cvtColor(pixels.reshape(height, width, 3), cv2.COLOR_RGB2BGR)
        else:
            frame = cv2.cvtColor(pixels.reshape(height, width, 2), cv2.COLOR_YUV2BGR_YUY2)

        self.width = width
        self.height = height
        self.pixel_id = pixel_id
        self.row_padding = padding
        return frame

    def _read_frame(self):
        while True:
            data = self._array()
            if data is None:
                return None
            try:
                return self._decode_array(data)
            except ValueError as exc:
                self.last_error = str(exc)
                self.dropped_arrays += 1

    def read(self):
        if self._file is None:
            return False, None
        try:
            frame = self._next_frame
            self._next_frame = None
            if frame is None:
                frame = self._read_frame()
            if frame is None:
                return False, None
            self.frames_read += 1
            return True, frame
        except (OSError, ValueError) as exc:
            self.last_error = str(exc)
            self.release()
            return False, None

    def get(self, prop):
        if prop == cv2.CAP_PROP_FRAME_WIDTH:
            return float(self.width)
        if prop == cv2.CAP_PROP_FRAME_HEIGHT:
            return float(self.height)
        if prop == cv2.CAP_PROP_FPS:
            return self.fps
        if prop == cv2.CAP_PROP_POS_FRAMES:
            return float(self.frames_read)
        if prop == cv2.CAP_PROP_POS_MSEC:
            return 1000.0 * self.frames_read / self.fps
        return 0.0

    def set(self, _prop, _value):
        return False

    def release(self):
        if self._file is not None:
            self._file.close()
            self._file = None


class LiveMikUdpCapture(UdpDumpCapture):
    """Live UDP receiver for the MIK packet payload used by port 59004."""

    def __init__(self, host="0.0.0.0", port=59004, fps=30.0, width=0, height=0):
        self.path = None
        self.host = str(host)
        self.port = int(port)
        self._init_decoder(fps, width, height)
        self._file = None
        self._next_frame = None
        self._assembler = MikPacketAssembler()
        self._socket = None
        try:
            self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self._socket.bind((self.host, self.port))
            self.port = int(self._socket.getsockname()[1])
            self._socket.settimeout(0.5)
        except OSError as exc:
            self.last_error = str(exc)
            self.release()

    def isOpened(self):
        return self._socket is not None

    def read(self):
        while self._socket is not None:
            try:
                payload, _address = self._socket.recvfrom(65535)
            except socket.timeout:
                continue
            except OSError as exc:
                self.last_error = str(exc)
                return False, None

            dropped_before = self._assembler.dropped_arrays
            try:
                data = self._assembler.push(payload)
            except ValueError as exc:
                self.last_error = str(exc)
                continue
            self.dropped_arrays += self._assembler.dropped_arrays - dropped_before
            if data is None:
                continue
            try:
                frame = self._decode_array(data)
            except ValueError as exc:
                self.last_error = str(exc)
                self.dropped_arrays += 1
                continue
            self.frames_read += 1
            return True, frame
        return False, None

    def release(self):
        sock, self._socket = getattr(self, "_socket", None), None
        if sock is not None:
            sock.close()