sync

cpp/models/FsNeRF.cpp

+#include "FsNeRF.h"
+
+namespace models
+{
+    FsNeRF::FsNeRF(const Args &args, int nRays)
+        : _nRays(nRays),
+          _nSamples(args.nSamples),
+          _xChns(3 - !args.withRadius),
+          _field(new fields::FsNeRF(args.modelPath)),
+          _sampler(new modules::Sampler({args.near, args.far}, "spherical_radius", args.nSamples, args.withRadius)),
+          _encoder(new modules::Encoder(args.xfreqs, 3 - !args.withRadius, false)),
+          _renderer(new modules::Renderer(args.whiteBg))
+    {
+        auto n = _nRays * _nSamples;
+        _x = darray<float>(new CudaArray<float>(n * _xChns));
+        _depths = darray<float>(new CudaArray<float>(n));
+        _encoded = darray<float>(new CudaArray<float>(n * _encoder->outChns()));
+        _rgbd = darray<glm::vec4>(new CudaArray<glm::vec4>(n));
+        _field->bindResources(_encoded.get(), _depths.get(), _rgbd.get());
+    }
+
+    FsNeRF::~FsNeRF()
+    {
+        delete _sampler;
+        delete _encoder;
+        delete _renderer;
+    }
+
+    void FsNeRF::operator()(darray<glm::vec4> o_colors, const darray<glm::vec3> dirs,
+                            glm::vec3 origin, bool showPerf = false)
+    {
+        CudaEvent eStart, eSampled, eEncoded, eInferred, eRendered;
+
+        cudaEventRecord(eStart);
+
+        (*_sampler)(_x, _depths, origin, darray<glm::vec3>(dirs->subArray(0, _nRays)));
+        CHECK_EX(cudaDeviceSynchronize());
+
+        cudaEventRecord(eSampled);
+
+        (*_encoder)(_encoded, _x);
+        CHECK_EX(cudaDeviceSynchronize());
+
+        cudaEventRecord(eEncoded);
+
+        _field->infer();
+        CHECK_EX(cudaDeviceSynchronize());
+
+        cudaEventRecord(eInferred);
+
+        (*_renderer)(darray<glm::vec4>(o_colors->subArray(0, _nRays)), _depths, _rgbd);
+        CHECK_EX(cudaDeviceSynchronize());
+
+        cudaEventRecord(eRendered);
+
+        if (showPerf)
+        {
+            CHECK_EX(cudaDeviceSynchronize());
+
+            float timeTotal, timeSample, timeEncode, timeInfer, timeRender;
+            cudaEventElapsedTime(&timeTotal, eStart, eRendered);
+            cudaEventElapsedTime(&timeSample, eStart, eSampled);
+            cudaEventElapsedTime(&timeEncode, eSampled, eEncoded);
+            cudaEventElapsedTime(&timeInfer, eEncoded, eInferred);
+            cudaEventElapsedTime(&timeRender, eInferred, eRendered);
+
+            std::ostringstream sout;
+            sout << "Infer pipeline: " << timeTotal << "ms (Sample: " << timeSample
+                 << "ms, Encode: " << timeEncode << "ms, Infer: " << timeInfer
+                 << "ms, Render: " << timeRender << "ms)";
+            Logger::instance.info(sout.str().c_str());
+        }
+
+        /*
+        {
+            std::ostringstream sout;
+            sout << "Rays:" << std::endl;
+            dumpArray<glm::vec3, float>(sout, *rays, 10);
+            Logger::instance.info(sout.str());
+        }
+        {
+            std::ostringstream sout;
+            sout << "Spherical coords:" << std::endl;
+            dumpArray(sout, *_coords, 10, _xChns * _nSamples);
+            Logger::instance.info(sout.str());
+        }
+        {
+            std::ostringstream sout;
+            sout << "Depths:" << std::endl;
+            dumpArray(sout, *_depths, 10, _nSamples);
+            Logger::instance.info(sout.str());
+        }
+        {
+            std::ostringstream sout;
+            sout << "Encoded:" << std::endl;
+            dumpArray(sout, *_encoded, 10, _encoder->outDim() * _nSamples);
+            Logger::instance.info(sout.str());
+        }
+        {
+            std::ostringstream sout;
+            sout << "Color:" << std::endl;
+            dumpArray<glm::vec4, float>(sout, *o_colors, 10);
+            Logger::instance.info(sout.str());
+        }
+        */
+    }
+}
\ No newline at end of file

cpp/models/FsNeRF.h

