借助 NVIDIA Jetson Thor 提升机器人感知效能

import vpi
import numpy as np
from PIL import Image
from argparse import ArgumentParser
 
# Create two streams for parallel processing
streamLeft = vpi.Stream()
streamRight = vpi.Stream()

# Load images and wrap them in VPI images
left_img = np.asarray(Image.open(args.left))
right_img = np.asarray(Image.open(args.right))
left = vpi.asimage(left_img)
right = vpi.asimage(right_img)
 
# Convert images to Y8_ER_BL format in parallel on different backends
left = left.convert(vpi.Format.Y8_ER_BL, scale=1, stream=streamLeft, backend=vpi.Backend.VIC)
right = right.convert(vpi.Format.Y8_ER_BL, scale=1, stream=streamRight, backend=vpi.Backend.CUDA)

# Synchronize streamLeft to ensure the left image is ready
streamLeft.sync()
 
# Submit the stereo disparity operation on streamRight
disparityS16 = vpi.stereodisp(left, right, backend=vpi.Backend.OFA|vpi.Backend.PVA|vpi.Backend.VIC, stream=streamRight)

# Post-process the disparity map 
# Convert Q10.5 to  U8 and scale for visualization
disparityU8 = disparityS16.convert(vpi.Format.U8, scale=255.0/(32*128), stream=streamRight, backend=vpi.Backend.CUDA)
 
# make accessible in cpu
disparityU8 = disparityU8.cpu()
 
#save with pillow
d_pil = Image.fromarray(disparityU8)
d_pil.save('./disparity.png')

int totalIterations = itersPerStream * numStreams;
std::vector<VPIImage> leftInputs(numStreams), rightInputs(numStreams), confidences(numStreams), leftTmps(numStreams), rightTmps(numStreams);
std::vector<VPIImage> leftOuts(numStreams), rightOuts(numStreams), disparities(numStreams);
std::vector<VPIPayload> stereoPayloads(numStreams);
std::vector<VPIStream> streamsLeft(numStreams), streamsRight(numStreams);
std::vector<VPIEvent> events(numStreams);
int width   = cvImageLeft.cols;
int height  = cvImageLeft.rows;
int vic_pva_ofa = VPI_BACKEND_VIC | VPI_BACKEND_OFA | VPI_BACKEND_PVA;
VPIStereoDisparityEstimatorCreationParams stereoPayloadParams;
VPIStereoDisparityEstimatorParams stereoParams;
CHECK_STATUS(vpiInitStereoDisparityEstimatorCreationParams(&stereoPayloadParams));
CHECK_STATUS(vpiInitStereoDisparityEstimatorParams(&stereoParams));
stereoPayloadParams.maxDisparity = 128;
stereoParams.maxDisparity= 128;
stereoParams.confidenceType  = VPI_STEREO_CONFIDENCE_RELATIVE;

for (int i = 0; i < numStreams; i++)
{
    CHECK_STATUS(vpiImageCreateWrapperOpenCVMat(cvImageLeft, 0, &leftInputs[i]));
    CHECK_STATUS(vpiImageCreateWrapperOpenCVMat(cvImageRight, 0, &rightInputs[i]));
    CHECK_STATUS(vpiStreamCreate(0, &streamsLeft[i]));
    CHECK_STATUS(vpiStreamCreate(0, &streamsRight[i]));
    CHECK_STATUS(vpiImageCreate(width, height, VPI_IMAGE_FORMAT_Y8_ER, 0, &leftTmps[i]));
    CHECK_STATUS(vpiImageCreate(width, height, VPI_IMAGE_FORMAT_NV12_BL, 0, &leftOuts[i]));
    CHECK_STATUS(vpiImageCreate(width, height, VPI_IMAGE_FORMAT_Y8_ER, 0, &rightTmps[i]));
    CHECK_STATUS(vpiImageCreate(width, height, VPI_IMAGE_FORMAT_NV12_BL, 0, &rightOuts[i]));
    CHECK_STATUS(vpiCreateStereoDisparityEstimator(vic_pva_ofa, width, height, VPI_IMAGE_FORMAT_NV12_BL,
    &stereoPayloadParams, &stereoPayloads[i]));
    CHECK_STATUS(vpiEventCreate(0, &events[i]));
}
int outCount = saveOutput ? (numStreams * itersPerStream) : numStreams;
disparities.resize(outCount);
confidences.resize(outCount);
for (int i = 0; i < outCount; i++)
{
    CHECK_STATUS(vpiImageCreate(width, height, VPI_IMAGE_FORMAT_S16, 0, &disparities[i]));
    CHECK_STATUS(vpiImageCreate(width, height, VPI_IMAGE_FORMAT_U16, 0, &confidences[i]));
}

