diff --git a/examples/01_vtk/vtk_cone_simple_video.py b/examples/01_vtk/vtk_cone_simple_video.py index 190f075..2384e98 100644 --- a/examples/01_vtk/vtk_cone_simple_video.py +++ b/examples/01_vtk/vtk_cone_simple_video.py @@ -22,6 +22,11 @@ class ConeApp(TrameApp): def __init__(self, server=None): super().__init__(server) + + self.server.cli.add_argument("--encoder", default="auto") + args, _ = self.server.cli.parse_known_args() + + self.encoder = args.encoder self.render_window = self.setup_vtk() self._build_ui() @@ -76,7 +81,9 @@ def _build_ui(self): with layout.content: view = rca.RemoteControlledArea(display="video-decoder") - self.view_handler = view.create_view_handler(self.render_window) + self.view_handler = view.create_view_handler( + self.render_window, encoder=self.encoder + ) # ----------------------------------------------------------------------------- diff --git a/src/trame_rca/encoders/__init__.py b/src/trame_rca/encoders/__init__.py index 4735fa4..30e0562 100644 --- a/src/trame_rca/encoders/__init__.py +++ b/src/trame_rca/encoders/__init__.py @@ -3,12 +3,14 @@ from .image_encoder import RcaImageEncoder try: - from .video_encoder import RcaVideoEncoder + from .video_encoder import RcaVideoEncoder, VideoEncoderType except ModuleNotFoundError: + RcaVideoEncoder = object + VideoEncoderType = object logger = logging.getLogger(__name__) logger.warning( "VTKStreaming Video encoding is NOT AVAILABLE (missing Python package)" ) -__all__ = ["RcaVideoEncoder", "RcaImageEncoder"] +__all__ = ["RcaVideoEncoder", "RcaImageEncoder", "VideoEncoderType"] diff --git a/src/trame_rca/encoders/video_encoder.py b/src/trame_rca/encoders/video_encoder.py index e69ce95..c017d8f 100644 --- a/src/trame_rca/encoders/video_encoder.py +++ b/src/trame_rca/encoders/video_encoder.py @@ -1,4 +1,5 @@ import logging +from enum import Enum, auto from time import time_ns from typing import Callable @@ -50,11 +51,18 @@ def encode( return (bytes(frame_data), meta, now_ms) +class VideoEncoderType(Enum): + AUTO = auto() + NVENC = auto() + VP9 = auto() + + class RcaVideoEncoder: def __init__( self, render_window: vtkRenderWindow, push_callback: Callable[[bytes, dict, int], None], + rca_encoder: VideoEncoderType | str | None = None, ) -> None: self._render_window = render_window self.encoder = None @@ -62,7 +70,14 @@ def __init__( self._push_callback = push_callback self._window_size = None - if vtkNvEncoderGL.CheckAvailability(): + rca_encoder = rca_encoder or VideoEncoderType.AUTO + if isinstance(rca_encoder, str): + rca_encoder = VideoEncoderType[rca_encoder.upper()] + + if vtkNvEncoderGL.CheckAvailability() and rca_encoder in [ + VideoEncoderType.AUTO, + VideoEncoderType.NVENC, + ]: logger.info("Using H264 through NVENC") self._create_encoder = create_nvenc_encoder else: diff --git a/src/trame_rca/schedulers/video_scheduler.py b/src/trame_rca/schedulers/video_scheduler.py index b4646bb..5533b45 100644 --- a/src/trame_rca/schedulers/video_scheduler.py +++ b/src/trame_rca/schedulers/video_scheduler.py @@ -10,7 +10,7 @@ from trame_rca.rca import VtkRemoteControlledArea -from trame_rca.encoders import RcaVideoEncoder +from trame_rca.encoders import RcaVideoEncoder, VideoEncoderType from trame_rca.rca import window_wrapper @@ -33,10 +33,11 @@ def __init__( *, push_callback: Callable[[bytes, dict]] | None = None, target_fps: float = 30.0, + rca_encoder: VideoEncoderType | None = None, ): self._rca: VtkRemoteControlledArea = window_wrapper(window) self._rca_encoder = RcaVideoEncoder( - self._rca.render_window, push_callback=self._push + self._rca.render_window, push_callback=self._push, rca_encoder=rca_encoder ) self._is_closing = False self._render_requested = False diff --git a/src/trame_rca/widgets/rca.py b/src/trame_rca/widgets/rca.py index 4a02c8a..9027cd5 100644 --- a/src/trame_rca/widgets/rca.py +++ b/src/trame_rca/widgets/rca.py @@ -11,7 +11,7 @@ if TYPE_CHECKING: from vtkmodules.vtkRenderingCore import vtkRenderWindow -from trame_rca.encoders import RcaImageEncoder +from trame_rca.encoders import RcaImageEncoder, VideoEncoderType from trame_rca.rca import RemoteControlledAreaProtocol, window_wrapper from trame_rca.schedulers import RcaImageRenderScheduler from trame_rca.utils import RcaViewAdapter @@ -101,7 +101,7 @@ def add_view_handler(self, view_handler): def create_view_handler( self, window: RemoteControlledAreaProtocol | vtkRenderWindow, - encoder: RcaImageEncoder | str | None = None, + encoder: RcaImageEncoder | VideoEncoderType | str | None = None, target_fps: float = 30.0, interactive_quality: int = 60, still_quality: int = 90, @@ -114,7 +114,9 @@ def create_view_handler( if not isinstance(window, VtkRemoteControlledArea): raise TypeError("Only VTK backends are supported by video decoder") - scheduler = RcaVideoRenderScheduler(window, target_fps=target_fps) + scheduler = RcaVideoRenderScheduler( + window, target_fps=target_fps, rca_encoder=encoder + ) elif encoder: scheduler = RcaImageRenderScheduler(