nedo-vision-worker-core 0.3.3__py3-none-any.whl → 0.3.4__py3-none-any.whl

nedo_vision_worker_core/streams/RTMPStreamer.py

@@ -4,281 +4,226 @@ import threading
 import time
 import numpy as np
 import os
+import sys
+import cv2
+import queue
 from typing import Optional
+from ..util.PlatformDetector import PlatformDetector
 
 class RTMPStreamer:
     """
-    Streams raw BGR frames to FFmpeg -> RTMP with:
-      - Internal pacing thread (exact fps)
-      - Latest-only frame buffer (no backlog)
-      - Resilient auto-restart on failure
+    Streams raw BGR frames to an RTMP server using a GStreamer-first approach
+    with a reliable FFmpeg subprocess fallback.
     """
 
     def __init__(self, pipeline_id: str, fps: int = 25, bitrate: str = "1500k"):
         self.rtmp_server = os.environ.get("RTMP_SERVER", "rtmp://localhost:1935/live")
         self.rtmp_url = f"{self.rtmp_server}/{pipeline_id}"
         self.fps = max(int(fps), 1)
-        self.bitrate = bitrate  # e.g. "1500k"
-
-        # VBV aligned to target bitrate by default (override via env if needed)
-        self.maxrate = os.environ.get("RTMP_MAXRATE", self.bitrate)
-        self.bufsize = os.environ.get("RTMP_BUFSIZE", f"{self._kbps(self.bitrate) * 2}k")
+        self.bitrate = self._kbps(bitrate) # Store as integer kbps
 
         self.width: Optional[int] = None
         self.height: Optional[int] = None
-        self.ffmpeg_process: Optional[subprocess.Popen] = None
-        self.started = False
-        self.active = False
+        self._platform = PlatformDetector()
+        
+        self._backend = None # "gstreamer" or "ffmpeg"
+        self._gstreamer_writer: Optional[cv2.VideoWriter] = None
+        self._ffmpeg_process: Optional[subprocess.Popen] = None
 
-        # Writer thread & latest-only buffer
+        self._frame_queue = queue.Queue(maxsize=2)
         self._writer_thread: Optional[threading.Thread] = None
         self._stop_evt = threading.Event()
-        self._lock = threading.Lock()
-        self._latest_frame: Optional[np.ndarray] = None  # last BGR frame set by push_frame()
 
-        # pacing
-        self._frame_period = 1.0 / self.fps
+    def _kbps(self, rate_str: str) -> int:
+        return int(str(rate_str).lower().replace("k", "").strip())
 
-    # -------------------- public API --------------------
+    # -------------------- Public API --------------------
 
     def is_active(self) -> bool:
-        return bool(self.active and self.ffmpeg_process and self.ffmpeg_process.poll() is None)
+        if self._backend == "gstreamer":
+            return self._gstreamer_writer is not None and self._gstreamer_writer.isOpened()
+        if self._backend == "ffmpeg":
+            return self._ffmpeg_process is not None and self._ffmpeg_process.poll() is None
+        return False
 
     def push_frame(self, frame: np.ndarray):
-        """
-        Provide a frame to the streamer. The internal writer thread will pick up
-        the latest frame at the correct fps. Older frames are dropped.
-        """
-        if frame is None or getattr(frame, "size", 0) == 0:
-            return
-        if frame.ndim != 3 or frame.shape[2] != 3:
+        if self._stop_evt.is_set():
             return
 
-        # On the very first frame, learn WxH and start ffmpeg + writer
-        if not self.started:
-            h, w = frame.shape[:2]
-            if w <= 0 or h <= 0:
-                return
-            self.width, self.height = w, h
-            self._start_ffmpeg_and_writer()
+        if self._writer_thread is None or not self._writer_thread.is_alive():
+            if frame is None: return
+            self.height, self.width = frame.shape[:2]
+            self._start_stream()
 
-        # Store latest only
-        with self._lock:
-            if frame.dtype != np.uint8:
-                frame = frame.astype(np.uint8, copy=False)
-            if not frame.flags['C_CONTIGUOUS']:
-                frame = np.ascontiguousarray(frame)
-            self._latest_frame = frame
-
-    def stop_stream(self):
-        """Stop writer thread and FFmpeg cleanly."""
-        self.active = False
-        self._stop_evt.set()
-
-        if self._writer_thread and self._writer_thread.is_alive():
+        try:
+            self._frame_queue.put_nowait(frame)
+        except queue.Full:
             try:
-                self._writer_thread.join(timeout=2.0)
-            except Exception as e:
-                logging.error(f"RTMP writer thread join error: {e}")
-        self._writer_thread = None
+                self._frame_queue.get_nowait()
+                self._frame_queue.put_nowait(frame)
+            except queue.Empty:
+                pass
 
