Pinjie/rebuild frame format manage

packages/on_demand_video_decoder/accvlab/on_demand_video_decoder/_internal/decoder.py

@@ -271,8 +271,9 @@ def CreateGopDecoder(
     Args:
         maxfiles: Maximum number of unique files that can be processed concurrently
         iGpu: GPU device ID to use for decoding (0 for primary GPU)
-        suppressNoColorRangeWarning: Suppress warning when no color range can be
-                                     extracted from video files (limited/MPEG range is assumed)
+        suppressNoColorRangeWarning: Suppress the warning emitted during RGB/BGR conversion
+                                     when the input color range is unspecified. Limited/MPEG
+                                     range is assumed regardless of this option.
         gopCacheCapacity: Maximum number of filepath entries kept in the Python GOP cache.
                           ``None`` defaults to ``maxfiles``. This capacity only affects
                           calls with ``useGOPCache=True``; each filepath stores the most

packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/CMakeLists.txt

@@ -49,6 +49,7 @@ set(PY_SOURCES
       src/PyNvOnDemandDecoder.cpp
       src/PyCAIMemoryView.cpp
       src/PyDecodedFrameExt.cpp
+      src/FrameOutput.cpp
       src/PyNvGopDecoder_common.cpp
       src/PyNvGopDecoder_constructors.cpp
       src/PyNvGopDecoder_random_decoder.cpp

packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/FrameOutput.hpp

@@ -0,0 +1,79 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <cuda.h>
+#include <cuviddec.h>
+#include <libavutil/pixfmt.h>
+
+class NvDecoder;
+struct DecodedFrameExt;
+class RGBFrame;
+
+enum Pixel_Format {
+    Pixel_Format_UNDEFINED = 0,
+    Pixel_Format_NV12 = 3,
+    Pixel_Format_YUV444 = 4,
+    Pixel_Format_P016 = 5,
+    Pixel_Format_YUV444_16Bit = 6
+
+};
+
+namespace accvlab::on_demand_video_decoder::internal::frame_output {
+
+// Tightly packed layouts exposed by the public frame-output APIs. RGB8 also
+// covers BGR8 because channel order does not change the allocation size.
+enum class FrameOutputFormat : uint8_t {
+    RGB8,
+    NV12,
+    P016,
+    YUV444,
+    YUV444_16Bit,
+};
+
+FrameOutputFormat output_format_from_pixel_format(Pixel_Format format);
+Pixel_Format pixel_format_from_surface(cudaVideoSurfaceFormat surface_format);
+FrameOutputFormat output_format_from_surface(cudaVideoSurfaceFormat surface_format);
+FrameOutputFormat output_format_from_av_pixel_format(AVPixelFormat pixel_format);
+
+// The single size-calculation entry point for all tightly packed RGB/YUV
+// frames exposed by this package.
+size_t frame_bytes(FrameOutputFormat format, size_t height, size_t width);
+
+// Convert one NVDEC surface into an RGB/BGR frame backed by output_buffer.
+// When is_async is true, work is only enqueued on decoder.GetStream(); the caller
+// owns the terminal synchronization and must keep both buffers alive until then.
+RGBFrame convert_decoded_frame_to_rgb(NvDecoder& decoder, const uint8_t* decoded_surface,
+                                      CUdeviceptr output_buffer, AVColorRange color_range, bool as_bgr,
+                                      bool is_async);
+
+// Copy one NVDEC surface into output_buffer and expose its native YUV planes.
+DecodedFrameExt copy_decoded_frame_to_yuv(NvDecoder& decoder, const uint8_t* decoded_surface,
+                                          CUdeviceptr output_buffer, AVColorRange color_range,
+                                          int64_t timestamp, bool is_async);
+
+// Copy an existing RGB/BGR frame into aggregator-owned storage.
+RGBFrame copy_rgb_frame(const RGBFrame& source, CUdeviceptr destination, CUstream destination_stream,
+                        bool as_bgr, bool is_async);
+
+// Copy an existing tightly packed YUV frame into aggregator-owned storage.
+DecodedFrameExt copy_yuv_frame(const DecodedFrameExt& source, CUdeviceptr destination,
+                               CUstream destination_stream, bool is_async);
+
+}  // namespace accvlab::on_demand_video_decoder::internal::frame_output

packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyCAIMemoryView.hpp

@@ -15,12 +15,12 @@
  */
 
 #pragma once
-#include "ExternalBuffer.hpp"
 #include "NvCodecUtils.h"
 #include "nvEncodeAPI.h"
 #include <cuda.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
+#include <memory>
 #include <sstream>
 #include <string>
 #include <vector>
@@ -78,15 +78,6 @@ class CuCtxGuard {
 #define DEF_CONSTANT(s) attr(ENUM_VALUE_STRINGIFY(s)) = py::cast(s)
 #define DEF_READWRITE(type, s) def_readwrite(ENUM_VALUE_STRINGIFY(s), &type::s)
 
-enum Pixel_Format {
-    Pixel_Format_UNDEFINED = 0,
-    Pixel_Format_NV12 = 3,
-    Pixel_Format_YUV444 = 4,
-    Pixel_Format_P016 = 5,
-    Pixel_Format_YUV444_16Bit = 6
-
-};
-
 struct CAIMemoryView {
     std::vector<size_t> shape;
     std::vector<size_t> stride;
@@ -147,139 +138,5 @@ struct CAIMemoryView {
     }
 };
 
-struct DecodedFrame {
-    int64_t timestamp;
-    std::vector<CAIMemoryView> views;
-    Pixel_Format format;
-    std::shared_ptr<ExternalBuffer> extBuf;
-    DecodedFrame() { extBuf = std::make_shared<ExternalBuffer>(); }
-    static void Export(py::module& m) {
-        py::class_<DecodedFrame, std::shared_ptr<DecodedFrame>>(m, "DecodedFrame", py::module_local())
-            .def_readonly("timestamp", &DecodedFrame::timestamp)
-            .def_readonly("format", &DecodedFrame::format)
-            .def("__repr__",
-                 [](std::shared_ptr<DecodedFrame>& self) {
-                     std::stringstream ss;
-                     ss << "<DecodedFrame [";
-                     ss << "timestamp=" << self->timestamp;
-                     ss << ", format=" << py::str(py::cast(self->format));
-                     ss << ", " << py::str(py::cast(self->views));
-                     ss << "]>";
-                     return ss.str();
-                 })
-            .def(
-                "framesize",
-                [](std::shared_ptr<DecodedFrame>& self) {
-                    int height = self->views.at(0).shape.at(0);
-                    int width = self->views.at(0).shape.at(1);
-                    int framesize = width * height * 1.5;
-                    switch (self->format) {
-                        case Pixel_Format_NV12:
-                            break;
-                        case Pixel_Format_P016:
-                            framesize = width * height * 3;
-                            break;
-                        case Pixel_Format_YUV444:
-                            framesize = width * height * 3;
-                            break;
-                        case Pixel_Format_YUV444_16Bit:
-                            framesize = width * height * 6;
-                            break;
-                        default:
-                            break;
-                    }
-                    return framesize;
-                },
-                R"pbdoc(
-            return underlying views which implement CAI
-            :param None: None
-            )pbdoc")
-            .def(
-                "cuda", [](std::shared_ptr<DecodedFrame>& self) { return self->views; },
-                R"pbdoc(
-            return underlying views which implement CAI
-            :param None: None
-            )pbdoc")
-            .def(
-                "nvcv_image",
-                [](std::shared_ptr<DecodedFrame>& self) {
-                    switch (self->format) {
-                        case Pixel_Format_NV12: {
-                            size_t width = self->views.at(0).shape[1];
-                            size_t height = self->views.at(0).shape[0] * 1.5;
-                            CUdeviceptr data = self->views.at(0).data;
-                            CUstream stream = self->views.at(0).stream;
-                            self->views.clear();
-                            self->views.push_back(
-                                CAIMemoryView{{height, width, 1},
-                                              {width, 2, 1},
-                                              "|u1",
-                                              reinterpret_cast<size_t>(stream),
-                                              (data),
-                                              false});  // hack for cvcuda tensor represenation
-                        } break;
-                        case Pixel_Format_YUV444: {
-                            size_t width = self->views.at(0).shape[1];
-                            size_t height = self->views.at(0).shape[0] * 3;
-                            CUdeviceptr data = self->views.at(0).data;
-                            CUstream stream = self->views.at(0).stream;
-                            self->views.clear();
-                            self->views.push_back(
-                                CAIMemoryView{{height, width, 1},
-                                              {width, 3, 1},
-                                              "|u1",
-                                              reinterpret_cast<size_t>(stream),
-                                              (data),
-                                              false});  // hack for cvcuda tensor represenation
-                        } break;
-                        default:
-                            throw std::invalid_argument("only nv12 and yuv444 supported as of now");
-                            break;
-                    }
-                    return self->views;
-                },
-                R"pbdoc(
-            return underlying views which implement CAI
-            :param None: None
-            )pbdoc")
-
-            // DL Pack Tensor
-            .def_property_readonly(
-                "shape", [](std::shared_ptr<DecodedFrame>& self) { return self->extBuf->shape(); },
-                "Get the shape of the buffer as an array")
-            .def_property_readonly(
-                "strides", [](std::shared_ptr<DecodedFrame>& self) { return self->extBuf->strides(); },
-                "Get the strides of the buffer")
-            .def_property_readonly(
-                "dtype", [](std::shared_ptr<DecodedFrame>& self) { return self->extBuf->dtype(); },
-                "Get the data type of the buffer")
-            .def(
-                "__dlpack__",
-                [](std::shared_ptr<DecodedFrame>& self, py::object stream) {
-                    return self->extBuf->dlpack(stream);
-                },
-                py::arg("stream") = NULL, "Export the buffer as a DLPack tensor")
-            .def(
-                "__dlpack_device__",
-                [](std::shared_ptr<DecodedFrame>& self) {
-                    // DLDevice ctx;
-                    // ctx.device_type = DLDeviceType::kDLCUDA;
-                    // ctx.device_id = 0;
-                    return py::make_tuple(py::int_(static_cast<int>(DLDeviceType::kDLCUDA)),
-                                          py::int_(static_cast<int>(0)));
-                },
-                "Get the device associated with the buffer")
-            .def(
-                "GetPtrToPlane",
-
-                [](std::shared_ptr<DecodedFrame>& self, int planeIdx) { return self->views[planeIdx].data; },
-                R"pbdoc(
-            return pointer to base address for plane index
-            :param planeIdx : index to the plane
-            )pbdoc");
-        // TODO add __iter__ interface on DecodedFrame
-    }
-};
-
 CAIMemoryView coerceToCudaArrayView(py::object cuda_array, NV_ENC_BUFFER_FORMAT bufferFormat, size_t width,
                                     size_t height, int planeIdx = 0);

packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyDecodedFrameExt.hpp

@@ -18,12 +18,14 @@
 
 #include <string>
 
+#include "ExternalBuffer.hpp"
+#include "FrameOutput.hpp"
 #include "PyCAIMemoryView.hpp"
 
 #include "cuviddec.h"
 #include <libavutil/pixfmt.h>
 
-struct DecodedFrameExt : public DecodedFrame {
+struct DecodedFrameExt {
     enum class VideoSurfaceFormat {
         UNSPECIFIED = 0,
         NV12 = 1,
@@ -46,8 +48,12 @@ struct DecodedFrameExt : public DecodedFrame {
         ColorRange_FULL = 2,
     };
 
+    int64_t timestamp;
+    std::vector<CAIMemoryView> views;
+    Pixel_Format format;
+    std::shared_ptr<ExternalBuffer> extBuf;
     ColorRange color_range = ColorRange::ColorRange_UNSPECIFIED;
-    DecodedFrameExt() = default;
+    DecodedFrameExt() { extBuf = std::make_shared<ExternalBuffer>(); }
     VideoSurfaceFormat GetVideoSurfaceFormat() const;
 
     void SetVideoSurfaceFormat(cudaVideoSurfaceFormat video_format_in);
@@ -93,27 +99,17 @@ struct DecodedFrameExt : public DecodedFrame {
                 [](std::shared_ptr<DecodedFrameExt>& self) {
                     int height = self->views.at(0).shape.at(0);
                     int width = self->views.at(0).shape.at(1);
-                    int framesize = width * height * 1.5;
-                    switch (self->format) {
-                        case Pixel_Format_NV12:
-                            break;
-                        case Pixel_Format_P016:
-                            framesize = width * height * 3;
-                            break;
-                        case Pixel_Format_YUV444:
-                            framesize = width * height * 3;
-                            break;
-                        case Pixel_Format_YUV444_16Bit:
-                            framesize = width * height * 6;
-                            break;
-                        default:
-                            break;
-                    }
+                    int framesize = static_cast<int>(
+                        accvlab::on_demand_video_decoder::internal::frame_output::frame_bytes(
+                            accvlab::on_demand_video_decoder::internal::frame_output::
+                                output_format_from_pixel_format(self->format),
+                            height, width));
                     return framesize;
                 },
                 R"pbdoc(
-            return underlying views which implement CAI
-            :param None: None
+            Return the total size in bytes of the tightly packed decoded frame buffer.
+
+            The size includes all YUV planes.
             )pbdoc")
             .def(
                 "cuda", [](std::shared_ptr<DecodedFrameExt>& self) { return self->views; },
@@ -127,7 +123,8 @@ struct DecodedFrameExt : public DecodedFrame {
                     switch (self->format) {
                         case Pixel_Format_NV12: {
                             size_t width = self->views.at(0).shape[1];
-                            size_t height = self->views.at(0).shape[0] * 1.5;
+                            size_t luma_height = self->views.at(0).shape[0];
+                            size_t height = luma_height + (luma_height + 1) / 2;
                             CUdeviceptr data = self->views.at(0).data;
                             CUstream stream = self->views.at(0).stream;
                             self->views.clear();

packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvBatchAsyncGopDecoder.hpp

@@ -140,15 +140,6 @@ class PyNvBatchAsyncGopDecoder {
                      std::vector<std::string> filepaths, std::vector<std::vector<int>> frame_ids_2d,
                      bool as_bgr, bool is_rgb);
 
-    // Returns the total contiguous byte size of one YUV frame for the given
-    // pixel format and dimensions.  Matches the layout produced by GetYUVFromFrame.
-    static size_t compute_yuv_frame_bytes(Pixel_Format fmt, size_t H, size_t W);
-
-    // Reconstruct a DecodedFrameExt whose views point into aggregator pool memory.
-    static void build_yuv_frame(Pixel_Format fmt, size_t H, size_t W, int64_t timestamp,
-                                DecodedFrameExt::ColorRange color_range, CUdeviceptr dst_ptr, CUstream stream,
-                                DecodedFrameExt& out);
-
    private:
     bool suppress_no_color_range_warning_ = false;
     bool destroy_context_ = false;