video_stage: add an optional GPU stage to the relay

CMakeLists.txt

@@ -73,6 +73,7 @@ option(BUILD_PLUGIN_WUJI_GLOVE "Build Wuji glove plugin (requires the wuji_sdk C
 option(BUILD_PLUGIN_SHARPA_AVATAR "Build Sharpa AvatarExo Glove plugin (Linux; skips if the Avatar SDK is missing)" OFF)
 option(BUILD_PLUGIN_SENSING "Build SENSING GMSL camera plugin (Jetson only; requires SIPL and NVENC)" OFF)
 option(BUILD_VIDEO_STAGE "Build the video stage: the Orin-to-PC video transport and the PC relay" OFF)
+option(BUILD_VIDEO_STAGE_GPU "Add the optional GPU stage to the video stage (needs a CUDA toolkit; its tests need a CUDA GPU)" OFF)
 option(BUILD_EXAMPLES "Build examples" ON)
 option(BUILD_EXAMPLE_TELEOP_ROS2 "Build only the teleop_ros2 ROS 2 reference integration (e.g. for Docker)" OFF)
 option(BUILD_TESTING "Build unit tests" ON)

src/video_stage/CMakeLists.txt

@@ -2,20 +2,53 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # ==============================================================================
-# Video stage: the Orin-to-PC video transport and the PC relay to LiveKit
+# Video stage: the Orin-to-PC video transport, the PC relay to LiveKit, and an optional PC GPU stage
 # ==============================================================================
 #   wire/    video_stage_wire   -- AVS1 records, clock exchange, H.264 access-unit helpers, frame correlator
 #   sender/  video_stage_sender -- the Orin-side sender of AVS1 records over a direct link
 #   app/     video_stage_worker -- the PC relay: receives the Orin's records and publishes them to LiveKit or a file
 #
-# No CUDA. The LiveKit SDK (src/livekit_publisher) is needed only by the relay.
+#   core/    video_stage_core   -- (GPU) the operation interface and the example sine-wave warp
+#   codec/   video_stage_codec  -- (GPU) NVDEC decoder and NVENC encoder
+#
+# No CUDA unless BUILD_VIDEO_STAGE_GPU is ON: the relay forwards the Orin's access units untouched. The LiveKit SDK
+# (src/livekit_publisher) is needed only by the relay.
 #
 # Configure with:
-#   cmake -B build -DBUILD_VIDEO_STAGE=ON
+#   cmake -B build -DBUILD_VIDEO_STAGE=ON [-DBUILD_VIDEO_STAGE_GPU=ON]
 # ==============================================================================
 
 add_subdirectory(wire)
 add_subdirectory(sender)
+
+if(BUILD_VIDEO_STAGE_GPU)
+    enable_language(CUDA)
+
+    if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+        # Real code for the common desktop architectures plus PTX for 90, which newer GPUs JIT-compile.
+        set(CMAKE_CUDA_ARCHITECTURES "80-real;86-real;89-real;90")
+    endif()
+
+    find_package(CUDAToolkit REQUIRED)
+
+    # NVDEC/NVENC API headers (MIT licence) with run-time loaders for libnvcuvid and libnvidia-encode, pinned by hash.
+    include(FetchContent)
+    FetchContent_Declare(
+        nv_codec_headers
+        URL https://github.com/FFmpeg/nv-codec-headers/archive/refs/tags/n12.2.72.0.tar.gz
+        URL_HASH SHA256=dbeaec433d93b850714760282f1d0992b1254fc3b5a6cb7d76fc1340a1e47563
+    )
+    FetchContent_GetProperties(nv_codec_headers)
+    if(NOT nv_codec_headers_POPULATED)
+        FetchContent_Populate(nv_codec_headers)
+    endif()
+    add_library(video_stage_nv_codec_headers INTERFACE)
+    target_include_directories(video_stage_nv_codec_headers SYSTEM INTERFACE "${nv_codec_headers_SOURCE_DIR}/include")
+
+    add_subdirectory(core)
+    add_subdirectory(codec)
+endif()
+
 add_subdirectory(app)
 
 if(BUILD_TESTING)

src/video_stage/README.md

@@ -6,8 +6,8 @@ SPDX-License-Identifier: Apache-2.0
 # Video stage
 
 Moves camera video from the Orin to the PC and on to LiveKit. The Orin sends its encoded H.264 access units to the PC
-over a direct link; the PC relay (`video_stage_worker`) hands them, untouched, to LiveKit or a file. There is no decode,
-encode or GPU in the relay, and it builds without CUDA.
+over a direct link; the PC relay (`video_stage_worker`) hands them, untouched, to LiveKit or a file. By default there is no decode,
+encode or GPU in the relay, and it builds without CUDA; an optional GPU stage is described below.
 
 ```text
 Orin producer --DATA (AVS1 records, H.264)--> PC relay --> LiveKit track (or an Annex-B file)
@@ -62,6 +62,49 @@ hello order, retry, keyframe start, bounded queue with drop-to-keyframe, keyfram
 offset clock), and feeds it Jetson-style H.264 (no AUDs, IDR every 5 s, constant bitrate). One `--port` per camera. Run it
 against a relay writing a file, then check the file with `ffprobe`.
 
+## The GPU stage (optional)
+
+Built only with `-DBUILD_VIDEO_STAGE_GPU=ON` (needs a CUDA toolkit; the tests need a CUDA GPU). `video_stage_worker --operation=sine-warp` then decodes each picture on NVDEC, keeps it in GPU memory through one GPU operation, and re-encodes it on NVENC, instead of forwarding the Orin's access units untouched. A build without it rejects `--operation=sine-warp` with a clear message and relays exactly as before.
+
+### The operation contract
+
+`core/inc/video_stage/gpu_operation.hpp` defines `GpuOperation`: NV12 in, NV12 out, fixed geometry.
+
+- `prepare(input_geometry)` runs once, outside the frame loop. Build and upload anything per-geometry here (a
+  rectification map, for example). Nothing in `enqueue` allocates or uploads.
+- `enqueue(in, out, stream)` only queues kernels on `stream` and returns. It does not synchronise the device or the
+  stream. After it returns, the caller records a CUDA event on `stream`; whoever reads `out` waits on that event.
+- `in` and `out` are different surfaces (the operation is not in-place). `out` has `output_geometry(in.geometry)` and
+  keeps the input's color description. Both are `Nv12View`s: device pointers, byte pitches, size.
+- Frame identity and timestamps are not the operation's concern; they travel with the frame.
+
+### The example: `SineWarp`
+
+`core/sine_warp.cu` is a horizontal sine ripple, `source_x = x + A*sin(2*pi*y/wavelength + phase)`,
+`source_y = y`, sampled bilinearly, with the border fill for the declared range (Y=16 limited, 0 full; chroma 128).
+Amplitude 0 is an identity. It is an example, not calibrated rectification: the warped picture does not keep the
+source camera's projection.
+
+The stage assumes the decoder's NV12 uses the H.264 default chroma siting described below; a stream that signals a
+different siting must be rejected at the decoder, not warped with the wrong sample positions.
+
+Chroma is not sampled at the luma coordinate divided by two. A chroma sample of H.264 NV12 sits on an even luma column
+and half way between two luma rows, so the kernel maps that luma-space position, then converts back to chroma-plane
+coordinates. Keep that when substituting an operation.
+
+### Substituting the fisheye rectification
+
+Replace the mapping, keep the rest: a class deriving from `GpuOperation` whose `prepare` generates the destination-to-
+source maps once (on the CPU, uploaded once) and whose `enqueue` launches a kernel that reads `map_x/map_y` instead of
+computing the sine. Sample luma at the map position and chroma at the position derived from the same map with the
+siting above. `output_geometry` returns the rectified output size. Nothing else in the stage changes.
+
+### Identity and overload
+
+Each picture carries a key through NVDEC and NVENC, and the source frame number and stamp are looked up by that key, so
+metadata follows its own frame through loss. A picture arriving while every encoder input surface is busy is dropped and
+counted, never queued.
+
 ## Build
 
 ```bash
@@ -70,4 +113,6 @@ cmake -B build -DBUILD_VIDEO_STAGE=ON -DBUILD_PLUGINS=OFF -DBUILD_VIZ=OFF -DBUIL
 cmake --build build && ctest --test-dir build -R 'video_stage|avs1|clock'
 ```
 
+Add `-DBUILD_VIDEO_STAGE_GPU=ON` for the GPU stage and its tests.
+
 The LiveKit SDK is fetched, or taken from `LIVEKIT_SDK_DIR`.

src/video_stage/app/CMakeLists.txt

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# The PC relay: receives AVS1 access units from the Orin over its direct link and hands them, untouched, to LiveKit or
-# a file.
+# The PC worker: receives AVS1 access units from the Orin and hands them, untouched, to LiveKit or a file. With
+# BUILD_VIDEO_STAGE_GPU it can also decode on NVDEC, apply a GPU operation and re-encode on NVENC first.
 add_executable(video_stage_worker
     main.cpp
     control_channel.cpp
@@ -16,3 +16,9 @@ target_link_libraries(video_stage_worker
         video_stage::wire
         sensing::livekit
 )
+
+if(BUILD_VIDEO_STAGE_GPU)
+    target_sources(video_stage_worker PRIVATE pipeline.cpp)
+    target_compile_definitions(video_stage_worker PRIVATE VIDEO_STAGE_GPU)
+    target_link_libraries(video_stage_worker PRIVATE video_stage::core video_stage::codec)
+endif()

src/video_stage/app/main.cpp

@@ -4,6 +4,10 @@
 #include "livekit_output.hpp"
 #include "stream_worker.hpp"
 
+#ifdef VIDEO_STAGE_GPU
+#    include <video_stage/cuda_context.hpp>
+#endif
+
 #include <atomic>
 #include <chrono>
 #include <csignal>
@@ -32,11 +36,13 @@ void usage(const char* program)
     std::cout
         << "Usage: " << program << " --stream=PORT:TRACK ... --livekit-url=URL --livekit-token-file=PATH\n"
         << "       " << program << " --stream=PORT:file=PATH.h264 ...\n"
-        << "       [--bitrate=BPS] [--max-upstream-uncertainty-ms=MS] [--stats-interval=S] [--once]\n"
+        << "       [--operation=none|sine-warp] [--amplitude=PX] [--bitrate=BPS] [--idr-period=N] [--max-upstream-uncertainty-ms=MS] [--stats-interval=S] [--once]\n"
         << "One --stream per camera. Each listens on PORT for the Orin's direct connection (AVS1 H.264 records plus a\n"
         << "control link) and publishes the access units as LiveKit track TRACK (all tracks under one participant) or\n"
-        << "writes them to a file. The Orin's access units are forwarded untouched, with no decode, encode or GPU. Serves\n"
-        << "one Orin session after another unless --once. Prints a summary at the end.\n";
+        << "writes them to a file. By default (--operation=none) the Orin's access units are forwarded untouched, with no\n"
+        << "decode, encode or GPU. --operation=sine-warp decodes on NVDEC, applies the example sine warp and re-encodes on\n"
+        << "NVENC, as a test of the GPU stage (a build with -DBUILD_VIDEO_STAGE_GPU=ON). Serves one Orin\n"
+        << "session after another unless --once. Prints per-stage percentiles at the end.\n";
 }
 
 std::string read_token(const std::string& path)
@@ -87,8 +93,33 @@ try
             stats_interval_s = std::stoi(arg.substr(17));
         else if (arg == "--once")
             settings.once = true;
+        else if (arg.rfind("--operation=", 0) == 0)
+        {
+            const std::string name = arg.substr(12);
+            if (name == "none")
+                settings.operation = Operation::None;
+            else if (name == "sine-warp")
+                settings.operation = Operation::SineWarp;
+            else
+            {
+                std::cerr << "--operation must be none or sine-warp\n";
+                return 1;
+            }
+        }
         else if (arg.rfind("--bitrate=", 0) == 0)
-            settings.livekit.max_bitrate_bps = static_cast<uint32_t>(std::stoul(arg.substr(10)));
+        {
+            const auto bitrate = static_cast<uint32_t>(std::stoul(arg.substr(10)));
+            settings.livekit.max_bitrate_bps = bitrate;
+#ifdef VIDEO_STAGE_GPU
+            settings.pipeline.bitrate_bps = bitrate;
+#endif
+        }
+#ifdef VIDEO_STAGE_GPU
+        else if (arg.rfind("--amplitude=", 0) == 0)
+            settings.warp.amplitude_px = std::stof(arg.substr(12));
+        else if (arg.rfind("--idr-period=", 0) == 0)
+            settings.pipeline.idr_period = static_cast<uint32_t>(std::stoul(arg.substr(13)));
+#endif
         else if (arg.rfind("--max-upstream-uncertainty-ms=", 0) == 0)
             settings.livekit.max_upstream_uncertainty_ns = static_cast<int64_t>(std::stod(arg.substr(30)) * 1e6);
         else
@@ -111,14 +142,36 @@ try
         return 1;
     }
 
+    // Only an operation that uses the GPU needs a CUDA context; the default forwards the Orin's access units untouched.
+#ifdef VIDEO_STAGE_GPU
+    std::unique_ptr<CudaContext> context;
+    if (settings.operation != Operation::None)
+    {
+        context = std::make_unique<CudaContext>();
+        std::cout << "GPU: " << context->device_name() << "\n";
+    }
+#else
+    if (settings.operation != Operation::None)
+    {
+        std::cerr << "this build has no GPU stage; configure with -DBUILD_VIDEO_STAGE_GPU=ON\n";
+        return 1;
+    }
+    CudaContext* const context = nullptr;
+#endif
     // One LiveKit participant for every stream, as the Orin producer joins once for all its cameras.
     std::unique_ptr<SharedLiveKit> livekit;
     if (any_livekit)
         livekit = std::make_unique<SharedLiveKit>(livekit_url, read_token(token_file));
 
+#ifdef VIDEO_STAGE_GPU
+    CudaContext* const context_ptr = context.get();
+#else
+    CudaContext* const context_ptr = context;
+#endif
     std::vector<std::unique_ptr<StreamWorker>> workers;
     for (const StreamSpec& spec : specs)
-        workers.push_back(std::make_unique<StreamWorker>(spec, settings, spec.file.empty() ? livekit.get() : nullptr));
+        workers.push_back(
+            std::make_unique<StreamWorker>(context_ptr, spec, settings, spec.file.empty() ? livekit.get() : nullptr));
 
     std::signal(SIGINT, on_signal);
     std::signal(SIGTERM, on_signal);