+#pragma once
+#include "../utils/common.h"
+#include "../modules/Sampler.h"
+#include "../modules/Encoder.h"
+#include "../modules/Renderer.h"
+#include "../fields/FsNeRF.h"
+#include "Model.h"
+
+namespace models
+{
+    class FsNeRF : public Model
+    {
+    public:
+        class Args
+        {
+        public:
+            std::string modelPath;
+            uint nSamples;
+            uint withRadius;
+            uint xfreqs;
+            float near;
+            float far;
+            bool whiteBg;
+        };
+
+        FsNeRF(const Args &args, int nRays);
+
+        ~FsNeRF();
+
+        virtual void operator()(darray<glm::vec4> o_colors, const darray<glm::vec3> dirs,
+                                glm::vec3 origin, bool showPerf = false);
+
+        uint nRays() const { return _nRays; }
+
+    private:
+        uint _nRays;
+        uint _nSamples;
+        uint _xChns;
+        fields::FsNeRF *_field;
+        modules::Sampler *_sampler;
+        modules::Encoder *_encoder;
+        modules::Renderer *_renderer;
+        darray<float> _x;
+        darray<float> _depths;
+        darray<float> _encoded;
+        darray<glm::vec4> _rgbd;
+    };
+}
\ No newline at end of file

cpp/models/Model.h

+#pragma once
+#include "../utils/common.h"
+
+namespace models
+{
+    class Model
+    {
+        virtual void operator()(darray<glm::vec4> o_colors, const darray<glm::vec3> dirs,
+                                glm::vec3 origin, bool showPerf = false) = 0;
+    };
+}
\ No newline at end of file

cpp/modules/Encoder.cu

+#include "Encoder.h"
+#include "../utils/cuda.h"
+
+/// idx3.z = 0: x, y, z, sin(x), sin(y), sin(z), cos(x), cos(y), cos(z)
+/// idx3.z = 1: sin(2x), sin(2y), sin(2z), cos(2x), cos(2y), cos(2z)
+/// ...
+/// idx3.z = n_freq-1: sin(2^(n_freq-1)x), sin(2^(n_freq-1)y), sin(2^(n_freq-1)z),
+///                    cos(2^(n_freq-1)x), cos(2^(n_freq-1)y), cos(2^(n_freq-1)z)
+/// Dispatch (n, in_chns, n_freqs)
+
+__global__ void cu_encode(float *output, const float *input, const float *freqs, uint n,
+                          bool includeInput)
+{
+    DEFINE_IDX3(i, chn, freq);
+    if (i >= n)
+        return;
+    uint offset = (uint)includeInput;
+    uint inChns = blockDim.y, nFreqs = blockDim.z;
+    uint elem = i * inChns + chn;
+    uint outChns = inChns * (nFreqs * 2 + offset);
+    uint base = i * outChns + chn;
+    if (freq == 0 && includeInput)
+        output[base] = input[elem];
+    float x = freqs[freq] * input[elem];
+    float s, c;
+    __sincosf(x, &s, &c);
+    output[base + inChns * (freq * 2 + offset)] = s;
+    output[base + inChns * (freq * 2 + offset + 1)] = c;
+}
+
+namespace modules
+{
+    /**
+     * @brief
+     *
+     * @param output (n, out_chns) encoded positions
+     * @param input (n, in_chns) positions
+     */
+    void Encoder::operator()(darray<float> output, const darray<float> input)
+    {
+        uint n = input->n() / _chns;
+        std::ostringstream sout;
+        sout << "Encoder => input size: (" << n << ", " << _chns << "), "
+             << "output size: (" << n << ", " << outChns() << ")";
+        // Logger::instance.info(sout.str());
+        dim3 blkSize(1024 / _chns / _multires, _chns, _multires);
+        dim3 grdSize(ceilDiv(n, blkSize.x), 1, 1);
+        CU_INVOKE(cu_encode)(*output, *input, *_freqs, n, _includeInput);
+        CHECK_EX(cudaGetLastError());
+    }
+
+    void Encoder::_genFreqArray()
+    {
+        std::vector<float> freqsHost(_multires);
+        freqsHost[0] = 1.0f;
+        for (auto i = 1; i < _multires; ++i)
+            freqsHost[i] = freqsHost[i - 1] * 2.0f;
+        _freqs = darray<float>(new CudaArray<float>(freqsHost));
+    }
+}
\ No newline at end of file

cpp/modules/Encoder.h

