MAI/delimited_frame_capture.py

import socket
from pathlib import Path

import cv2
import numpy as np


RAW_FRAME_CHANNELS = {
    "bgr24": 3,
    "rgb24": 3,
    "gray8": 1,
    "gray16": 2,
    "yuyv422": 2,
}


def raw_frame_size(encoding, width, height):
    channels = RAW_FRAME_CHANNELS.get(str(encoding or "").lower())
    return max(1, int(width)) * max(1, int(height)) * channels if channels else 0


def decode_frame_data(data, encoding, width, height):
    encoding = str(encoding or "auto").lower()
    width = max(1, int(width))
    height = max(1, int(height))
    raw = np.frombuffer(data, dtype=np.uint8)
    if encoding == "auto":
        frame = cv2.imdecode(raw, cv2.IMREAD_COLOR)
        if frame is None:
            raise ValueError("frame is not JPEG/PNG")
        return frame

    shapes = {
        "bgr24": (height, width, 3),
        "rgb24": (height, width, 3),
        "gray8": (height, width),
        "gray16": (height, width),
        "yuyv422": (height, width, 2),
    }
    if encoding not in shapes:
        raise ValueError(f"unsupported frame encoding: {encoding}")
    dtype = np.dtype("<u2") if encoding == "gray16" else np.dtype(np.uint8)
    expected = int(np.prod(shapes[encoding]) * dtype.itemsize)
    if len(data) != expected:
        raise ValueError(f"raw frame size {len(data)} != {expected}")
    pixels = np.frombuffer(data, dtype=dtype).reshape(shapes[encoding])
    if encoding == "bgr24":
        return pixels.copy()
    if encoding == "rgb24":
        return cv2.cvtColor(pixels, cv2.COLOR_RGB2BGR)
    if encoding == "yuyv422":
        return cv2.cvtColor(pixels, cv2.COLOR_YUV2BGR_YUY2)
    if encoding == "gray16":
        pixels = (pixels >> 8).astype(np.uint8)
    return cv2.cvtColor(pixels, cv2.COLOR_GRAY2BGR)


class DelimitedFrameCapture:
    """Reads encoded or raw frames separated by one byte."""

    MAX_FRAME_SIZE = 64 * 1024 * 1024

    def __init__(
        self,
        source=None,
        host="0.0.0.0",
        port=59005,
        separator=0,
        encoding="auto",
        width=1280,
        height=720,
        fps=30.0,
    ):
        self.source = Path(source) if source is not None else None
        self.host = str(host)
        self.port = int(port)
        self.separator = bytes((int(separator) & 0xFF,))
        self.encoding = str(encoding or "auto").lower()
        self.width = max(1, int(width))
        self.height = max(1, int(height))
        self.fps = max(1.0, float(fps))
        self.frames_read = 0
        self.dropped_frames = 0
        self.last_error = ""
        self._buffer = bytearray()
        self._frames = []
        self._file = None
        self._socket = None
        try:
            if self.source is not None:
                self._file = self.source.open("rb", buffering=8 * 1024 * 1024)
            else:
                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._file is not None or self._socket is not None

    def _next_chunk(self):
        if self._file is not None:
            return self._file.read(1024 * 1024)
        while self._socket is not None:
            try:
                return self._socket.recvfrom(65535)[0]
            except socket.timeout:
                continue
            except OSError as exc:
                self.last_error = str(exc)
                return b""
        return b""

    def _split(self, chunk):
        expected = self._raw_frame_size()
        if self._socket is not None and expected:
            if chunk == self.separator:
                if self._buffer:
                    self._buffer.clear()
                    self.dropped_frames += 1
                return
            self._buffer.extend(chunk)
            while len(self._buffer) >= expected:
                self._frames.append(bytes(self._buffer[:expected]))
                del self._buffer[:expected]
                if self._buffer.startswith(self.separator):
                    del self._buffer[:1]
            if len(self._buffer) > self.MAX_FRAME_SIZE:
                self._buffer.clear()
                self.dropped_frames += 1
            return

        self._buffer.extend(chunk)
        parts = self._buffer.split(self.separator)
        self._buffer = bytearray(parts.pop())
        self._frames.extend(part for part in parts if part)
        if len(self._buffer) > self.MAX_FRAME_SIZE:
            self._buffer.clear()
            self.dropped_frames += 1

    def _raw_frame_size(self):
        return raw_frame_size(self.encoding, self.width, self.height)

    def _decode(self, data):
        frame = decode_frame_data(data, self.encoding, self.width, self.height)
        self.height, self.width = frame.shape[:2]
        return frame

    def read(self):
        while self.isOpened():
            if not self._frames:
                chunk = self._next_chunk()
                if not chunk:
                    if self._file is not None and self._buffer:
                        self._frames.append(bytes(self._buffer))
                        self._buffer.clear()
                    else:
                        return False, None
                else:
                    self._split(chunk)
                    continue
            data = self._frames.pop(0)
            try:
                frame = self._decode(data)
            except ValueError as exc:
                self.last_error = str(exc)
                self.dropped_frames += 1
                continue
            self.frames_read += 1
            return True, frame
        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
        sock, self._socket = self._socket, None
        if sock is not None:
            sock.close()