MAI/ffmpeg_capture.py

import json
import shutil
import subprocess

import cv2
import numpy as np


def _parse_rate(value):
    try:
        numerator, denominator = str(value).split("/", 1)
        denominator = float(denominator)
        return float(numerator) / denominator if denominator else 0.0
    except (TypeError, ValueError):
        return 0.0


def _parse_int(value):
    try:
        return int(value)
    except (TypeError, ValueError):
        return 0


def _parse_float(value):
    try:
        return float(value)
    except (TypeError, ValueError):
        return 0.0


class FFmpegCapture:
    """Sequential file reader using FFmpeg's more tolerant decoder."""

    def __init__(self, source):
        self.source = str(source)
        self.width = 0
        self.height = 0
        self.fps = 0.0
        self.frame_count = 0
        self.frames_read = 0
        self.process = None

        if not shutil.which("ffmpeg") or not shutil.which("ffprobe"):
            return
        try:
            probe = subprocess.run(
                [
                    "ffprobe",
                    "-v",
                    "error",
                    "-select_streams",
                    "v:0",
                    "-show_entries",
                    "stream=width,height,avg_frame_rate,nb_frames,duration",
                    "-of",
                    "json",
                    self.source,
                ],
                check=True,
                capture_output=True,
                text=True,
                timeout=10,
            )
            streams = json.loads(probe.stdout).get("streams") or []
            if not streams:
                raise ValueError("no video stream")
            stream = streams[0]
            self.width = _parse_int(stream.get("width"))
            self.height = _parse_int(stream.get("height"))
            if self.width <= 0 or self.height <= 0:
                raise ValueError("invalid video dimensions")
            self.fps = _parse_rate(stream.get("avg_frame_rate"))
            self.frame_count = _parse_int(stream.get("nb_frames"))
            if not self.frame_count and self.fps > 0.0:
                self.frame_count = int(round(_parse_float(stream.get("duration")) * self.fps))
            self._start()
        except (KeyError, ValueError, OSError, subprocess.SubprocessError, json.JSONDecodeError):
            self.release()

    def _start(self):
        self.process = subprocess.Popen(
            [
                "ffmpeg",
                "-hide_banner",
                "-loglevel",
                "fatal",
                "-err_detect",
                "ignore_err",
                "-i",
                self.source,
                "-map",
                "0:v:0",
                "-an",
                "-sn",
                "-dn",
                "-vsync",
                "0",
                "-pix_fmt",
                "bgr24",
                "-f",
                "rawvideo",
                "pipe:1",
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            bufsize=max(1024 * 1024, self.width * self.height * 3),
        )

    def isOpened(self):
        return (
            self.process is not None
            and self.process.stdout is not None
            and not self.process.stdout.closed
        )

    def read(self):
        if not self.isOpened() or self.process.stdout is None:
            return False, None
        expected = self.width * self.height * 3
        data = bytearray()
        while len(data) < expected:
            chunk = self.process.stdout.read(expected - len(data))
            if not chunk:
                return False, None
            data.extend(chunk)
        self.frames_read += 1
        frame = np.frombuffer(data, dtype=np.uint8).reshape(self.height, self.width, 3)
        return True, frame

    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 float(self.fps)
        if prop == cv2.CAP_PROP_FRAME_COUNT:
            return float(self.frame_count)
        if prop == cv2.CAP_PROP_POS_FRAMES:
            return float(self.frames_read)
        if prop == cv2.CAP_PROP_POS_MSEC and self.fps > 0.0:
            return 1000.0 * self.frames_read / self.fps
        return 0.0

    def set(self, prop, value):
        if prop == cv2.CAP_PROP_POS_FRAMES and int(value) == 0:
            self.release()
            self.frames_read = 0
            self._start()
            return self.isOpened()
        return prop == cv2.CAP_PROP_BUFFERSIZE

    def release(self):
        process, self.process = self.process, None
        if process is None:
            return
        if process.stdout is not None:
            process.stdout.close()
        if process.poll() is None:
            process.terminate()
            try:
                process.wait(timeout=1)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait(timeout=1)