Method RunImpl

plugins/video/pkg_src/src/decoder/video_decoder_mixed.cc:70–127

}

void VideoDecoderMixed::RunImpl(dali::Workspace &ws) {
  auto &output = ws.Output<dali::GPUBackend>(0);
  const auto &input = ws.Input<dali::CPUBackend>(0);
  int batch_size = input.num_samples();

  CUcontext cuContext = nullptr;
  CUstream cuStream = ws.stream();
  cuCtxGetCurrent(&cuContext);

  if (!cuContext) {
    throw std::runtime_error("Failed to create a cuda context");
  }

  for (int i = 0; i < batch_size; i++) {
    auto output_sample = output[i];
    uint8_t *output_data = output_sample.template mutable_data<uint8_t>();

    auto &sample = samples_[i];
    sample.decoder_ = std::make_unique<NvDecoder>(
        cuStream, cuContext, true, FFmpeg2NvCodecId(sample.demuxer_->GetVideoCodec()), false,
        false /*_enableasyncallocations*/, false);


    uint8_t *pVideo = NULL, *pFrame = nullptr;
    int nVideoBytes = 0, nFrameReturned = 0, nFrame = 0;

    int num_frames = 0;

    do {
      sample.demuxer_->Demux(&pVideo, &nVideoBytes);
      nFrameReturned = sample.decoder_->Decode(pVideo, nVideoBytes);

      for (int i = 0; i < nFrameReturned; i++) {
        pFrame = sample.decoder_->GetFrame();
        CUDA_CALL(cudaStreamSynchronize(cuStream));

        uint8_t *dpFrame = output_data + num_frames * sample.demuxer_->GetHeight() *
                                             sample.demuxer_->GetWidth() * 3;
        int nWidth = sample.decoder_->GetWidth();
        int nPitch = sample.decoder_->GetWidth();
        int iMatrix =
            sample.decoder_->GetVideoFormatInfo().video_signal_description.matrix_coefficients;
        bool full_range =
            sample.decoder_->GetVideoFormatInfo().video_signal_description.video_full_range_flag;


        yuv_to_rgb(pFrame, nPitch, reinterpret_cast<uint8_t *>(dpFrame),
                   sample.decoder_->GetWidth() * 3, sample.decoder_->GetWidth(),
                   sample.decoder_->GetHeight(), full_range, cuStream);

        ++num_frames;
        if (end_frame_ > 0 && num_frames >= end_frame_) {
          return;
        }
      }
    } while (nVideoBytes);
  }
}