Method extract_flow

src/opr/impl/nvof/denseflownvidia.cpp:122–181

}

void NVFlowExtractor::extract_flow(
        unsigned char* frames, std::vector<size_t>& shape, int16_t* result_out_ptr) {
    auto batch_size = shape[0];
    auto temporal_size = shape[1];
    auto height = shape[2];
    auto width = shape[3];
    auto channel = shape[4];
    auto temporal_len = height * width * channel;
    auto batch_len = temporal_size * height * width * channel;

    init_nvof_engine();

    auto src_mem_type = get_mem_type(reinterpret_cast<CUdeviceptr>(frames));
    auto out_mem_type = get_mem_type(reinterpret_cast<CUdeviceptr>(result_out_ptr));

    if ((height != m_height || width != m_width) ||
        (m_temporal_size != temporal_size)) {
        mgb_log_warn("We do not support dynamic shape at mgb side");
        mgb_throw(MegBrainError, "NVOF: Nvof err shap!!!! err type: NV_OF_ERR_GENERIC");
    }

    for (size_t batch_idx = 0; batch_idx < batch_size; batch_idx++) {
        auto input_buffer_batch_offsect = buffer_pool_size * batch_idx;
        auto output_buffer_batch_offsect = (buffer_pool_size - 1) * batch_idx;
        input_buffers[input_buffer_batch_offsect]->UploadData(
                (unsigned char*)(frames + batch_idx * batch_len), src_mem_type);

        for (size_t temporal_idx = 1; temporal_idx < temporal_size; temporal_idx++) {
            input_buffers[input_buffer_batch_offsect + temporal_idx % buffer_pool_size]->UploadData(
                    (unsigned char*)(frames + batch_idx * batch_len + temporal_idx * temporal_len),
                    src_mem_type);

            nv_optical_flow->Execute(
                    input_buffers
                            [input_buffer_batch_offsect +
                             (temporal_idx - 1) % buffer_pool_size]
                                    .get(),
                    input_buffers
                            [input_buffer_batch_offsect +
                             temporal_idx % buffer_pool_size]
                                    .get(),
                    output_buffers
                            [output_buffer_batch_offsect +
                             (temporal_idx - 1) % (buffer_pool_size - 1)]
                                    .get(),
                    nullptr, nullptr);

            output_buffers
                    [output_buffer_batch_offsect +
                     (temporal_idx - 1) % (buffer_pool_size - 1)]
                            ->DownloadData(
                                    result_out_ptr +
                                            batch_idx * (temporal_size - 1) * out_size +
                                            (temporal_idx - 1) * out_size,
                                    out_mem_type);
        }
    }