+#pragma once
+#include "../utils/common.h"
+
+namespace modules
+{
+    class Encoder
+    {
+    public:
+        Encoder(unsigned int multires, unsigned int chns, bool includeInput)
+            : _multires(multires), _chns(chns), _includeInput(includeInput)
+        {
+            _genFreqArray();
+        }
+
+        unsigned int outChns() const { return _chns * (_includeInput + _multires * 2); }
+
+        void operator()(darray<float> output, const darray<float> input);
+
+    private:
+        unsigned int _multires;
+        unsigned int _chns;
+        bool _includeInput;
+        darray<float> _freqs;
+
+        void _genFreqArray();
+    };
+}
\ No newline at end of file

cpp/modules/Renderer.cu

+#include "Renderer.h"
+#include "../utils/cuda.h"
+
+/// Dispatch (nRays)
+__global__ void cu_render(glm::vec4 *o_colors, const float *depths, const glm::vec4 *rgbd,
+                          uint nSamples, uint nRays)
+{
+    DEFINE_IDX(rayIdx);
+    if (rayIdx >= nRays)
+        return;
+    glm::vec4 outColor;
+    auto depth = 1e10f;
+    for (int si = nSamples - 1, i = (rayIdx + 1) * nSamples - 1; si >= 0; --si, --i)
+    {
+        auto depth1 = depth;
+        auto c = rgbd[i];
+        depth = depths[i];
+        c.a = 1.0f - exp(-max(c.a, 0.f) * (depth1 - depth));
+        outColor = outColor * (1 - c.a) + c * c.a;
+    }
+    outColor.a = 1.0f;
+    o_colors[rayIdx] = outColor;
+}
+
+namespace modules
+{
+    void Renderer::operator()(darray<glm::vec4> o_colors, const darray<float> depths,
+                              const darray<glm::vec4> rgbd)
+    {
+        uint n = o_colors->n();
+        dim3 blkSize(1024);
+        dim3 grdSize(ceilDiv(n, blkSize.x));
+        CU_INVOKE(cu_render)(*o_colors, *depths, *rgbd, rgbd->n() / n, n);
+        CHECK_EX(cudaGetLastError());
+    }
+}
\ No newline at end of file

cpp/modules/Renderer.h

+#pragma once
+#include "../utils/common.h"
+
+namespace modules
+{
+    class Renderer
+    {
+    public:
+        Renderer(bool whiteBg) : _whiteBg(whiteBg) {}
+
+        /**
+         * @brief
+         *
+         * @param o_colors
+         * @param layeredColors
+         */
+        void operator()(darray<glm::vec4> o_colors, const darray<float> depths,
+                        const darray<glm::vec4> rgbd);
+
+    private:
+        bool _whiteBg;
+    };
+}
\ No newline at end of file

cpp/modules/Sampler.cu

