ResampleSinogram3D.h

Go to the documentation of this file.

#pragma once
 
#include "Common.h"
 
namespace reg23 {
 
at::Tensor ResampleSinogram3D_CPU(const at::Tensor &sinogram3d, const std::string &sinogramType, double rSpacing,
                                  const at::Tensor &projectionMatrix, const at::Tensor &phiValues,
                                  const at::Tensor &rValues, c10::optional<at::Tensor> out);
 
__host__ at::Tensor ResampleSinogram3D_CUDA(const at::Tensor &sinogram3d, const std::string &sinogramType,
                                            double rSpacing, const at::Tensor &projectionMatrix,
                                            const at::Tensor &phiValues, const at::Tensor &rValues,
                                            c10::optional<at::Tensor> out);
 
__host__ at::Tensor ResampleSinogram3DCUDATexture(int64_t sinogram3dTextureHandle, int64_t sinogramWidth,
                                                  int64_t sinogramHeight, int64_t sinogramDepth,
                                                  const std::string &sinogramType, double rSpacing,
                                                  const at::Tensor &projectionMatrix, const at::Tensor &phiValues,
                                                  const at::Tensor &rValues, c10::optional<at::Tensor> out);
 
 
struct ResampleSinogram3D {
 
    enum class SinogramType { CLASSIC, HEALPIX };
 
    struct ConstantGeometry {
        Vec<float, 2> originProjection{};
        float squareRadius{};
        Vec<Vec<float, 4>, 4> projectionMatrixTranspose{};
    };
 
    struct CommonData {
        SinogramType sinogramType{};
        ConstantGeometry geometry{};
        at::Tensor flatOutput{};
    };
 
    __host__ static CommonData Common(const std::string &sinogramType, const at::Tensor &projectionMatrix,
                                      const at::Tensor &phiValues, const at::Tensor &rValues, at::DeviceType device,
                                      c10::optional<at::Tensor> out) {
        // projectionMatrix should be of size (4, 4), contain floats and be on the chosen device
        TORCH_CHECK(projectionMatrix.sizes() == at::IntArrayRef({4, 4}))
        TORCH_CHECK(projectionMatrix.dtype() == at::kFloat)
        TORCH_INTERNAL_ASSERT(projectionMatrix.device().type() == device)
        // phiValues and rValues should be of the same size, contain floats and be on the chosen device
        TORCH_CHECK(phiValues.sizes() == rValues.sizes())
        TORCH_CHECK(phiValues.dtype() == at::kFloat)
        TORCH_CHECK(rValues.dtype() == at::kFloat)
        TORCH_INTERNAL_ASSERT(phiValues.device().type() == device)
        TORCH_INTERNAL_ASSERT(rValues.device().type() == device);
        if (out) {
            TORCH_CHECK(out.value().sizes() == phiValues.sizes())
            TORCH_CHECK(out.value().dtype() == at::kFloat)
            TORCH_INTERNAL_ASSERT(out.value().device().type() == device)
        }
 
        SinogramType st = SinogramType::CLASSIC;
        if (sinogramType == "healpix") {
            st = SinogramType::HEALPIX;
        } else if (sinogramType != "classic") {
            TORCH_WARN("Invalid sinogram type string given. Valid values are: 'classic', 'healpix'. Assuming default "
                "value: 'classic'.")
        }
 
        const at::Tensor originProjectionHomogeneous = matmul(projectionMatrix,
                                                              torch::tensor(
                                                                  {{0.f, 0.f, 0.f, 1.f}},
                                                                  projectionMatrix.options()).t());
 
        CommonData ret{};
        ret.sinogramType = st;
        ret.geometry.originProjection = Vec<float, 2>{originProjectionHomogeneous[0].item().toFloat(),
                                                      originProjectionHomogeneous[1].item().toFloat()} /
                                        originProjectionHomogeneous[3].item().toFloat();
        ret.geometry.squareRadius = .25f * ret.geometry.originProjection.Apply<float>(&Square<float>).Sum();
        ret.geometry.projectionMatrixTranspose = Vec<Vec<float, 4>, 4>::FromTensor2D(projectionMatrix.t());
        ret.flatOutput = out
                             ? out.value().view({-1})
                             : torch::zeros(at::IntArrayRef({phiValues.numel()}), at::TensorOptions{device});
        return ret;
    }
 
    template <typename sinogram_t> __host__ __device__ static float ResamplePlane(
        const sinogram_t &sinogram, const ConstantGeometry &geometry, float phi, float r) {
        const float cp = std::cos(phi);
        const float sp = std::sin(phi);
        const Vec<float, 4> intermediate = MatMul(geometry.projectionMatrixTranspose, Vec<float, 4>{cp, sp, 0.f, -r});
        const Vec<float, 3> intermediate3 = intermediate.XYZ();
        const Vec<float, 3> posCartesian = -intermediate.W() * intermediate3 / intermediate3.Apply<float>(
                                               &Square<float>).Sum();
 
        Vec<double, 3> rThetaPhi{};
        rThetaPhi[2] = std::atan2(posCartesian.Y(), posCartesian.X());
        const float magXY = posCartesian.X() * posCartesian.X() + posCartesian.Y() * posCartesian.Y();
        rThetaPhi[1] = std::atan2(posCartesian.Z(), std::sqrt(magXY));
        rThetaPhi[0] = std::sqrt(magXY + posCartesian.Z() * posCartesian.Z());
        rThetaPhi = UnflipSphericalCoordinate(rThetaPhi);
 
        float ret = sinogram.Sample(rThetaPhi);
 
        if ((r * Vec<float, 2>{cp, sp} - .5f * geometry.originProjection).Apply<float>(&Square<float>).Sum() < geometry.
            squareRadius) {
            ret *= -1.f;
        }
        return ret;
    }
};
 
} // namespace reg23