-        # tear down ffmpeg
-        proc = self.ffmpeg_process
-        self.ffmpeg_process = None
-        if not proc:
-            logging.info("RTMP streaming process already stopped.")
+    def stop_stream(self):
+        if self._stop_evt.is_set():
             return
-
+        logging.info(f"Stopping RTMP stream for {self.rtmp_url}")
+        self._stop_evt.set()
+        
         try:
-            if proc.stdin:
-                try:
-                    proc.stdin.flush()
-                except Exception:
-                    pass
-                try:
-                    proc.stdin.close()
-                except Exception:
-                    pass
-            proc.terminate()
-            proc.wait(timeout=5)
-        except Exception as e:
-            logging.error(f"Error stopping RTMP stream: {e}")
+            self._frame_queue.put_nowait(None)
+        except queue.Full:
+            pass
+
+        if self._writer_thread and self._writer_thread.is_alive() and threading.current_thread() != self._writer_thread:
+            self._writer_thread.join(timeout=2.0)
+
+        if self._gstreamer_writer:
+            self._gstreamer_writer.release()
+            self._gstreamer_writer = None
+            logging.info("GStreamer writer released.")
+        
+        if self._ffmpeg_process:
             try:
-                proc.kill()
+                if self._ffmpeg_process.stdin: self._ffmpeg_process.stdin.close()
+                self._ffmpeg_process.terminate()
+                self._ffmpeg_process.wait(timeout=2.0)
             except Exception:
-                pass
-        finally:
-            logging.info("RTMP streaming process stopped.")
-            self.started = False
-
-    # -------------------- internals --------------------
-
-    def _kbps(self, rate_str: str) -> int:
-        return int(str(rate_str).lower().replace("k", "").strip())
-
-    def _build_ffmpeg_cmd(self) -> list[str]:
-        # GOP = 1 second for faster recovery; disable scene-cut; repeat headers
-        gop = max(self.fps, 1)
-        return [
-            "ffmpeg",
-            "-y",
-            "-loglevel", os.environ.get("FFLOG", "error"),
-            "-nostats",
-            "-hide_banner",
-
-            # raw frames via stdin
-            "-f", "rawvideo",
-            "-pixel_format", "bgr24",
-            "-video_size", f"{self.width}x{self.height}",
-            "-framerate", str(self.fps),
-            "-i", "-",
-
-            # encoder
-            "-c:v", "libx264",
-            "-tune", "zerolatency",
-            "-preset", "ultrafast",
-            "-profile:v", "main",
-            "-pix_fmt", "yuv420p",
-
-            # CBR-like VBV
-            "-b:v", self.bitrate,
-            "-maxrate", self.maxrate,
-            "-bufsize", self.bufsize,
-
-            # GOP / keyframes
-            "-g", str(gop),
-            "-keyint_min", str(gop),
-            "-sc_threshold", "0",
-            "-x264-params", "open_gop=0:aud=1:repeat-headers=1:nal-hrd=cbr",
-            # force an IDR every 1s for faster player join/recovery
-            "-force_key_frames", "expr:gte(t,n_forced*1)",
-
-            # single scaling location (if needed downstream)
-            "-vf", "scale='min(1024,iw)':-2",
-
-            # no audio
-            "-an",
-
-            # reduce container buffering
-            "-flvflags", "no_duration_filesize",
-            "-flush_packets", "1",
-            "-rtmp_live", "live",
-            "-muxpreload", "0",
-            "-muxdelay", "0",
-
-            "-f", "flv",
-            self.rtmp_url,
-        ]
+                if self._ffmpeg_process: self._ffmpeg_process.kill()
+            self._ffmpeg_process = None
+            logging.info("FFmpeg process stopped.")
+        
+        self._backend = None
+
+    # -------------------- Internal Stream Management --------------------
+
+    def _start_stream(self):
+        self._stop_evt.clear()
+        
+        gstreamer_pipeline = self._build_gstreamer_pipeline()
+        self._gstreamer_writer = cv2.VideoWriter(gstreamer_pipeline, cv2.CAP_GSTREAMER, 0, self.fps, (self.width, self.height))
+
+        if self._gstreamer_writer.isOpened():
+            self._backend = "gstreamer"
+            self._writer_thread = threading.Thread(target=self._gstreamer_writer_loop, daemon=True)
+            self._writer_thread.start()
+            logging.info(f"RTMP streaming started with GStreamer (HW-Accel): {self.rtmp_url}")
+            return
+        
+        logging.warning("GStreamer VideoWriter failed to open. Falling back to FFmpeg subprocess.")
+        self._gstreamer_writer = None
 
