First-class spherical undistort support

src/core/camera.cpp

@@ -145,10 +145,13 @@ namespace lfs::core {
           _T(std::move(other._T)),
           _radial_distortion(std::move(other._radial_distortion)),
           _tangential_distortion(std::move(other._tangential_distortion)),
+          _camera_model_type(other._camera_model_type),
           _image_path(std::move(other._image_path)),
           _image_name(std::move(other._image_name)),
           _mask_path(std::move(other._mask_path)),
           _depth_path(std::move(other._depth_path)),
+          _cube_face_projection(std::move(other._cube_face_projection)),
+          _has_alpha(other._has_alpha),
           _split(other._split),
           _camera_width(other._camera_width),
           _camera_height(other._camera_height),
@@ -194,10 +197,13 @@ namespace lfs::core {
             _T = std::move(other._T);
             _radial_distortion = std::move(other._radial_distortion);
             _tangential_distortion = std::move(other._tangential_distortion);
+            _camera_model_type = other._camera_model_type;
             _image_path = std::move(other._image_path);
             _image_name = std::move(other._image_name);
             _mask_path = std::move(other._mask_path);
             _depth_path = std::move(other._depth_path);
+            _cube_face_projection = std::move(other._cube_face_projection);
+            _has_alpha = other._has_alpha;
             _split = other._split;
             _camera_width = other._camera_width;
             _camera_height = other._camera_height;
@@ -241,6 +247,8 @@ namespace lfs::core {
           _image_path(other._image_path),
           _mask_path(other._mask_path),
           _depth_path(other._depth_path),
+          _cube_face_projection(other._cube_face_projection),
+          _has_alpha(other._has_alpha),
           _split(other._split),
           _camera_width(other._camera_width),
           _camera_height(other._camera_height),
@@ -281,6 +289,10 @@ namespace lfs::core {
         return {_focal_x * x_scale, _focal_y * y_scale, _center_x * x_scale, _center_y * y_scale};
     }
 
+    void Camera::set_cube_face_projection(CubeFaceProjection projection) {
+        _cube_face_projection = std::move(projection);
+    }
+
     Tensor Camera::load_and_get_image(int resize_factor, int max_width, const bool output_uint8,
                                       const bool update_dimensions) {
         const ImageLoadParams params{
@@ -310,7 +322,10 @@ namespace lfs::core {
 
     void Camera::load_image_size(int resize_factor, int max_width) {
         int w, h;
-        if (_undistort_prepared) {
+        if (_cube_face_projection) {
+            w = _cube_face_projection->face_size;
+            h = _cube_face_projection->face_size;
+        } else if (_undistort_prepared) {
             w = _camera_width;
             h = _camera_height;
         } else {
@@ -356,11 +371,11 @@ namespace lfs::core {
     }
 
     size_t Camera::get_num_bytes_from_file(int resize_factor, int max_width) const {
-        auto result = get_image_info(_image_path);
-
-        int w = std::get<0>(result);
-        int h = std::get<1>(result);
-        int c = std::get<2>(result);
+        auto [w, h, c] = get_image_info(_image_path);
+        if (_cube_face_projection) {
+            w = _cube_face_projection->face_size;
+            h = _cube_face_projection->face_size;
+        }
 
         if (resize_factor > 0) {
             w = w / resize_factor;
@@ -383,6 +398,10 @@ namespace lfs::core {
 
     size_t Camera::get_num_bytes_from_file() const {
         auto [w, h, c] = get_image_info(_image_path);
+        if (_cube_face_projection) {
+            w = _cube_face_projection->face_size;
+            h = _cube_face_projection->face_size;
+        }
         size_t num_bytes = w * h * c * sizeof(uint8_t);
         return num_bytes;
     }

src/core/include/core/camera.hpp

@@ -5,6 +5,7 @@
 #pragma once
 
 #include "core/camera_types.h"
+#include "core/cube_face_projection.hpp"
 #include "core/cuda/undistort/undistort.hpp"
 #include "core/export.hpp"
 #include "core/tensor.hpp"
@@ -13,6 +14,7 @@
 #include <cuda_runtime.h>
 #include <filesystem>
 #include <future>
+#include <optional>
 #include <string>
 
 namespace lfs::core {
@@ -119,6 +121,9 @@ namespace lfs::core {
         const std::filesystem::path& image_path() const noexcept { return _image_path; }
         const std::filesystem::path& mask_path() const noexcept { return _mask_path; }
         const std::filesystem::path& depth_path() const noexcept { return _depth_path; }
+        void set_cube_face_projection(CubeFaceProjection projection);
+        const std::optional<CubeFaceProjection>& cube_face_projection() const noexcept { return _cube_face_projection; }
+        bool has_cube_face_projection() const noexcept { return _cube_face_projection.has_value(); }
         bool has_in_memory_mask() const noexcept { return _in_memory_mask_raw.is_valid(); }
         bool has_mask() const noexcept {
             return has_in_memory_mask() ||
@@ -174,6 +179,7 @@ namespace lfs::core {
         std::filesystem::path _image_path;
         std::filesystem::path _mask_path;
         std::filesystem::path _depth_path;
+        std::optional<CubeFaceProjection> _cube_face_projection;
         bool _has_alpha = false;
         CameraSplit _split = CameraSplit::Train;
         int _camera_width = 0;

src/core/include/core/cube_face_projection.hpp

@@ -0,0 +1,19 @@
+/* SPDX-FileCopyrightText: 2025 LichtFeld Studio Authors
+ *
+ * SPDX-License-Identifier: GPL-3.0-or-later */
+
+#pragma once
+
+#include <array>
+
+namespace lfs::core {
+
+    struct CubeFaceProjection {
+        std::array<float, 9> pano_to_face{};
+        int face_size = 0;
+        int source_width = 0;
+        int source_height = 0;
+        float fov_degrees = 90.0f;
+    };
+
+} // namespace lfs::core

src/io/CMakeLists.txt

@@ -213,6 +213,8 @@ add_library(lfs_io STATIC
         mesh/texture_loader.hpp
         mesh/texture_loader.cpp
         cache_image_loader.cpp
+        cube_face_projection.cpp
+        include/io/cube_face_projection.hpp
         nvcodec_image_loader.hpp
         nvcodec_image_loader.cpp
         pipelined_image_loader.cpp

src/io/cube_face_projection.cpp

@@ -0,0 +1,114 @@
+/* SPDX-FileCopyrightText: 2025 LichtFeld Studio Authors
+ *
+ * SPDX-License-Identifier: GPL-3.0-or-later */
+
+#include "io/cube_face_projection.hpp"
+
+#include <algorithm>
+#include <cmath>
+#include <stdexcept>
+
+namespace lfs::io {
+
+    namespace {
+        constexpr float kPi = 3.14159265358979323846f;
+
+        [[nodiscard]] float wrap_source_x(float x, const int width) {
+            const float w = static_cast<float>(width);
+            x = std::fmod(x, w);
+            if (x < 0.0f)
+                x += w;
+            return x;
+        }
+
+        [[nodiscard]] float projection_focal(const lfs::core::CubeFaceProjection& projection,
+                                             const int output_size) {
+            const float fov = std::clamp(projection.fov_degrees, 1.0f, 179.0f);
+            return 0.5f * static_cast<float>(output_size) /
+                   std::tan(0.5f * fov * kPi / 180.0f);
+        }
+    } // namespace
+
+    int projected_face_output_size(
+        const lfs::core::CubeFaceProjection& projection,
+        const int resize_factor,
+        const int max_width) {
+        int size = std::max(1, projection.face_size);
+        if (resize_factor > 1) {
+            size = std::max(1, size / resize_factor);
+        }
+        if (max_width > 0 && size > max_width) {
+            size = max_width;
+        }
+        return size;
+    }
+
+    std::vector<uint8_t> project_cube_face_hwc(
+        const uint8_t* source,
+        const int width,
+        const int height,
+        const int channels,
+        const lfs::core::CubeFaceProjection& projection,
+        const int output_size) {
+        if (!source || width <= 0 || height <= 0 || channels <= 0 || output_size <= 0) {
+            throw std::runtime_error("Invalid spherical projection input");
+        }
+
+        std::vector<uint8_t> projected(static_cast<size_t>(output_size) * output_size * channels);
+
+        const float focal = projection_focal(projection, output_size);
+        const float inv_focal = 1.0f / focal;
+        const float center = 0.5f * static_cast<float>(output_size);
+        const auto& m = projection.pano_to_face;
+
+        for (int y = 0; y < output_size; ++y) {
+            const float face_y = (static_cast<float>(y) + 0.5f - center) * inv_focal;
+            for (int x = 0; x < output_size; ++x) {
+                const float face_x = (static_cast<float>(x) + 0.5f - center) * inv_focal;
+                constexpr float face_z = 1.0f;
+
+                float pano_x = m[0] * face_x + m[3] * face_y + m[6] * face_z;
+                float pano_y = m[1] * face_x + m[4] * face_y + m[7] * face_z;
+                float pano_z = m[2] * face_x + m[5] * face_y + m[8] * face_z;
+
+                const float inv_len = 1.0f / std::sqrt(pano_x * pano_x + pano_y * pano_y + pano_z * pano_z);
+                pano_x *= inv_len;
+                pano_y *= inv_len;
+                pano_z *= inv_len;
+
+                const float azimuth = std::atan2(pano_x, pano_z);
+                const float elevation = std::asin(std::clamp(pano_y, -1.0f, 1.0f));
+                const float src_x = wrap_source_x((azimuth / (2.0f * kPi) + 0.5f) * width - 0.5f, width);
+                const float src_y = std::clamp((elevation / kPi + 0.5f) * height - 0.5f,
+                                               0.0f,
+                                               static_cast<float>(height - 1));
+
+                const int x0 = static_cast<int>(std::floor(src_x));
+                const int y0 = static_cast<int>(std::floor(src_y));
+                const int x1 = (x0 + 1) % width;
+                const int y1 = std::min(y0 + 1, height - 1);
+                const float tx = src_x - static_cast<float>(x0);
+                const float ty = src_y - static_cast<float>(y0);
+
+                const size_t dst_base = (static_cast<size_t>(y) * output_size + x) * channels;
+                const size_t p00 = (static_cast<size_t>(y0) * width + x0) * channels;
+                const size_t p10 = (static_cast<size_t>(y0) * width + x1) * channels;
+                const size_t p01 = (static_cast<size_t>(y1) * width + x0) * channels;
+                const size_t p11 = (static_cast<size_t>(y1) * width + x1) * channels;
+
+                for (int c = 0; c < channels; ++c) {
+                    const float top = static_cast<float>(source[p00 + c]) * (1.0f - tx) +
+                                      static_cast<float>(source[p10 + c]) * tx;
+                    const float bottom = static_cast<float>(source[p01 + c]) * (1.0f - tx) +
+                                         static_cast<float>(source[p11 + c]) * tx;
+                    const float value = top * (1.0f - ty) + bottom * ty;
+                    projected[dst_base + c] = static_cast<uint8_t>(
+                        std::clamp(static_cast<int>(std::lround(value)), 0, 255));
+                }
+            }
+        }
+
+        return projected;
+    }
+
+} // namespace lfs::io

src/io/include/io/cube_face_projection.hpp

@@ -0,0 +1,28 @@
+/* SPDX-FileCopyrightText: 2025 LichtFeld Studio Authors
+ *
+ * SPDX-License-Identifier: GPL-3.0-or-later */
+
+#pragma once
+
+#include "core/cube_face_projection.hpp"
+#include "core/export.hpp"
+
+#include <cstdint>
+#include <vector>
+
+namespace lfs::io {
+
+    [[nodiscard]] LFS_IO_API int projected_face_output_size(
+        const lfs::core::CubeFaceProjection& projection,
+        int resize_factor,
+        int max_width);
+
+    [[nodiscard]] LFS_IO_API std::vector<uint8_t> project_cube_face_hwc(
+        const uint8_t* source,
+        int width,
+        int height,
+        int channels,
+        const lfs::core::CubeFaceProjection& projection,
+        int output_size);
+
+} // namespace lfs::io

src/io/include/io/pipelined_image_loader.hpp

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include "core/cube_face_projection.hpp"
 #include "core/export.hpp"
 #include "core/tensor.hpp"
 #include "io/cache_image_loader.hpp"
@@ -23,6 +24,10 @@
 struct CUstream_st;
 using cudaStream_t = CUstream_st*;
 
+namespace lfs::core {
+    class Camera;
+}
+
 namespace lfs::io {
 
     namespace config {
@@ -73,6 +78,7 @@ namespace lfs::io {
         bool extract_alpha_as_mask = false;
         MaskParams alpha_mask_params;
         const lfs::core::UndistortParams* undistort = nullptr;
+        std::optional<lfs::core::CubeFaceProjection> cube_face_projection;
     };
 
     struct ReadyImage {
@@ -147,6 +153,17 @@ namespace lfs::io {
 
         lfs::core::Tensor load_image_immediate(
             const std::filesystem::path& path, const LoadParams& params);
+        lfs::core::Tensor load_camera_image_immediate(
+            const lfs::core::Camera& camera,
+            const LoadParams& params);
+        lfs::core::Tensor load_camera_image_immediate(
+            const lfs::core::Camera& camera,
+            const std::filesystem::path& path,
+            const LoadParams& params);
+        lfs::core::Tensor load_image_immediate(
+            const std::filesystem::path& path,
+            const LoadParams& params,
+            const std::optional<lfs::core::CubeFaceProjection>& cube_face_projection);
 
         size_t ready_count() const;
         size_t in_flight_count() const;
@@ -174,6 +191,7 @@ namespace lfs::io {
             bool alpha_as_mask = false;
             MaskParams alpha_mask_params;
             const lfs::core::UndistortParams* undistort = nullptr;
+            std::optional<lfs::core::CubeFaceProjection> cube_face_projection;
         };
 
         struct CachedJpegHit {
@@ -285,7 +303,10 @@ namespace lfs::io {
         void gpu_batch_decode_thread_func();
         void cold_process_thread_func();
 
-        std::string make_cache_key(const std::filesystem::path& path, const LoadParams& params) const;
+        std::string make_cache_key(
+            const std::filesystem::path& path,
+            const LoadParams& params,
+            const std::optional<lfs::core::CubeFaceProjection>& cube_face_projection = std::nullopt) const;
         std::filesystem::path get_fs_cache_path(const std::string& cache_key) const;
         bool is_jpeg_data(const std::vector<uint8_t>& data) const;
         std::vector<uint8_t> read_file(const std::filesystem::path& path) const;
@@ -302,7 +323,8 @@ namespace lfs::io {
         // Auxiliary image pairing helpers
         std::string make_mask_cache_key(
             const std::filesystem::path& path,
-            const LoadParams& params) const;
+            const LoadParams& params,
+            const std::optional<lfs::core::CubeFaceProjection>& cube_face_projection = std::nullopt) const;
         void try_complete_pair(
             size_t sequence_id,
             std::optional<lfs::core::Tensor> image,