ProjectDRRCuboidMaskCPU.h

Go to the documentation of this file.

#pragma once
 
#include "Common.h"
 
namespace reg23 {
 
at::Tensor ProjectDRRCuboidMask_CPU(const at::Tensor &volumeSize, const at::Tensor &voxelSpacing,
                                    const at::Tensor &homographyMatrixInverse, double sourceDistance,
                                    int64_t outputWidth, int64_t outputHeight, const at::Tensor &outputOffset,
                                    const at::Tensor &detectorSpacing);
 
__host__ at::Tensor ProjectDRRCuboidMask_CUDA(const at::Tensor &volumeSize, const at::Tensor &voxelSpacing,
                                              const at::Tensor &homographyMatrixInverse, double sourceDistance,
                                              int64_t outputWidth, int64_t outputHeight, const at::Tensor &outputOffset,
                                              const at::Tensor &detectorSpacing);
 
 
template <typename T, std::size_t faceCount> __host__ __device__ T RayConvexPolyhedronDistance(
    const std::array<Vec<T, 3>, faceCount> &facePoints, const std::array<Vec<T, 3>, faceCount> &faceOutUnitNormals,
    const Vec<T, 3> &rayPoint, const Vec<T, 3> &rayUnitDirection) {
 
    constexpr T epsilon = 1e-12;
 
    T entryLambda = -std::numeric_limits<T>::infinity();
    T exitLambda = std::numeric_limits<T>::infinity();
    // std::size_t entryIndex = 0;
    for (std::size_t i = 0; i < faceCount; ++i) {
        const T dot = VecDot(faceOutUnitNormals[i], rayUnitDirection);
        const T numer = VecDot(faceOutUnitNormals[i], facePoints[i] - rayPoint);
        // If dot is 0, the ray is parallel to the face; if dot is negative, increasing lambda moves you along the ray
        // in the 'inward' direction relative to this face, and vice versa if dot is positive.
        if (std::abs(dot) < epsilon) {
            if (numer < -epsilon) return 0.0f;
            continue;
        }
        T lambda = numer / dot;
        if (dot > epsilon) {
            // Moving in 'outward' direction, so could be exiting the polyhedron through this face
            if (lambda < exitLambda) exitLambda = lambda;
        } else {
            // Moving in 'inward' direction, so could be entering the polyhedron through this face
            if (lambda > entryLambda) entryLambda = lambda;
        }
        // early rejection
        if (entryLambda > exitLambda) return 0.0f;
    }
    // `entryLambda` is now the largest value for which ray is moving 'into' face, and `entryIndex` is index of
    // corresponding face. `exitLambda` is the smallest value for which ray is moving 'out of' face. This means that,
    // **if the ray intersects with the polyhedron**, it will enter at `entryLambda` and leave at `exitLambda`.
 
    return exitLambda > entryLambda ? exitLambda - entryLambda : 0.0f;
 
    // const Vec<T, 3> centrePoint = rayPoint + 0.5 * (entryLambda + exitLambda) * rayUnitDirection;
    // // Checking that the centre point between the entry and exit points lies 'within' all the other faces of the
    // // polyhedron. If it lies 'outside' of any,  it doesn't intersect the polyhedron. Technically, could just check the
    // // entry point against only the faces adjacent to the intersected face, but that would require expensive analysis of
    // // the face intersections.
    // for (std::size_t i = 0; i < faceCount; ++i) {
    //  // if (i == entryIndex) continue;
    //  if (VecDot(centrePoint - facePoints[i], faceOutUnitNormals[i]) > epsilon) return 0.f;
    // }
    //
    // return std::abs(exitLambda - entryLambda);
}
 
 
struct ProjectDRRCuboidMask {
 
    struct Cuboid {
        Vec<Vec<float, 3>, 6> facePoints{};
        Vec<Vec<float, 3>, 6> faceOutUnitNormals{};
    };
 
    struct CommonData {
        Vec<Vec<float, 4>, 4> homographyMatrixInverse{};
        Vec<float, 2> outputOffset{};
        Vec<float, 2> detectorSpacing{};
        Vec<float, 3> sourcePositionTransformed{};
        Cuboid cuboidIn{};
        Cuboid cuboidAbove{};
        Cuboid cuboidBelow{};
        at::Tensor flatOutput{};
    };
 