-    def _start_ffmpeg_and_writer(self):
-        """Start ffmpeg process and the pacing writer thread."""
         cmd = self._build_ffmpeg_cmd()
         try:
-            with open(os.devnull, "w") as devnull:
-                self.ffmpeg_process = subprocess.Popen(
-                    cmd,
-                    stdin=subprocess.PIPE,
-                    stdout=devnull,
-                    stderr=devnull if os.environ.get("FFSILENT", "1") == "1" else None,
-                    bufsize=0,  # unbuffered for low latency
-                )
-            self.started = True
-            self.active = True
-            self._stop_evt.clear()
-            self._writer_thread = threading.Thread(
-                target=self._writer_loop, name=f"rtmp-writer-{os.getpid()}",
-                daemon=True
-            )
+            self._ffmpeg_process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+            self._backend = "ffmpeg"
+            self._writer_thread = threading.Thread(target=self._ffmpeg_pacing_loop, daemon=True)
             self._writer_thread.start()
-            logging.info(f"RTMP streaming started: {self.rtmp_url} ({self.width}x{self.height}@{self.fps}fps)")
+            logging.info(f"RTMP streaming started with FFmpeg (Fallback): {self.rtmp_url}")
         except Exception as e:
-            logging.error(f"Failed to start FFmpeg: {e}")
-            self.ffmpeg_process = None
-            self.active = False
-            self.started = False
-
-    def _writer_loop(self):
-        """Paces frames at exact fps, writing latest frame only. Auto-restarts on failure."""
-        next_deadline = time.monotonic() + self._frame_period
-        idle_frame: Optional[np.ndarray] = None  # reuse last good frame if no new one arrived
+            logging.error(f"Failed to start FFmpeg fallback: {e}")
+            self._ffmpeg_process = None
 
+    def _gstreamer_writer_loop(self):
         while not self._stop_evt.is_set():
             try:
-                # pacing
-                now = time.monotonic()
-                sleep_for = next_deadline - now
-                if sleep_for > 0:
-                    time.sleep(sleep_for)
-                next_deadline += self._frame_period
-                if next_deadline < now - self._frame_period:
-                    next_deadline = now + self._frame_period
-
-                # get the freshest frame
-                with self._lock:
-                    frame = self._latest_frame
-
+                frame = self._frame_queue.get(timeout=1.0)
                 if frame is None:
-                    # no new frame—if we have an idle one, repeat it (prevents RTMP underflow)
-                    frame_to_send = idle_frame
-                    if frame_to_send is None:
-                        continue  # nothing to send yet
-                else:
-                    frame_to_send = frame
-                    idle_frame = frame
-
-                if not self.is_active():
-                    raise BrokenPipeError("FFmpeg not active")
-
-                # write raw BGR24
-                self.ffmpeg_process.stdin.write(frame_to_send.tobytes())
-
-            except BrokenPipeError:
-                # logging.error("RTMP pipe broken. Restarting encoder.")
-                self._restart_ffmpeg()
-                # On restart we keep width/height; writer will continue
+                    break
+                self._gstreamer_writer.write(frame)
+            except queue.Empty:
+                continue
             except Exception as e:
-                # Any unexpected error: try a restart, but do not spin
-                # logging.error(f"RTMP writer error: {e}")
-                time.sleep(0.1)
-                self._restart_ffmpeg()
+                logging.error(f"Error in GStreamer writer loop: {e}. Stopping thread.")
+                self._stop_evt.set()
+                break
+
+    # -------------------- Pipeline Builders --------------------
+
+    def _build_gstreamer_pipeline(self) -> str:
+        if self._platform.is_jetson():
+            encoder = "nvv4l2h264enc insert-sps-pps=true"
+            converter = "nvvidconv ! video/x-raw(memory:NVMM) ! "
+        else:
+            encoder = "nvh264enc preset=low-latency-hq"
+            converter = "videoconvert"
+
+        pipeline = (
+            f"appsrc ! "
+            f"video/x-raw,format=BGR,width={self.width},height={self.height},framerate={self.fps}/1 ! "
+            f"{converter} ! {encoder} bitrate={self.bitrate} ! "
+            f"h264parse ! flvmux ! "
+            f"rtmpsink location=\"{self.rtmp_url}\""
+        )
+        return pipeline
+
+    # --- FFmpeg Fallback Methods ---
+    
+    def _ffmpeg_pacing_loop(self):
+        frame_period = 1.0 / self.fps
+        last_frame_sent = None
 
-    def _restart_ffmpeg(self):
-        """Restart ffmpeg while keeping width/height and thread alive."""
-        self.active = False
-        proc = self.ffmpeg_process
-        self.ffmpeg_process = None
-        if proc:
+        while not self._stop_evt.is_set():
+            start_time = time.monotonic()
+            
             try:
