sync

cpp/fnr_core/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=(1, nRays/*)
+ * Index with (sample_idx, ray_idx)
+ */
+__global__ void cu_sampleOnRays(float *o_coords, float *o_depths, glm::vec3 *rays, uint nRays,
+                                glm::vec3 origin, Range range, bool outputRadius) {
+    glm::uvec3 idx3 = IDX3;
+    uint idx = flattenIdx(idx3);
+    uint sampleIdx = idx3.x;
+    uint rayIdx = idx3.y;
+    if (rayIdx >= nRays)
+        return;
+    float r_reciprocal = range.get(sampleIdx);
+    glm::vec3 p = _raySphereIntersect(origin, rays[rayIdx], 1.0f / r_reciprocal, o_depths[idx]);
+    glm::vec3 sp(r_reciprocal, _getAngle(p.x, p.z), acos(p.y * r_reciprocal));
+    if (outputRadius)
+        ((glm::vec3 *)o_coords)[idx] = sp;
+    else
+        ((glm::vec2 *)o_coords)[idx] = {sp.y, sp.z};
+}
+
+void Sampler::sampleOnRays(sptr<CudaArray<float>> o_coords, sptr<CudaArray<float>> o_depths,
+                           sptr<CudaArray<glm::vec3>> rays, glm::vec3 rayCenter) {
+    dim3 blkSize(_dispRange.steps(), 1024 / _dispRange.steps());
+    dim3 grdSize(1, (uint)ceil(rays->n() / (float)blkSize.y));
+    CU_INVOKE(cu_sampleOnRays)
+    (*o_coords, *o_depths, *rays, rays->n(), rayCenter, _dispRange, _outputRadius);
+    CHECK_EX(cudaGetLastError());
+}
\ No newline at end of file

cpp/fnr_core/View.cu

+#include "View.h"
+#define _USE_MATH_DEFINES
+#include <math.h>
+#include <cuda_runtime.h>
+#include <fstream>
+#include "../utils/cuda.h"
+
+__global__ void cu_transPoints(glm::vec3 *o_vecs, glm::vec3 *vecs, glm::vec3 t, glm::mat3 rot_t,
+                               uint n) {
+    uint idx = flattenIdx();
+    if (idx >= n)
+        return;
+    o_vecs[idx] = vecs[idx] * rot_t + t;
+}
+
+__global__ void cu_transPoints(glm::vec3 *o_vecs, glm::vec3 *vecs, glm::vec3 t, glm::mat3 rot_t,
+                               uint n, int *indices) {
+    uint idx = flattenIdx();
+    if (idx >= n)
+        return;
+    o_vecs[idx] = vecs[indices[idx]] * rot_t + t;
+}
+
+__global__ void cu_transPointsInverse(glm::vec3 *o_pts, glm::vec3 *pts, glm::vec3 t,
+                                      glm::mat3 inv_rot_t, uint n) {
+    uint idx = flattenIdx();
+    if (idx >= n)
+        return;
+    o_pts[idx] = (pts[idx] - t) * inv_rot_t;
+}
+
+__global__ void cu_transPointsInverse(glm::vec3 *o_pts, glm::vec3 *pts, glm::vec3 t,
+                                      glm::mat3 inv_rot_t, uint n, int *indices) {
+    uint idx = flattenIdx();
+    if (idx >= n)
+        return;
+    o_pts[idx] = (pts[indices[idx]] - t) * inv_rot_t;
+}
+
+__global__ void cu_transVectors(glm::vec3 *o_vecs, glm::vec3 *vecs, glm::mat3 rot_t, uint n) {
+    uint idx = flattenIdx();
+    if (idx >= n)
+        return;
+    o_vecs[idx] = vecs[idx] * rot_t;
+}
+
+__global__ void cu_transVectors(glm::vec3 *o_vecs, glm::vec3 *vecs, glm::mat3 rot_t, uint n,
+                                int *indices) {
+    uint idx = flattenIdx();
+    if (idx >= n)
+        return;
+    o_vecs[idx] = vecs[indices[idx]] * rot_t;
+}
+
+__global__ void cu_genLocalRays(glm::vec3 *o_rays, glm::vec2 f, glm::vec2 c, glm::uvec2 res) {
+    glm::uvec2 idx2 = IDX2;
+    if (idx2.x >= res.x || idx2.y >= res.y)
+        return;
+    uint idx = idx2.x + idx2.y * res.x;
+    o_rays[idx] = glm::vec3((glm::vec2(idx2) - c) / f, 1.0f);
+}
+
+__global__ void cu_genLocalRaysNormed(glm::vec3 *o_rays, glm::vec2 f, glm::vec2 c, glm::uvec2 res) {
+    glm::uvec2 idx2 = IDX2;
+    if (idx2.x >= res.x || idx2.y >= res.y)
+        return;
+    uint idx = idx2.x + idx2.y * res.x;
+    o_rays[idx] = glm::normalize(glm::vec3((glm::vec2(idx2) - c) / f, 1.0f));
+}
+
+__global__ void cu_camGetRays(glm::vec3 *o_vecs, glm::vec3 *vecs, glm::vec3 offset, glm::mat3 rot_t, uint n) {
+	uint idx = flattenIdx();
+	if (idx >= n)
+		return;
+	o_vecs[idx] = glm::normalize(vecs[idx] + offset) * rot_t;
+}
+
+__global__ void cu_camGetRays(glm::vec3 *o_vecs, glm::vec3 *vecs, glm::vec3 offset, glm::mat3 rot_t, uint n,
+	int *indices) {
+	uint idx = flattenIdx();
+	if (idx >= n)
+		return;
+	o_vecs[idx] = (vecs[indices[idx]] + offset) * rot_t;
+}
+
+__global__ void cu_indexedCopy(glm::vec4 *o_colors, glm::vec4 *colors, int *indices, uint n) {
+    uint idx = flattenIdx();
+    if (idx >= n)
+        return;
+    int srcIdx = indices[idx];
+    o_colors[idx] = srcIdx >= 0 ? colors[srcIdx] : glm::vec4();
+}
+
+void View::transPoints(sptr<CudaArray<glm::vec3>> results, sptr<CudaArray<glm::vec3>> points,
+                       sptr<CudaArray<int>> indices, bool inverse) {
+    glm::mat3 r_t = inverse ? _r : glm::transpose(_r);
+    dim3 blkSize(1024);
+    dim3 grdSize(ceilDiv(results->n(), blkSize.x));
+    if (inverse) {
+        if (indices == nullptr)
+            CU_INVOKE(cu_transPointsInverse)(*results, *points, _t, r_t, points->n());
+        else
+            CU_INVOKE(cu_transPointsInverse)(*results, *points, _t, r_t, points->n(), *indices);
+    } else {
+        if (indices == nullptr)
+            CU_INVOKE(cu_transPoints)(*results, *points, _t, r_t, results->n());
+        else
+            CU_INVOKE(cu_transPoints)(*results, *points, _t, r_t, results->n(), *indices);
+    }
+}
+
+void View::transVectors(sptr<CudaArray<glm::vec3>> results, sptr<CudaArray<glm::vec3>> vectors,
+                        sptr<CudaArray<int>> indices, bool inverse) {
+    glm::mat3 r_t = inverse ? _r : glm::transpose(_r);
+    dim3 blkSize(1024);
+    dim3 grdSize(ceilDiv(results->n(), blkSize.x));
+    if (indices == nullptr)
+        CU_INVOKE(cu_transVectors)(*results, *vectors, r_t, results->n());
+    else
+        CU_INVOKE(cu_transVectors)(*results, *vectors, r_t, results->n(), *indices);
+}
+
+View View::getStereoEye(float ipd, Eye eye) {
+    glm::vec3 eyeOffset((eye == Eye_Left ? -ipd : ipd) / 2.0f, 0.0f, 0.0f);
+    return View(_r * eyeOffset + _t, _r);
+}
+
+Camera::Camera(glm::uvec2 res, float fov) {
+    _f.x = _f.y = 0.5f * res.y / tan(fov * (float)M_PI / 360.0f);
+    _f.y *= -1.0f;
+    _c = res / 2u;
+    _res = res;
+}
+
+Camera::Camera(glm::uvec2 res, float fov, glm::vec2 c) {
+    _f.x = _f.y = 0.5f * res.y / tan(fov * (float)M_PI / 360.0f);
+    _f.y *= -1.0f;
+    _c = c;
+    _res = res;
+}
+
+sptr<CudaArray<glm::vec3>> Camera::localRays() {
+    if (_localRays == nullptr)
+        _genLocalRays(false);
+    return _localRays;
+}
+
+bool Camera::loadMaskData(std::string filepath) {
+    std::ifstream fin(filepath, std::ios::binary);
+    if (!fin)
+        return false;
+    int n;
+
+    fin.read((char *)&n, sizeof(n));
+    std::vector<int> subsetIndicesBuffer(n);
+    fin.read((char *)subsetIndicesBuffer.data(), sizeof(int) * n);
+    _subsetIndices.reset(new CudaArray<int>(subsetIndicesBuffer));
+
+    fin.read((char *)&n, sizeof(n));
+    std::vector<int> subsetInverseIndicesBuffer(n);
+    fin.read((char *)subsetInverseIndicesBuffer.data(), sizeof(int) * n);
+    _subsetInverseIndices.reset(new CudaArray<int>(subsetInverseIndicesBuffer));
+
+    if (!fin) {
+        _subsetIndices = nullptr;
+        _subsetInverseIndices = nullptr;
+        return false;
+    }
+    Logger::instance.info("Mask data loaded. Subset indices: %d, subset inverse indices: %d",
+        _subsetIndices->n(), _subsetInverseIndices->n());
+    return true;
+}
+
+void Camera::getRays(sptr<CudaArray<glm::vec3>> o_rays, View &view, glm::vec3 offset) {
+	glm::mat3 r_t = glm::transpose(view.r());
+	dim3 blkSize(1024);
+	dim3 grdSize(ceilDiv(nRays(), blkSize.x));
+	if (_subsetIndices == nullptr)
+		CU_INVOKE(cu_camGetRays)(*o_rays, *localRays(), offset, r_t, nRays());
+	else
+		CU_INVOKE(cu_camGetRays)(*o_rays, *localRays(), offset, r_t, nRays(), *_subsetIndices);
+}
+
+void Camera::restoreImage(sptr<CudaArray<glm::vec4>> o_imgData, sptr<CudaArray<glm::vec4>> colors) {
+    if (_subsetInverseIndices == nullptr) {
+        cudaMemcpy(o_imgData->getBuffer(), colors->getBuffer(), nPixels() * sizeof(glm::vec4),
+                   cudaMemcpyDeviceToDevice);
+    } else {
+        dim3 blkSize(1024);
+        dim3 grdSize(ceilDiv(nPixels(), blkSize.x));
+        CU_INVOKE(cu_indexedCopy)(*o_imgData, *colors, *_subsetInverseIndices, nPixels());
+    }
+}
+
+void Camera::_genLocalRays(bool norm) {
+    _localRays = sptr<CudaArray<glm::vec3>>(new CudaArray<glm::vec3>(nPixels()));
+    dim3 blkSize(32, 32);
+    dim3 grdSize(ceilDiv(_res.x, blkSize.x), ceilDiv(_res.y, blkSize.y));
+    if (norm)
+        CU_INVOKE(cu_genLocalRaysNormed)(*_localRays, _f, _c, _res);
+    else
+        CU_INVOKE(cu_genLocalRays)(*_localRays, _f, _c, _res);
+}

cpp/fnr_core/View.h

+#pragma once
+#include "../utils/common.h"
+
+
+class View {
+public:
+    View(glm::vec3 t, glm::mat3 r) : _t(t), _r(r) {}
+
+    glm::vec3 t() const { return _t; }
+    glm::mat3 r() const { return _r; }
+
+    void transPoints(sptr<CudaArray<glm::vec3>> results, sptr<CudaArray<glm::vec3>> points,
+                     sptr<CudaArray<int>> indices = nullptr, bool inverse = false);
+
+    void transVectors(sptr<CudaArray<glm::vec3>> results, sptr<CudaArray<glm::vec3>> vectors,
+                      sptr<CudaArray<int>> indices = nullptr, bool inverse = false);
+
+    View getStereoEye(float ipd, Eye eye);
+
+private:
+    glm::vec3 _t;
+    glm::mat3 _r;
+};
+
+class Camera {
+public:
+    Camera(glm::uvec2 res, float fov);
+    Camera(glm::uvec2 res, float fov, glm::vec2 c);
+
+    bool loadMaskData(std::string filepath);
+    sptr<CudaArray<glm::vec3>> localRays();
+    void getRays(sptr<CudaArray<glm::vec3>> o_rays, View &view, glm::vec3 offset);
+    void restoreImage(sptr<CudaArray<glm::vec4>> o_imgData, sptr<CudaArray<glm::vec4>> colors);
+    glm::vec2 f() const { return _f; }
+    glm::vec2 c() const { return _c; }
+    glm::uvec2 res() const { return _res; }
+    uint nPixels() const { return _res.x * _res.y; }
+    uint nRays() const { return _subsetIndices == nullptr ? nPixels() : (uint)_subsetIndices->n(); }
+    sptr<CudaArray<int>> subsetIndices() const { return _subsetIndices; };
+    sptr<CudaArray<int>> subsetInverseIndices() const { return _subsetInverseIndices; }
+
+private:
+    glm::vec2 _f;
+    glm::vec2 _c;
+    glm::uvec2 _res;
+    sptr<CudaArray<glm::vec3>> _localRays;
+    sptr<CudaArray<int>> _subsetIndices;
+    sptr<CudaArray<int>> _subsetInverseIndices;
+
+    void _genLocalRays(bool norm);
+};

cpp/fnr_core/Voxels.cu

+#include "Voxels.h"
+
+Voxels::Voxels(const std::string &path) { }
+
+void Voxels::toVoxelIndices(sptr<CudaArray<int>> gi, sptr<CudaArray<int>> o_vi) const {
+
+}
\ No newline at end of file

cpp/fnr_core/Voxels.h

+#pragma once
+#include "../utils/common.h"
+
+class Voxels {
+public:
+    Voxels(const std::string &path);
+
+    int nVoxels() const { return _voxels->n(); }
+
+    int nCorners() const { return _corners->n(); }
+
+    int nGrids() const { return _steps[0] * _steps[1] * _steps[2]; }
+
+    glm::vec3 voxelSize() const { return (_bbox[1] - _bbox[0]) / (glm::vec3)_steps; }
+
+    void toVoxelIndices(sptr<CudaArray<int>> gi, sptr<CudaArray<int>> o_vi) const;
+
+private:
+    glm::vec3 _bbox[2];
+    glm::ivec3 _steps;
+    sptr<CudaArray<glm::vec3>> _voxels;
+    sptr<CudaArray<int>> _cornerIndices;
+    sptr<CudaArray<glm::vec3>> _corners;
+    sptr<CudaArray<int>> _voxelIndicesInGrid;
+};
\ No newline at end of file