    __host__ static CommonData Common(const at::Tensor &volumeSize, const at::Tensor &voxelSpacing,
                                      const at::Tensor &homographyMatrixInverse, double sourceDistance,
                                      int64_t outputWidth, int64_t outputHeight, const at::Tensor &outputOffset,
                                      const at::Tensor &detectorSpacing, at::DeviceType device) {
        // volumeSize should be a 1D tensor of 3 longs on the chosen device
        TORCH_CHECK(volumeSize.sizes() == at::IntArrayRef{3});
        TORCH_CHECK(volumeSize.dtype() == at::kLong);
        TORCH_INTERNAL_ASSERT(volumeSize.device().type() == device);
        // voxelSpacing should be a 1D tensor of 3 doubles on the chosen device
        TORCH_CHECK(voxelSpacing.sizes() == at::IntArrayRef{3});
        TORCH_CHECK(voxelSpacing.dtype() == at::kDouble);
        TORCH_INTERNAL_ASSERT(voxelSpacing.device().type() == device);
        // homographyMatrixInverse should be of size (4, 4), contain doubles and be on the chosen device
        TORCH_CHECK(homographyMatrixInverse.sizes() == at::IntArrayRef({4, 4}));
        TORCH_CHECK(homographyMatrixInverse.dtype() == at::kDouble);
        TORCH_INTERNAL_ASSERT(homographyMatrixInverse.device().type() == device);
        // outputOffset should be a 1D tensor of 2 doubles on the chosen device
        TORCH_CHECK(outputOffset.sizes() == at::IntArrayRef{2});
        TORCH_CHECK(outputOffset.dtype() == at::kDouble);
        TORCH_INTERNAL_ASSERT(outputOffset.device().type() == device);
        // detectorSpacing should be a 1D tensor of 2 doubles on the chosen device
        TORCH_CHECK(detectorSpacing.sizes() == at::IntArrayRef{2});
        TORCH_CHECK(detectorSpacing.dtype() == at::kDouble);
        TORCH_INTERNAL_ASSERT(detectorSpacing.device().type() == device);
 
        CommonData ret{};
        ret.homographyMatrixInverse = Vec<Vec<float, 4>, 4>::FromTensor2D(homographyMatrixInverse);
        const Vec<float, 6> planeSigns = Vec<int, 6>{1, 1, 1, -1, -1, -1}.StaticCast<float>();
        const Vec<float, 3> cuboidHalfSize = 0.5f * Vec<float, 3>::FromTensor(voxelSpacing) * Vec<
            int64_t, 3>::FromTensor(volumeSize).StaticCast<float>();
        ret.cuboidIn.facePoints = VecOuter(cuboidHalfSize, planeSigns);
        ret.cuboidIn.faceOutUnitNormals = VecCat(Vec<Vec<float, 3>, 3>::Identity(),
                                                 -1.f * Vec<Vec<float, 3>, 3>::Identity());
 
        const Vec<float, 3> aboveBelowHalfSize = Vec<float, 3>{1.f, 1.f, 4.f} * cuboidHalfSize;
        ret.cuboidAbove.facePoints = VecOuter(aboveBelowHalfSize, planeSigns);
        ret.cuboidBelow.facePoints = ret.cuboidAbove.facePoints;
        const float zSum = aboveBelowHalfSize.Z() + cuboidHalfSize.Z();
        for (Vec<float, 3> &v : ret.cuboidAbove.facePoints) v.Z() += zSum;
        for (Vec<float, 3> &v : ret.cuboidBelow.facePoints) v.Z() -= zSum;
        ret.cuboidAbove.faceOutUnitNormals = ret.cuboidIn.faceOutUnitNormals;
        ret.cuboidBelow.faceOutUnitNormals = ret.cuboidIn.faceOutUnitNormals;
 
        ret.sourcePositionTransformed = MatMul(ret.homographyMatrixInverse,
                                               Vec<float, 4>{0.0f, 0.0f, static_cast<float>(sourceDistance), 1.0f}).
            XYZ();
        ret.outputOffset = Vec<float, 2>::FromTensor(outputOffset);
        ret.detectorSpacing = Vec<float, 2>::FromTensor(detectorSpacing);
        ret.flatOutput = torch::zeros(at::IntArrayRef({outputWidth * outputHeight}),
                                      at::TensorOptions{at::Device{device}});
        return ret;
    }
};
 
} // namespace reg23