for (int i = 0; i < numStreams; i++)
{
    CHECK_STATUS(vpiSubmitConvertImageFormat(streamsLeft[i], VPI_BACKEND_CPU, leftInputs[i], leftTmps[i], NULL));
    CHECK_STATUS(vpiSubmitConvertImageFormat(streamsLeft[i], VPI_BACKEND_VIC, leftTmps[i], leftOuts[i], NULL));
    CHECK_STATUS(vpiEventRecord(events[i], streamsLeft[i]));
    CHECK_STATUS(vpiSubmitConvertImageFormat(streamsRight[i], VPI_BACKEND_CPU, rightInputs[i], rightTmps[i], NULL));
    CHECK_STATUS(vpiSubmitConvertImageFormat(streamsRight[i], VPI_BACKEND_VIC, rightTmps[i], rightOuts[i], NULL));
    CHECK_STATUS(vpiStreamWaitEvent(streamsRight[i], events[i]));
}
for (int i = 0; i < numStreams; i++)
{
    CHECK_STATUS(vpiStreamSync(streamsLeft[i]));
    CHECK_STATUS(vpiStreamSync(streamsRight[i]));
}

auto benchmarkStart = std::chrono::high_resolution_clock::now();
for (int iter = 0; iter < itersPerStream; iter++)
{
    for (int i = 0; i < numStreams; i++)
    {
        int dispIdx = saveOutput ? (i * itersPerStream + iter) : i;
        CHECK_STATUS(vpiSubmitStereoDisparityEstimator(streamsRight[i], vic_pva_ofa, stereoPayloads[i], leftOuts[i],
                                                     rightOuts[i], disparities[dispIdx], confidences[dispIdx],
                                                     &stereoParams));
    }
}
// ====================
// End Benchmarking
for (int i = 0; i < numStreams; i++)
{
    CHECK_STATUS(vpiStreamSync(streamsRight[i]));
}
auto benchmarkEnd = std::chrono::high_resolution_clock::now();

// ====================
// Save Outputs
if (saveOutput)
{
    for (int i = 0; i < numStreams * itersPerStream; i++)
    {
        VPIImageData dispData, confData;
        cv::Mat cvDisparity, cvDisparityColor, cvConfidence, cvMask;
        CHECK_STATUS(
        vpiImageLockData(disparities[i], VPI_LOCK_READ, VPI_IMAGE_BUFFER_HOST_PITCH_LINEAR, &dispData));
        vpiImageDataExportOpenCVMat(dispData, &cvDisparity);
        cvDisparity.convertTo(cvDisparity, CV_8UC1, 255.0 / (32 * stereoParams.maxDisparity), 0);
        applyColorMap(cvDisparity, cvDisparityColor, cv::COLORMAP_JET);
        CHECK_STATUS(vpiImageUnlock(disparities[i]));
        std::ostringstream fpStream;
        fpStream << "stream_" << i / itersPerStream << "_iter_" << i % itersPerStream << "_disparity.png";
        imwrite(fpStream.str(), cvDisparityColor);

        // Confidence output (U16 -> scale to 8-bit and save)
        CHECK_STATUS(
        vpiImageLockData(confidences[i], VPI_LOCK_READ, VPI_IMAGE_BUFFER_HOST_PITCH_LINEAR, &confData));
        vpiImageDataExportOpenCVMat(confData, &cvConfidence);
        cvConfidence.convertTo(cvConfidence, CV_8UC1, 255.0 / 65535.0, 0);
        CHECK_STATUS(vpiImageUnlock(confidences[i]));
        std::ostringstream fpStreamConf;
        fpStreamConf << "stream_" << i / itersPerStream << "_iter_" << i % itersPerStream << "_confidence.png";
        imwrite(fpStreamConf.str(), cvConfidence);
    }
}

// ====================
// Clean Up VPI Objects
for (int i = 0; i < numStreams; i++)
{
    CHECK_STATUS(vpiStreamSync(streamsLeft[i]));
    CHECK_STATUS(vpiStreamSync(streamsRight[i]));
    vpiStreamDestroy(streamsLeft[i]);
    vpiStreamDestroy(streamsRight[i]);
    vpiImageDestroy(rightInputs[i]);
    vpiImageDestroy(leftInputs[i]);
    vpiImageDestroy(leftTmps[i]);
    vpiImageDestroy(leftOuts[i]);
    vpiImageDestroy(rightTmps[i]);
    vpiImageDestroy(rightOuts[i]);
    vpiPayloadDestroy(stereoPayloads[i]);
    vpiEventDestroy(events[i]);
}
// Destroy all disparity and confidence images
for (int i = 0; i < (int)disparities.size(); i++)
{
    vpiImageDestroy(disparities[i]);
}
for (int i = 0; i < (int)confidences.size(); i++)
{
    vpiImageDestroy(confidences[i]);
}

double totalTimeSeconds = totalTime / 1000000.0;
double avgTimePerFrame  = totalTimeSeconds / totalIterations;
double throughputFPS= totalIterations / totalTimeSeconds;

std::cout << "\n" << std::string(70, '=') << std::endl;
std::cout << "SIMPLE MULTI-STREAM RESULTS" << std::endl;
std::cout << std::string(70, '=') << std::endl;
std::cout << "Input: RGB8 -> Y8_BL_ER" << std::endl;
std::cout << "Total time: " << totalTimeSeconds << " seconds" << std::endl;
std::cout << "Avg time per frame: " << (avgTimePerFrame * 1000) << " ms" << std::endl;
std::cout << "THROUGHPUT: " << throughputFPS << " FPS" << std::endl;
std::cout << std::string(70, '=') << std::endl;

std::cout << "THROUGHPUT: " << throughputFPS << " FPS" << std::endl;
std::cout << std::string(70, '=') << std::endl;