Method __call__

bench/python/nvcodec_utils.py:246–317

(self, batch)

    # docs_tag: begin_call_videobatchencoder_pyvideocodec
    def __call__(self, batch):
        self.cvcuda_perf.push_range("encoder.pyVideoCodec")

        # Get the name of the original video file read by the decoder. We would use
        # the same filename to save the output video.
        file_name = os.path.splitext(os.path.basename(batch.fileinfo))[0]
        self.output_file_name = os.path.join(self.output_path, "out_%s.mp4" % file_name)

        assert isinstance(batch.data, torch.Tensor)

        # docs_tag: begin_alloc_cvcuda_videobatchencoder_pyvideocodec
        # Check if we need to allocate the encoder for its first use.
        if self.encoder is None:
            self.encoder = nvVideoEncoder(
                self.device_id,
                batch.data.shape[3],
                batch.data.shape[2],
                self.fps,
                self.output_file_name,
                self.cuda_ctx,
                self.cuda_stream,
                "NV12",
            )
        # docs_tag: end_alloc_cvcuda_videobatchencoder_pyvideocodec

        # docs_tag: begin_convert_videobatchencoder_pyvideocodec

        # Create 2 CVCUDA tensors: reformat NCHW->NHWC and color conversion RGB->YUV
        current_batch_size = batch.data.shape[0]
        height, width = batch.data.shape[2], batch.data.shape[3]

        # Allocate only for the first time or for the last batch.
        if (
            not self.cvcuda_HWCtensor_batch
            or current_batch_size != self.cvcuda_HWCtensor_batch.shape[0]
        ):
            self.cvcuda_HWCtensor_batch = cvcuda.Tensor(
                (current_batch_size, height, width, 3),
                cvcuda.Type.U8,
                cvcuda.TensorLayout.NHWC,
            )
            self.cvcuda_YUVtensor_batch = cvcuda.Tensor(
                (current_batch_size, (height // 2) * 3, width, 1),
                cvcuda.Type.U8,
                cvcuda.TensorLayout.NHWC,
            )

        # Convert RGB to NV12, in batch, before sending it over to pyVideoCodec.
        # Convert to CVCUDA tensor
        cvcuda_tensor = cvcuda.as_tensor(batch.data, cvcuda.TensorLayout.NCHW)

        # Reformat NCHW to NHWC
        cvcuda.reformat_into(self.cvcuda_HWCtensor_batch, cvcuda_tensor)

        # Color convert from RGB to YUV_NV12
        cvcuda.cvtcolor_into(
            self.cvcuda_YUVtensor_batch,
            self.cvcuda_HWCtensor_batch,