-                if proc.stdin:
-                    try:
-                        proc.stdin.flush()
-                    except Exception:
-                        pass
-                    try:
-                        proc.stdin.close()
-                    except Exception:
-                        pass
-                proc.terminate()
-                proc.wait(timeout=3)
-            except Exception:
-                try:
-                    proc.kill()
-                except Exception:
-                    pass
+                frame = self._frame_queue.get_nowait()
+                last_frame_sent = frame
+            except queue.Empty:
+                frame = last_frame_sent
+
+            if frame is None:
+                time.sleep(frame_period)
+                continue
 
-        # restart only if we have size info
-        if self.width and self.height:
             try:
-                with open(os.devnull, "w") as devnull:
-                    self.ffmpeg_process = subprocess.Popen(
-                        self._build_ffmpeg_cmd(),
-                        stdin=subprocess.PIPE,
-                        stdout=devnull,
-                        stderr=devnull if os.environ.get("FFSILENT", "1") == "1" else None,
-                        bufsize=0,
-                    )
-                self.active = True
-                # logging.info("RTMP encoder restarted.")
-            except Exception as e:
-                logging.error(f"Failed to restart FFmpeg: {e}")
-                self.active = False
+                if not self.is_active():
+                    raise BrokenPipeError("FFmpeg process is not active")
+                self._ffmpeg_process.stdin.write(frame.tobytes())
+            except (BrokenPipeError, OSError) as e:
+                logging.error(f"FFmpeg process pipe broken: {e}. Stopping thread.")
+                self._stop_evt.set()
+                break
+            
+            elapsed = time.monotonic() - start_time
+            sleep_duration = max(0, frame_period - elapsed)
+            time.sleep(sleep_duration)
+
+    def _build_ffmpeg_cmd(self) -> list[str]:
+        encoder_args = self._select_ffmpeg_encoder()
+        encoder_name = encoder_args[1]
+
+        # Base command arguments
+        cmd = [
+            'ffmpeg', '-y', '-loglevel', 'error', '-nostats', '-hide_banner',
+            '-f', 'rawvideo', '-pixel_format', 'bgr24',
+            '-video_size', f'{self.width}x{self.height}',
+            '-framerate', str(self.fps), '-i', '-',
+        ]
+
+        # Add the selected encoder
+        cmd.extend(encoder_args)
+
+        # Add common arguments for all encoders
+        cmd.extend([
+            '-profile:v', 'main', '-pix_fmt', 'yuv420p',
+            '-b:v', f"{self.bitrate}k", '-maxrate', f"{self.bitrate}k", '-bufsize', f"{self.bitrate*2}k",
+            '-g', str(self.fps), '-keyint_min', str(self.fps),
+            '-force_key_frames', 'expr:gte(t,n_forced*1)', '-an',
+            '-flvflags', 'no_duration_filesize', '-f', 'flv', self.rtmp_url,
+        ])
+        
+        # Conditionally add arguments that are ONLY valid for the libx264 encoder
+        if encoder_name == "libx264":
+            cmd.extend([
+                "-tune", "zerolatency",
+                "-x264-params", "open_gop=0:aud=1:repeat-headers=1:nal-hrd=cbr",
+            ])
+        
+        return cmd
+
+    def _select_ffmpeg_encoder(self) -> list:
+        if sys.platform == "darwin": return ["-c:v", "h264_videotoolbox"]
+        if os.environ.get("NVIDIA_VISIBLE_DEVICES") is not None or os.path.exists("/proc/driver/nvidia/version"):
+            return ["-c:v", "h264_nvenc", "-preset", "llhp"]
+        return ["-c:v", "libx264"]

nedo_vision_worker_core/streams/SharedVideoDeviceManager.py

@@ -8,6 +8,10 @@ from .VideoStream import VideoStream
 from ..services.SharedVideoStreamServer import get_shared_stream_server
 from ..services.VideoSharingDaemonManager import get_daemon_manager
 
+import numpy as np
+from numpy.typing import NDArray
+MatLike = NDArray[np.uint8] 
+
 try:
     from nedo_vision_worker_core.services.VideoSharingDaemon import VideoSharingClient
 except ImportError:
@@ -77,7 +81,7 @@ class SharedVideoDeviceManager:
                 pass
         return False, None
     
-    def subscribe_to_device(self, source, subscriber_id: str, callback: Callable[[cv2.Mat], None]) -> bool:
+    def subscribe_to_device(self, source, subscriber_id: str, callback: Callable[[MatLike], None]) -> bool:
         """
         Subscribe to a direct video device.