Example #1
0
TensorWrapper::TensorWrapper(cuda::GpuMat & mat, THCState *state) {

    this->definedInLua = false;

    if (mat.empty()) {
        this->typeCode = CV_CUDA;
        this->tensorPtr = nullptr;
        return;
    }

    this->typeCode = CV_CUDA;

    THCudaTensor *outputPtr = THCudaTensor_new(state);

    // Build new storage on top of the Mat
    outputPtr->storage = THCudaStorage_newWithData(
            state,
            reinterpret_cast<float *>(mat.data),
            mat.step * mat.rows * mat.channels() / cv::getElemSize(mat.depth())
    );

    int sizeMultiplier;
    if (mat.channels() == 1) {
        outputPtr->nDimension = 2;
        sizeMultiplier = cv::getElemSize(mat.depth());
    } else {
        outputPtr->nDimension = 3;
        sizeMultiplier = mat.elemSize1();
    }

    outputPtr->size = static_cast<long *>(THAlloc(sizeof(long) * outputPtr->nDimension));
    outputPtr->stride = static_cast<long *>(THAlloc(sizeof(long) * outputPtr->nDimension));

    if (mat.channels() > 1) {
        outputPtr->size[2] = mat.channels();
        outputPtr->stride[2] = 1;
    }

    outputPtr->size[0] = mat.rows;
    outputPtr->size[1] = mat.cols;

    outputPtr->stride[0] = mat.step / sizeMultiplier;
    outputPtr->stride[1] = mat.channels();

    outputPtr->storageOffset = 0;

    // Make OpenCV treat underlying data as user-allocated
    mat.refcount = nullptr;

    this->tensorPtr = reinterpret_cast<THByteTensor *>(outputPtr);
}