+#include "Sampler.h"
+#define _USE_MATH_DEFINES
+#include <math.h>
+#include "../utils/cuda.h"
+
+__device__ glm::vec3 _raySphereIntersect(glm::vec3 p, glm::vec3 v, float r, float &o_depth)
+{
+    float pp = glm::dot(p, p);
+    float vv = glm::dot(v, v);
+    float pv = glm::dot(p, v);
+    o_depth = (sqrtf(pv * pv - vv * (pp - r * r)) - pv) / vv;
+    return p + o_depth * v;
+}
+
+__device__ float _getAngle(float x, float y)
+{
+    return -atan(x / y) - (y < 0) * (float)M_PI + 0.5f * (float)M_PI;
+}
+
+/**
+ * Dispatch with block_size=(*, n_samples), grid_size=(nRays/*, 1)
+ * Index with (ray_idx, sample_idx)
+ */
+__global__ void cu_sampleXyz(glm::vec3 *o_coords, float *o_depths, glm::vec3 origin,
+                             const glm::vec3 *dirs, uint nRays, uint nSamples, glm::vec2 range)
+{
+    DEFINE_IDX2(rayIdx, sampleIdx);
+    if (rayIdx >= nRays)
+        return;
+    uint idx = rayIdx * nSamples + sampleIdx;
+    float z = (range.y - range.x) / (nSamples - 1) * sampleIdx + range.x;
+    o_coords[idx] = origin + dirs[rayIdx] * z;
+    o_depths[idx] = z;
+}
+
+/**
+ * Dispatch with block_size=(*, n_samples), grid_size=(nRays/*, 1)
+ * Index with (ray_idx, sample_idx)
+ */
+__global__ void cu_sampleXyzDisp(glm::vec3 *o_coords, float *o_depths, glm::vec3 origin,
+                                 const glm::vec3 *dirs, uint nRays, uint nSamples, glm::vec2 range)
+{
+    DEFINE_IDX2(rayIdx, sampleIdx);
+    if (rayIdx >= nRays)
+        return;
+    uint idx = rayIdx * nSamples + sampleIdx;
+    float z = 1.f / ((range.y - range.x) / (nSamples - 1) * sampleIdx + range.x);
+    o_coords[idx] = origin + dirs[rayIdx] * z;
+    o_depths[idx] = z;
+}
+
+/**
+ * Dispatch with block_size=(*, n_samples), grid_size=(nRays/*, 1)
+ * Index with (ray_idx, sample_idx)
+ */
+__global__ void cu_sampleSpherical(glm::vec3 *o_coords, float *o_depths, glm::vec3 origin,
+                                   const glm::vec3 *dirs, uint nRays, uint nSamples,
+                                   glm::vec2 range)
+{
+    DEFINE_IDX2(rayIdx, sampleIdx);
+    if (rayIdx >= nRays)
+        return;
+    uint idx = rayIdx * nSamples + sampleIdx;
+    float z = 1.f / ((range.y - range.x) / (nSamples - 1) * sampleIdx + range.x);
+    glm::vec3 p = origin + dirs[rayIdx] * z;
+    float r_reci = 1. / glm::length(p);
+    float theta = _getAngle(p.x, p.z);
+    float phi = acos(p.y * r_reci);
+    o_coords[idx] = {r_reci, theta, phi};
+    o_depths[idx] = z;
+}
+
+/**
+ * Dispatch with block_size=(*, n_samples), grid_size=(nRays/*, 1)
+ * Index with (ray_idx, sample_idx)
+ */
+__global__ void cu_sampleSphere(float *o_coords, float *o_depths, glm::vec3 origin,
+                                const glm::vec3 *dirs, uint nRays, uint nSamples,
+                                glm::vec2 range, bool withRadius)
+{
+    DEFINE_IDX2(rayIdx, sampleIdx);
+    if (rayIdx >= nRays)
+        return;
+    uint idx = rayIdx * nSamples + sampleIdx;
+    float r_reci = (range.y - range.x) / (nSamples - 1) * sampleIdx + range.x;
+    float r = 1.f / r_reci;
+    glm::vec3 p = _raySphereIntersect(origin, dirs[rayIdx], r, o_depths[idx]);
+    float theta = _getAngle(p.x, p.z);
+    float phi = acos(p.y * r_reci);
+    if (withRadius)
+        ((glm::vec3 *)o_coords)[idx] = {r_reci, theta, phi};
+    else
+        ((glm::vec2 *)o_coords)[idx] = {theta, phi};
+}
+
+namespace modules
+{
+    void Sampler::operator()(darray<float> o_coords, darray<float> o_depths, glm::vec3 origin,
+                             const darray<glm::vec3> dirs)
+    {
+        auto n = dirs->n();
+        dim3 blkSize(1024 / _nSamples, _nSamples);
+        dim3 grdSize((uint)ceil(n / (float)blkSize.y), 1);
+        if (_mode == "xyz_disp")
+            CU_INVOKE(cu_sampleXyzDisp)((glm::vec3 *)*o_coords, *o_depths, origin, *dirs, n,
+                                        _nSamples, 1.f / _range);
+        else if (_mode == "spherical")
+            CU_INVOKE(cu_sampleSpherical)((glm::vec3 *)*o_coords, *o_depths, origin, *dirs, n,
+                                          _nSamples, 1.f / _range);
+        else if (_mode == "spherical_radius")
+            CU_INVOKE(cu_sampleSphere)(*o_coords, *o_depths, origin, *dirs, n, _nSamples,
+                                       1.f / _range, _withRadius);
+        else
+            CU_INVOKE(cu_sampleXyz)((glm::vec3 *)*o_coords, *o_depths, origin, *dirs, n,
+                                    _nSamples, _range);
+        CHECK_EX(cudaGetLastError());
+    }
+}
\ No newline at end of file

cpp/modules/Sampler.h

+#pragma once
+#include "../utils/common.h"
+
+namespace modules
+{
+    class Sampler
+    {
+    public:
+        Sampler(glm::vec2 range, std::string mode, uint nSamples, bool withRadius) : _range(range), _mode(mode), _nSamples(nSamples), _withRadius(withRadius) {}
+
+        void operator()(darray<float> o_coords, darray<float> o_depths, glm::vec3 origin,
+                        const darray<glm::vec3> dirs);
+
+    private:
+        glm::vec2 _range;
+        std::string _mode;
+        uint _nSamples;
+        bool _withRadius;
+    };
+}
\ No newline at end of file