vogelalbum/detector/stream.py

"""Framequellen.

Der Pi Zero kann auf mehrere Arten streamen und wir wissen noch nicht, welche
Software drauf läuft. Deshalb vier Quellen hinter einem Interface — die
Pipeline interessiert sich nur für ``for frame in source``.

MJPEG und Snapshot brauchen nur die Stdlib. RTSP ruft ffmpeg auf; das muss
dann auf dem Host installiert sein (``apt install ffmpeg``).
"""

from __future__ import annotations

import io
import subprocess
import time
import urllib.request
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Iterator

from PIL import Image

JPEG_SOI = b"\xff\xd8"
JPEG_EOI = b"\xff\xd9"


class FrameSource(ABC):
    restartable = True
    """Ob es sinnvoll ist, die Quelle nach dem Ende neu zu öffnen.

    Bei Netzquellen ja — der Pi startet neu, das WLAN hustet. Bei einem
    endlichen Ordner nein, sonst läuft ``watch`` in einer Endlosschleife.
    """

    letztes_jpeg: bytes | None = None
    """Die unveränderten JPEG-Bytes des zuletzt gelieferten Bildes.

    Wer das Bild weitergeben will — etwa der Live-Verteiler —, spart damit das
    Umkodieren: bei MJPEG und Einzelbild-URLs liegen die Bytes ohnehin vor.
    Quellen ohne fertiges JPEG lassen das Feld auf ``None``.
    """

    @abstractmethod
    def frames(self) -> Iterator[Image.Image]: ...

    def __iter__(self) -> Iterator[Image.Image]:
        return self.frames()


class MJPEGStream(FrameSource):
    """multipart/x-mixed-replace — was mjpg-streamer, motion und die meisten
    picamera2-Beispiele ausliefern.

    Wir suchen die JPEG-Marker direkt im Bytestrom, statt die MIME-Boundaries
    zu parsen: Boundary-Formate sind zwischen den Servern uneinheitlich, die
    Marker sind es nicht.
    """

    def __init__(self, url: str, timeout: int = 15, max_buffer: int = 8 << 20):
        self.url = url
        self.timeout = timeout
        self.max_buffer = max_buffer

    def frames(self) -> Iterator[Image.Image]:
        req = urllib.request.Request(self.url, headers={"User-Agent": "vogelalbum/1.0"})
        with urllib.request.urlopen(req, timeout=self.timeout) as resp:
            buf = bytearray()
            while True:
                chunk = resp.read(16384)
                if not chunk:
                    return
                buf.extend(chunk)

                while True:
                    start = buf.find(JPEG_SOI)
                    if start < 0:
                        break
                    end = buf.find(JPEG_EOI, start + 2)
                    if end < 0:
                        break
                    raw = bytes(buf[start : end + 2])
                    del buf[: end + 2]
                    try:
                        bild = Image.open(io.BytesIO(raw)).convert("RGB")
                    except OSError:
                        continue  # halbes Bild, nächstes nehmen
                    self.letztes_jpeg = raw
                    yield bild

                # Reißt die Verbindung mitten im Frame ab, wächst der Puffer
                # sonst unbegrenzt weiter.
                if len(buf) > self.max_buffer:
                    del buf[: -(1 << 20)]


class SnapshotPoller(FrameSource):
    """Einzelbild-URL im Intervall abholen (z.B. /snapshot.jpg bei motion,
    oder ein Cron-Job, der auf den Pi schreibt). Genügsamer als MJPEG und für
    einen Nistkasten völlig ausreichend."""

    def __init__(self, url: str, interval: float = 1.0, timeout: int = 10):
        self.url = url
        self.interval = interval
        self.timeout = timeout

    def frames(self) -> Iterator[Image.Image]:
        req = urllib.request.Request(self.url, headers={"User-Agent": "vogelalbum/1.0"})
        while True:
            t0 = time.monotonic()
            try:
                with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                    raw = resp.read()
                bild = Image.open(io.BytesIO(raw)).convert("RGB")
            except (OSError, ValueError) as exc:
                print(f"[stream] Snapshot fehlgeschlagen: {exc}")
                time.sleep(min(self.interval * 5, 30))
                continue
            self.letztes_jpeg = raw
            yield bild
            time.sleep(max(0.0, self.interval - (time.monotonic() - t0)))


class RTSPStream(FrameSource):
    """RTSP über ffmpeg. Braucht ffmpeg auf dem Host."""

    def __init__(self, url: str, fps: float = 2.0, width: int = 1280):
        self.url = url
        self.fps = fps
        self.width = width

    def frames(self) -> Iterator[Image.Image]:
        cmd = [
            "ffmpeg", "-loglevel", "error",
            "-rtsp_transport", "tcp", "-i", self.url,
            "-vf", f"fps={self.fps},scale={self.width}:-2",
            "-f", "image2pipe", "-vcodec", "mjpeg", "-q:v", "4", "-",
        ]
        try:
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, bufsize=1 << 16)
        except FileNotFoundError:
            raise RuntimeError("ffmpeg nicht gefunden — 'apt install ffmpeg'") from None

        try:
            buf = bytearray()
            assert proc.stdout is not None
            while True:
                chunk = proc.stdout.read(16384)
                if not chunk:
                    return
                buf.extend(chunk)
                while True:
                    start = buf.find(JPEG_SOI)
                    end = buf.find(JPEG_EOI, start + 2) if start >= 0 else -1
                    if start < 0 or end < 0:
                        break
                    raw = bytes(buf[start : end + 2])
                    del buf[: end + 2]
                    try:
                        yield Image.open(io.BytesIO(raw)).convert("RGB")
                    except OSError:
                        continue
        finally:
            proc.terminate()
            try:
                proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                proc.kill()


class FolderSource(FrameSource):
    """Bilder aus einem Ordner — zum Entwickeln und für Regressionstests ohne
    laufende Kamera."""

    def __init__(self, folder: str | Path, loop: bool = False, delay: float = 0.0):
        self.folder = Path(folder)
        self.loop = loop
        self.delay = delay
        self.restartable = loop

    def frames(self) -> Iterator[Image.Image]:
        exts = {".jpg", ".jpeg", ".png", ".webp"}
        files = sorted(p for p in self.folder.iterdir() if p.suffix.lower() in exts)
        if not files:
            raise RuntimeError(f"Keine Bilder in {self.folder}")
        while True:
            for p in files:
                yield Image.open(p).convert("RGB")
                if self.delay:
                    time.sleep(self.delay)
            if not self.loop:
                return


def open_source(url: str, **kw) -> FrameSource:
    """Baut die passende Quelle aus einer URL bzw. einem Pfad."""
    if url.startswith("rtsp://"):
        return RTSPStream(url, fps=kw.get("fps", 2.0))
    if url.startswith(("http://", "https://")):
        # Heuristik: eine URL auf .jpg ist ein Einzelbild, alles andere ein Stream.
        path = url.split("?", 1)[0].lower()
        if path.endswith((".jpg", ".jpeg", ".png")) or kw.get("snapshot"):
            return SnapshotPoller(url, interval=kw.get("interval", 1.0))
        return MJPEGStream(url)
    return FolderSource(url, loop=kw.get("loop", False), delay=kw.get("delay", 0.0))