diff --git a/src/core/camera.cpp b/src/core/camera.cpp index 1083b0c65..60027b73a 100644 --- a/src/core/camera.cpp +++ b/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; } diff --git a/src/core/include/core/camera.hpp b/src/core/include/core/camera.hpp index 10853e4ff..6c405bd72 100644 --- a/src/core/include/core/camera.hpp +++ b/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 #include #include +#include #include 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& 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 _cube_face_projection; bool _has_alpha = false; CameraSplit _split = CameraSplit::Train; int _camera_width = 0; diff --git a/src/core/include/core/cube_face_projection.hpp b/src/core/include/core/cube_face_projection.hpp new file mode 100644 index 000000000..e13d0c6e2 --- /dev/null +++ b/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 + +namespace lfs::core { + + struct CubeFaceProjection { + std::array pano_to_face{}; + int face_size = 0; + int source_width = 0; + int source_height = 0; + float fov_degrees = 90.0f; + }; + +} // namespace lfs::core diff --git a/src/io/CMakeLists.txt b/src/io/CMakeLists.txt index 27964ca05..9f13e61b7 100644 --- a/src/io/CMakeLists.txt +++ b/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 diff --git a/src/io/cube_face_projection.cpp b/src/io/cube_face_projection.cpp new file mode 100644 index 000000000..4327d13fa --- /dev/null +++ b/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 +#include +#include + +namespace lfs::io { + + namespace { + constexpr float kPi = 3.14159265358979323846f; + + [[nodiscard]] float wrap_source_x(float x, const int width) { + const float w = static_cast(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(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 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 projected(static_cast(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(output_size); + const auto& m = projection.pano_to_face; + + for (int y = 0; y < output_size; ++y) { + const float face_y = (static_cast(y) + 0.5f - center) * inv_focal; + for (int x = 0; x < output_size; ++x) { + const float face_x = (static_cast(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(height - 1)); + + const int x0 = static_cast(std::floor(src_x)); + const int y0 = static_cast(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(x0); + const float ty = src_y - static_cast(y0); + + const size_t dst_base = (static_cast(y) * output_size + x) * channels; + const size_t p00 = (static_cast(y0) * width + x0) * channels; + const size_t p10 = (static_cast(y0) * width + x1) * channels; + const size_t p01 = (static_cast(y1) * width + x0) * channels; + const size_t p11 = (static_cast(y1) * width + x1) * channels; + + for (int c = 0; c < channels; ++c) { + const float top = static_cast(source[p00 + c]) * (1.0f - tx) + + static_cast(source[p10 + c]) * tx; + const float bottom = static_cast(source[p01 + c]) * (1.0f - tx) + + static_cast(source[p11 + c]) * tx; + const float value = top * (1.0f - ty) + bottom * ty; + projected[dst_base + c] = static_cast( + std::clamp(static_cast(std::lround(value)), 0, 255)); + } + } + } + + return projected; + } + +} // namespace lfs::io diff --git a/src/io/include/io/cube_face_projection.hpp b/src/io/include/io/cube_face_projection.hpp new file mode 100644 index 000000000..d8b5dbd18 --- /dev/null +++ b/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 +#include + +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 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 diff --git a/src/io/include/io/pipelined_image_loader.hpp b/src/io/include/io/pipelined_image_loader.hpp index b565886a8..1627a4527 100644 --- a/src/io/include/io/pipelined_image_loader.hpp +++ b/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 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& 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 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& 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& data) const; std::vector 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& cube_face_projection = std::nullopt) const; void try_complete_pair( size_t sequence_id, std::optional image, diff --git a/src/io/pipelined_image_loader.cpp b/src/io/pipelined_image_loader.cpp index b6fc0edf9..3ecd597ce 100644 --- a/src/io/pipelined_image_loader.cpp +++ b/src/io/pipelined_image_loader.cpp @@ -2,6 +2,7 @@ * SPDX-License-Identifier: GPL-3.0-or-later */ #include "io/pipelined_image_loader.hpp" +#include "core/camera.hpp" #include "core/cuda/lanczos_resize/lanczos_resize.hpp" #include "core/cuda/undistort/undistort.hpp" #include "core/image_io.hpp" @@ -11,14 +12,17 @@ #include "core/tensor/internal/memory_pool.hpp" #include "cuda/image_format_kernels.cuh" #include "diagnostics/vram_profiler.hpp" +#include "io/cube_face_projection.hpp" #include "io/nvcodec_image_loader.hpp" #include #include #include +#include #include #include +#include #include namespace lfs::io { @@ -173,6 +177,22 @@ namespace lfs::io { return lfs::core::path_to_utf8(path); } + void append_cube_face_projection_key( + std::string& key, + const std::optional& cube_face_projection) { + if (!cube_face_projection) + return; + + key += "_cube" + std::to_string(cube_face_projection->face_size) + + "_src" + std::to_string(cube_face_projection->source_width) + + "x" + std::to_string(cube_face_projection->source_height) + + "_fov" + std::to_string(static_cast( + std::lround(cube_face_projection->fov_degrees * 100.0f))); + for (const float value : cube_face_projection->pano_to_face) { + key += "_" + std::to_string(static_cast(std::lround(value * 1000.0f))); + } + } + void apply_requested_undistort(lfs::core::Tensor& tensor, const LoadParams& params) { if (!params.undistort) return; @@ -275,6 +295,97 @@ namespace lfs::io { return {data, w, h}; } + struct FreeImageDeleter { + void operator()(uint8_t* data) const noexcept { + if (data) + lfs::core::free_image(data); + } + }; + + struct StbiImageDeleter { + void operator()(uint8_t* data) const noexcept { + if (data) + stbi_image_free(data); + } + }; + + using FreeImagePtr = std::unique_ptr; + using StbiImagePtr = std::unique_ptr; + + [[nodiscard]] lfs::core::Tensor upload_hwc_uint8_to_chw( + const uint8_t* data, + const size_t height, + const size_t width, + const size_t channels, + const bool output_uint8, + cudaStream_t stream) { + auto cpu_tensor = lfs::core::Tensor::from_blob( + const_cast(data), + lfs::core::TensorShape({height, width, channels}), + lfs::core::Device::CPU, + lfs::core::DataType::UInt8); + auto gpu_uint8 = cpu_tensor.to(lfs::core::Device::CUDA); + + if (output_uint8) { + auto decoded = lfs::core::Tensor::empty( + lfs::core::TensorShape({channels, height, width}), + lfs::core::Device::CUDA, + lfs::core::DataType::UInt8); + cuda::launch_uint8_hwc_to_uint8_chw( + gpu_uint8.ptr(), + decoded.ptr(), + height, + width, + channels, + stream); + return decoded; + } + + auto decoded = lfs::core::Tensor::zeros( + lfs::core::TensorShape({channels, height, width}), + lfs::core::Device::CUDA, + lfs::core::DataType::Float32); + cuda::launch_uint8_hwc_to_float32_chw( + gpu_uint8.ptr(), + decoded.ptr(), + height, + width, + channels, + stream); + return decoded; + } + + [[nodiscard]] lfs::core::Tensor load_projected_cube_face_tensor( + const std::filesystem::path& path, + const LoadParams& params, + const lfs::core::CubeFaceProjection& projection) { + auto [img_data, width, height, channels] = lfs::core::load_image(path, 1, 0); + if (!img_data) { + throw std::runtime_error("Failed to decode image: " + lfs::core::path_to_utf8(path)); + } + const FreeImagePtr image_guard(img_data); + + const int output_size = projected_face_output_size( + projection, + params.resize_factor, + params.max_width); + auto projected = project_cube_face_hwc( + img_data, + width, + height, + channels, + projection, + output_size); + + return upload_hwc_uint8_to_chw( + projected.data(), + static_cast(output_size), + static_cast(output_size), + static_cast(channels), + params.output_uint8, + static_cast(params.cuda_stream)); + } + } // namespace PipelinedImageLoader::PipelinedImageLoader(PipelinedLoaderConfig config) @@ -470,7 +581,27 @@ namespace lfs::io { lfs::core::Tensor PipelinedImageLoader::load_image_immediate( const std::filesystem::path& path, const LoadParams& params) { - const auto cache_key = make_cache_key(path, params); + return load_image_immediate(path, params, std::nullopt); + } + + lfs::core::Tensor PipelinedImageLoader::load_camera_image_immediate( + const lfs::core::Camera& camera, + const LoadParams& params) { + return load_image_immediate(camera.image_path(), params, camera.cube_face_projection()); + } + + lfs::core::Tensor PipelinedImageLoader::load_camera_image_immediate( + const lfs::core::Camera& camera, + const std::filesystem::path& path, + const LoadParams& params) { + return load_image_immediate(path, params, camera.cube_face_projection()); + } + + lfs::core::Tensor PipelinedImageLoader::load_image_immediate( + const std::filesystem::path& path, + const LoadParams& params, + const std::optional& cube_face_projection) { + const auto cache_key = make_cache_key(path, params, cube_face_projection); const auto base_key = make_base_cache_key(path); auto decode_cached_hit = [&](const std::shared_ptr>& jpeg_data, const bool cached_blob_is_base) -> lfs::core::Tensor { @@ -493,6 +624,26 @@ namespace lfs::io { } } + if (cube_face_projection) { + auto decoded = load_projected_cube_face_tensor(path, params, *cube_face_projection); + apply_requested_undistort(decoded, params); + + if (is_nvcodec_available()) { + try { + auto nvcodec = acquire_nvcodec_loader(config_.decoder_pool_size); + auto jpeg_bytes = nvcodec->encode_to_jpeg(decoded, config_.cache_jpeg_quality, params.cuda_stream); + save_to_fs_cache(cache_key, jpeg_bytes); + put_in_jpeg_cache(cache_key, std::make_shared>(std::move(jpeg_bytes))); + } catch (...) { + LOG_DEBUG("[PipelinedImageLoader] Immediate spherical projection cache write skipped for {}: {}", + lfs::core::path_to_utf8(path), + describe_current_exception("non-standard nvImageCodec exception")); + } + } + + return decoded; + } + if (auto jpeg_data = load_cached_jpeg_blob(base_key)) { if (auto tensor = decode_cached_hit(jpeg_data, true); tensor.is_valid() && tensor.numel() > 0) { @@ -632,21 +783,27 @@ namespace lfs::io { return decoded; } - std::string PipelinedImageLoader::make_cache_key(const std::filesystem::path& path, const LoadParams& params) const { + std::string PipelinedImageLoader::make_cache_key( + const std::filesystem::path& path, + const LoadParams& params, + const std::optional& cube_face_projection) const { auto key = lfs::core::path_to_utf8(path) + ":rf" + std::to_string(params.resize_factor) + "_mw" + std::to_string(params.max_width); if (params.undistort) key += "_ud"; + append_cube_face_projection_key(key, cube_face_projection); return key; } std::string PipelinedImageLoader::make_mask_cache_key( const std::filesystem::path& path, - const LoadParams& params) const { + const LoadParams& params, + const std::optional& cube_face_projection) const { auto key = lfs::core::path_to_utf8(path) + ":mask_rf" + std::to_string(params.resize_factor) + "_mw" + std::to_string(params.max_width); if (params.undistort) key += "_ud"; + append_cube_face_projection_key(key, cube_face_projection); return key; } @@ -984,8 +1141,8 @@ namespace lfs::io { } if (request.extract_alpha_as_mask) { - const auto rgb_key = make_cache_key(request.path, request.params); - const auto alpha_key = make_mask_cache_key(request.path, request.params); + const auto rgb_key = make_cache_key(request.path, request.params, request.cube_face_projection); + const auto alpha_key = make_mask_cache_key(request.path, request.params, request.cube_face_projection); auto cached_rgb = get_from_jpeg_cache(rgb_key); auto cached_alpha = get_from_jpeg_cache(alpha_key); @@ -996,6 +1153,7 @@ namespace lfs::io { img_item.cache_key = rgb_key; img_item.jpeg_data = cached_rgb; img_item.is_cache_hit = true; + img_item.cube_face_projection = request.cube_face_projection; hot_queue_.push(std::move(img_item)); PrefetchedImage mask_item; @@ -1005,6 +1163,7 @@ namespace lfs::io { mask_item.jpeg_data = cached_alpha; mask_item.is_mask = true; mask_item.is_cache_hit = true; + mask_item.cube_face_projection = request.cube_face_projection; hot_queue_.push(std::move(mask_item)); std::lock_guard lock(stats_mutex_); @@ -1019,6 +1178,7 @@ namespace lfs::io { result.alpha_mask_params = request.alpha_mask_params; result.needs_processing = true; result.undistort = request.undistort; + result.cube_face_projection = request.cube_face_projection; cold_queue_.push(std::move(result)); std::lock_guard lock(stats_mutex_); @@ -1032,21 +1192,38 @@ namespace lfs::io { result.sequence_id = request.sequence_id; result.path = request.path; result.params = request.params; - result.cache_key = make_cache_key(request.path, request.params); + result.cache_key = make_cache_key(request.path, request.params, request.cube_face_projection); result.is_mask = false; result.undistort = request.undistort; + result.cube_face_projection = request.cube_face_projection; try { const bool needs_requested_processing = load_params_need_processing(request.params); + const bool needs_cube_projection = request.cube_face_projection.has_value(); const auto base_key = make_base_cache_key(request.path); - if (auto cached = find_cached_jpeg(result.cache_key, base_key)) { + std::optional cached = needs_cube_projection + ? std::optional{} + : find_cached_jpeg(result.cache_key, base_key); + if (needs_cube_projection) { + if (auto projected = load_cached_jpeg_blob(result.cache_key)) { + cached = CachedJpegHit{.data = std::move(projected), .from_base_key = false}; + } + } + + if (cached) { result.jpeg_data = std::move(cached->data); result.is_cache_hit = true; result.needs_processing = cached->from_base_key && needs_requested_processing; hot_queue_.push(std::move(result)); std::lock_guard lock(stats_mutex_); ++stats_.hot_path_hits; + } else if (needs_cube_projection) { + result.is_cache_hit = false; + result.needs_processing = true; + cold_queue_.push(std::move(result)); + std::lock_guard lock(stats_mutex_); + ++stats_.cold_path_misses; } else { result.raw_bytes = read_file(request.path); result.is_original_jpeg = is_jpeg_data(result.raw_bytes); @@ -1057,7 +1234,7 @@ namespace lfs::io { stats_.total_bytes_read += result.raw_bytes.size(); } - if (result.is_original_jpeg && !needs_requested_processing) { + if (result.is_original_jpeg && !needs_requested_processing && !needs_cube_projection) { auto data = std::make_shared>(std::move(result.raw_bytes)); put_in_jpeg_cache(result.cache_key, data); result.jpeg_data = data; @@ -1088,10 +1265,11 @@ namespace lfs::io { mask_result.sequence_id = request.sequence_id; mask_result.path = *request.mask_path; mask_result.params = request.params; - mask_result.cache_key = make_mask_cache_key(*request.mask_path, request.params); + mask_result.cache_key = make_mask_cache_key(*request.mask_path, request.params, request.cube_face_projection); mask_result.is_mask = true; mask_result.mask_params = request.mask_params; mask_result.undistort = request.undistort; + mask_result.cube_face_projection = request.cube_face_projection; try { if (auto cached = get_from_jpeg_cache(mask_result.cache_key)) { @@ -1100,6 +1278,11 @@ namespace lfs::io { hot_queue_.push(std::move(mask_result)); std::lock_guard lock(stats_mutex_); ++stats_.mask_cache_hits; + } else if (request.cube_face_projection) { + mask_result.needs_processing = true; + cold_queue_.push(std::move(mask_result)); + std::lock_guard lock(stats_mutex_); + ++stats_.mask_cache_misses; } else if (config_.use_filesystem_cache) { const auto fs_path = get_fs_cache_path(mask_result.cache_key); auto done_path = fs_path; @@ -1333,19 +1516,40 @@ namespace lfs::io { if (item.alpha_as_mask) { auto [img_data, width, height, channels] = lfs::core::load_image_with_alpha( - item.path, item.params.resize_factor, item.params.max_width); + item.path, + item.cube_face_projection ? 1 : item.params.resize_factor, + item.cube_face_projection ? 0 : item.params.max_width); + const FreeImagePtr image_guard(img_data); if (!img_data || channels != 4) throw std::runtime_error("Failed to load RGBA image"); + std::vector projected_rgba; + const uint8_t* rgba_data = image_guard.get(); + if (item.cube_face_projection) { + const int output_size = projected_face_output_size( + *item.cube_face_projection, + item.params.resize_factor, + item.params.max_width); + projected_rgba = project_cube_face_hwc( + image_guard.get(), + width, + height, + channels, + *item.cube_face_projection, + output_size); + width = output_size; + height = output_size; + rgba_data = projected_rgba.data(); + } + const size_t H = static_cast(height); const size_t W = static_cast(width); auto cpu_tensor = lfs::core::Tensor::from_blob( - img_data, lfs::core::TensorShape({H, W, 4}), + const_cast(rgba_data), lfs::core::TensorShape({H, W, 4}), lfs::core::Device::CPU, lfs::core::DataType::UInt8); auto gpu_uint8 = cpu_tensor.to(lfs::core::Device::CUDA); - lfs::core::free_image(img_data); auto rgb = item.params.output_uint8 ? lfs::core::Tensor::empty( @@ -1404,7 +1608,10 @@ namespace lfs::io { put_in_jpeg_cache(item.cache_key, std::make_shared>(std::move(rgb_jpeg))); - const auto alpha_key = make_mask_cache_key(item.path, item.params); + const auto alpha_key = make_mask_cache_key( + item.path, + item.params, + item.cube_face_projection); auto alpha_jpeg = nvcodec->encode_grayscale_to_jpeg( alpha, config_.cache_jpeg_quality, nullptr); save_to_fs_cache(alpha_key, alpha_jpeg); @@ -1425,7 +1632,7 @@ namespace lfs::io { lfs::core::Tensor aux_tensor; bool used_gpu = false; - if (is_nvcodec_available()) { + if (is_nvcodec_available() && !item.cube_face_projection) { try { aux_tensor = nvcodec->load_image_gpu( item.path, item.params.resize_factor, item.params.max_width, @@ -1437,16 +1644,35 @@ namespace lfs::io { if (!used_gpu) { const auto [gray_data, src_w, src_h] = load_grayscale_stb(item.path); + const StbiImagePtr gray_guard(gray_data); if (!gray_data) throw std::runtime_error(item.is_mask ? "Failed to decode mask" : "Failed to decode depth"); int target_w = src_w; int target_h = src_h; - if (item.params.resize_factor > 1) { + std::vector projected_gray; + const uint8_t* gray_source = gray_guard.get(); + + if (item.cube_face_projection) { + const int output_size = projected_face_output_size( + *item.cube_face_projection, + item.params.resize_factor, + item.params.max_width); + projected_gray = project_cube_face_hwc( + gray_guard.get(), + src_w, + src_h, + 1, + *item.cube_face_projection, + output_size); + gray_source = projected_gray.data(); + target_w = output_size; + target_h = output_size; + } else if (item.params.resize_factor > 1) { target_w /= item.params.resize_factor; target_h /= item.params.resize_factor; } - const int max_w = item.params.max_width; + const int max_w = item.cube_face_projection ? 0 : item.params.max_width; if (max_w > 0 && (target_w > max_w || target_h > max_w)) { if (target_w > target_h) { target_h = std::max(1, max_w * target_h / target_w); @@ -1458,12 +1684,13 @@ namespace lfs::io { } const auto cpu_tensor = lfs::core::Tensor::from_blob( - gray_data, lfs::core::TensorShape({static_cast(src_h), static_cast(src_w)}), + const_cast(gray_source), + lfs::core::TensorShape({static_cast(item.cube_face_projection ? target_h : src_h), + static_cast(item.cube_face_projection ? target_w : src_w)}), lfs::core::Device::CPU, lfs::core::DataType::UInt8); const auto gpu_uint8 = cpu_tensor.to(lfs::core::Device::CUDA); - stbi_image_free(gray_data); - if (target_w != src_w || target_h != src_h) { + if (!item.cube_face_projection && (target_w != src_w || target_h != src_h)) { aux_tensor = lfs::core::lanczos_resize_grayscale(gpu_uint8, target_h, target_w, 2, nullptr); } else { aux_tensor = lfs::core::Tensor::zeros( @@ -1524,7 +1751,7 @@ namespace lfs::io { lfs::core::Tensor decoded; bool used_gpu = false; - if (is_nvcodec_available() && item.is_original_jpeg) { + if (is_nvcodec_available() && item.is_original_jpeg && !item.cube_face_projection) { try { decoded = nvcodec->load_image_gpu( item.path, item.params.resize_factor, item.params.max_width, @@ -1536,44 +1763,48 @@ namespace lfs::io { if (!used_gpu) { auto [img_data, width, height, channels] = lfs::core::load_image( - item.path, item.params.resize_factor, item.params.max_width); + item.path, + item.cube_face_projection ? 1 : item.params.resize_factor, + item.cube_face_projection ? 0 : item.params.max_width); if (!img_data) throw std::runtime_error("Failed to decode image"); + const FreeImagePtr image_guard(img_data); + + std::vector projected_rgb; + const uint8_t* image_data = image_guard.get(); + if (item.cube_face_projection) { + const int output_size = projected_face_output_size( + *item.cube_face_projection, + item.params.resize_factor, + item.params.max_width); + projected_rgb = project_cube_face_hwc( + image_guard.get(), + width, + height, + channels, + *item.cube_face_projection, + output_size); + width = output_size; + height = output_size; + image_data = projected_rgb.data(); + } const size_t H = static_cast(height); const size_t W = static_cast(width); const size_t C = static_cast(channels); - auto cpu_tensor = lfs::core::Tensor::from_blob( - img_data, lfs::core::TensorShape({H, W, C}), - lfs::core::Device::CPU, lfs::core::DataType::UInt8); - - auto gpu_uint8 = cpu_tensor.to(lfs::core::Device::CUDA); - lfs::core::free_image(img_data); - - if (item.params.output_uint8) { - decoded = lfs::core::Tensor::empty( - lfs::core::TensorShape({C, H, W}), - lfs::core::Device::CUDA, lfs::core::DataType::UInt8); - cuda::launch_uint8_hwc_to_uint8_chw( - reinterpret_cast(gpu_uint8.data_ptr()), - reinterpret_cast(decoded.data_ptr()), - H, W, C, nullptr); - } else { - decoded = lfs::core::Tensor::zeros( - lfs::core::TensorShape({C, H, W}), - lfs::core::Device::CUDA, lfs::core::DataType::Float32); - cuda::launch_uint8_hwc_to_float32_chw( - reinterpret_cast(gpu_uint8.data_ptr()), - reinterpret_cast(decoded.data_ptr()), - H, W, C, nullptr); - } + decoded = upload_hwc_uint8_to_chw( + image_data, + H, + W, + C, + item.params.output_uint8, + nullptr); if (const cudaError_t err = cudaDeviceSynchronize(); err != cudaSuccess) { throw std::runtime_error(std::string("CUDA sync failed: ") + cudaGetErrorString(err)); } - gpu_uint8 = lfs::core::Tensor(); } if (item.undistort) { @@ -1620,16 +1851,16 @@ namespace lfs::io { item.path, item.params.resize_factor, item.params.max_width); if (!img_data) throw std::runtime_error("RGB fallback also failed"); + const FreeImagePtr image_guard(img_data); const size_t H = static_cast(height); const size_t W = static_cast(width); const size_t C = static_cast(channels); auto cpu_tensor = lfs::core::Tensor::from_blob( - img_data, lfs::core::TensorShape({H, W, C}), + image_guard.get(), lfs::core::TensorShape({H, W, C}), lfs::core::Device::CPU, lfs::core::DataType::UInt8); auto gpu_uint8 = cpu_tensor.to(lfs::core::Device::CUDA); - lfs::core::free_image(img_data); auto decoded = item.params.output_uint8 ? lfs::core::Tensor::empty( diff --git a/src/python/lfs/py_scene.cpp b/src/python/lfs/py_scene.cpp index 5455d2568..e27e13300 100644 --- a/src/python/lfs/py_scene.cpp +++ b/src/python/lfs/py_scene.cpp @@ -684,6 +684,19 @@ namespace lfs::python { return result; } + size_t PyScene::active_training_image_count(const bool undistort) const { + size_t count = 0; + for (const auto* node : scene_->getNodes()) { + if (node->type != core::NodeType::CAMERA || !node->camera || !node->training_enabled) { + continue; + } + count += undistort && node->camera->camera_model_type() == core::CameraModelType::EQUIRECTANGULAR + ? 12 + : 1; + } + return count; + } + void PyScene::set_camera_training_enabled(const std::string& name, bool enabled) { if (auto* const scene_manager = get_scene_manager()) { const auto* node = scene_->getNode(name); @@ -1366,6 +1379,10 @@ namespace lfs::python { // Camera training control .def("set_camera_training_enabled", &PyScene::set_camera_training_enabled, nb::arg("name"), nb::arg("enabled"), "Enable or disable a camera for training by name") .def_prop_ro("active_camera_count", &PyScene::active_camera_count, "Number of cameras enabled for training") + .def("active_training_image_count", + &PyScene::active_training_image_count, + nb::arg("undistort"), + "Effective number of training images after internal image expansion") .def("get_active_cameras", &PyScene::get_active_cameras, "Get camera nodes enabled for training") // Training data .def("has_training_data", &PyScene::has_training_data, "Check if training dataset is loaded") diff --git a/src/python/lfs/py_scene.hpp b/src/python/lfs/py_scene.hpp index 7a39c22d7..157b25675 100644 --- a/src/python/lfs/py_scene.hpp +++ b/src/python/lfs/py_scene.hpp @@ -514,6 +514,7 @@ namespace lfs::python { // Camera training control void set_camera_training_enabled(const std::string& name, bool enabled); size_t active_camera_count() const { return scene_->getActiveCameraCount(); } + size_t active_training_image_count(bool undistort) const; std::vector get_active_cameras(); // Training data diff --git a/src/python/lfs_plugins/training_panel.py b/src/python/lfs_plugins/training_panel.py index 09897c2f8..fcd8822a7 100644 --- a/src/python/lfs_plugins/training_panel.py +++ b/src/python/lfs_plugins/training_panel.py @@ -2357,7 +2357,12 @@ def _try_auto_scale_steps(self, params): scene = lf.get_scene() if scene is None: return False - camera_count = scene.active_camera_count + try: + camera_count = scene.active_training_image_count( + bool(getattr(params, "undistort", False)) + ) + except (AttributeError, TypeError, ValueError, RuntimeError): + camera_count = scene.active_camera_count if camera_count == 0 or camera_count == self._auto_scaled_for_cameras: return False self._auto_scaled_for_cameras = camera_count diff --git a/src/python/stubs/lichtfeld/scene.pyi b/src/python/stubs/lichtfeld/scene.pyi index 49045de49..e3be1de72 100644 --- a/src/python/stubs/lichtfeld/scene.pyi +++ b/src/python/stubs/lichtfeld/scene.pyi @@ -663,6 +663,9 @@ class Scene: def active_camera_count(self) -> int: """Number of cameras enabled for training""" + def active_training_image_count(self, undistort: bool) -> int: + """Effective number of training images after internal image expansion""" + def get_active_cameras(self) -> list[SceneNode]: """Get camera nodes enabled for training""" diff --git a/src/training/dataset.hpp b/src/training/dataset.hpp index abe264fd3..c40dc4da9 100644 --- a/src/training/dataset.hpp +++ b/src/training/dataset.hpp @@ -551,10 +551,19 @@ namespace lfs::training { std::optional next() { if (shutdown_) return std::nullopt; - prefetch_next_batch(); try { - auto ready = loader_->get(); + std::optional ready_opt; + while (!ready_opt) { + prefetch_next_batch(); + if (sampler_exhausted_ && loader_->in_flight_count() == 0) { + sequence_to_camera_.clear(); + return std::nullopt; + } + ready_opt = loader_->try_get_for(config_.output_wait_timeout); + } + + auto ready = std::move(*ready_opt); const auto it = sequence_to_camera_.find(ready.sequence_id); if (it == sequence_to_camera_.end()) { LOG_ERROR("[PipelinedDataLoader] Unknown sequence_id: {}", ready.sequence_id); @@ -613,6 +622,7 @@ namespace lfs::training { sampler_.reset(); sequence_to_camera_.clear(); next_sequence_id_ = 0; + sampler_exhausted_ = false; prefetch_next_batch(); } @@ -630,10 +640,15 @@ namespace lfs::training { private: void prefetch_next_batch() { + if (sampler_exhausted_) + return; + while (loader_->in_flight_count() < config_.prefetch_count) { const auto indices = sampler_.next(1); - if (!indices || indices->empty()) + if (!indices || indices->empty()) { + sampler_exhausted_ = true; break; + } const size_t local_idx = (*indices)[0]; const size_t camera_idx = dataset_->local_to_source(local_idx); @@ -651,6 +666,7 @@ namespace lfs::training { request.undistort = &cam->undistort_params(); request.params.undistort = request.undistort; } + request.cube_face_projection = cam->cube_face_projection(); if (aux_config_.load_masks && cam->has_mask()) { request.mask_path = cam->mask_path(); @@ -677,6 +693,7 @@ namespace lfs::training { std::unordered_map sequence_to_camera_; size_t next_sequence_id_ = 0; + bool sampler_exhausted_ = false; bool shutdown_ = false; }; diff --git a/src/training/metrics/metrics.cpp b/src/training/metrics/metrics.cpp index 3b5295612..0e5446ada 100644 --- a/src/training/metrics/metrics.cpp +++ b/src/training/metrics/metrics.cpp @@ -5,13 +5,11 @@ #include "metrics.hpp" #include "../rasterization/fast_rasterizer.hpp" #include "../rasterization/gsplat_rasterizer.hpp" -#include "core/cuda/undistort/undistort.hpp" #include "core/events.hpp" #include "core/image_io.hpp" #include "core/logger.hpp" #include "core/path_utils.hpp" #include "core/splat_data.hpp" -#include "io/cuda/image_format_kernels.cuh" #include "lfs/kernels/ssim.cuh" #include #include @@ -352,91 +350,6 @@ namespace lfs::training { return colormap.to(depth_normalized.device()); } - lfs::core::Tensor MetricsEvaluator::load_eval_mask(lfs::core::Camera* cam, - lfs::core::Tensor& gt_image, - const bool alpha_as_mask) const { - if (cam->has_mask()) { - bool is_segment_and_ignore = _params.optimization.mask_mode == lfs::core::param::MaskMode::SegmentAndIgnore; - auto m = cam->load_and_get_mask( - _params.dataset.resize_factor, - _params.dataset.max_width, - _params.optimization.invert_masks, - _params.optimization.mask_threshold, - !is_segment_and_ignore); - if (is_segment_and_ignore) { - m = m.gt(250).to(lfs::core::DataType::UInt8).contiguous(); - } - return m; - } - - if (!alpha_as_mask) - return {}; - - // Re-load from disk because the dataloader strips alpha to produce RGB gt_image. - // We need the original alpha channel as the mask, with undistortion applied consistently. - auto [img_data, width, height, channels] = lfs::core::load_image_with_alpha( - cam->image_path(), _params.dataset.resize_factor, _params.dataset.max_width); - - if (!img_data || channels != 4) { - if (img_data) - lfs::core::free_image(img_data); - return {}; - } - - const auto H = static_cast(height); - const auto W = static_cast(width); - - auto cpu_tensor = lfs::core::Tensor::from_blob( - img_data, lfs::core::TensorShape({H, W, 4}), - lfs::core::Device::CPU, lfs::core::DataType::UInt8); - auto gpu_uint8 = cpu_tensor.to(lfs::core::Device::CUDA); - lfs::core::free_image(img_data); - - auto rgb = lfs::core::Tensor::zeros( - lfs::core::TensorShape({3, H, W}), - lfs::core::Device::CUDA, lfs::core::DataType::UInt8); - auto mask = lfs::core::Tensor::zeros( - lfs::core::TensorShape({H, W}), - lfs::core::Device::CUDA, lfs::core::DataType::Float32); - - lfs::io::cuda::launch_uint8_rgba_split_to_uint8_rgb_and_float32_alpha( - gpu_uint8.ptr(), rgb.ptr(), mask.ptr(), - H, W, nullptr); - gpu_uint8 = lfs::core::Tensor(); - - if (_params.optimization.invert_masks) - lfs::io::cuda::launch_mask_invert(mask.ptr(), H, W, nullptr); - if (_params.optimization.mask_threshold > 0) - lfs::io::cuda::launch_mask_threshold( - mask.ptr(), H, W, _params.optimization.mask_threshold, nullptr); - - if (cam->is_undistort_prepared()) { - const auto scaled = lfs::core::scale_undistort_params( - cam->undistort_params(), - static_cast(W), static_cast(H)); - auto rgb_float = rgb.to(lfs::core::DataType::Float32) / 255.0f; - rgb_float = lfs::core::undistort_image(rgb_float, scaled, nullptr); - auto rgb_uint8 = lfs::core::Tensor::empty( - rgb_float.shape(), lfs::core::Device::CUDA, lfs::core::DataType::UInt8); - lfs::io::cuda::launch_float32_chw_to_uint8_chw( - rgb_float.ptr(), - rgb_uint8.ptr(), - rgb_float.shape()[1], - rgb_float.shape()[2], - rgb_float.shape()[0], - nullptr); - rgb = std::move(rgb_uint8); - mask = lfs::core::undistort_mask(mask, scaled, nullptr); - } - - gt_image = std::move(rgb); - return mask.ge(0.5f).to(lfs::core::DataType::UInt8).contiguous(); - } - - auto MetricsEvaluator::make_dataloader(std::shared_ptr dataset, const int workers) const { - return create_dataloader_from_dataset(dataset, workers); - } - EvalMetrics MetricsEvaluator::evaluate(const int iteration, const lfs::core::SplatData& splatData, std::shared_ptr val_dataset, @@ -449,8 +362,6 @@ namespace lfs::training { result.num_gaussians = static_cast(splatData.size()); result.iteration = iteration; - const auto val_dataloader = make_dataloader(val_dataset); - std::vector psnr_values, ssim_values; const auto start_time = std::chrono::steady_clock::now(); @@ -472,9 +383,17 @@ namespace lfs::training { mask_mode == lfs::core::param::MaskMode::Ignore || mask_mode == lfs::core::param::MaskMode::SegmentAndIgnore; - while (auto batch_opt = val_dataloader->next()) { - auto& batch = *batch_opt; - auto camera_with_image = batch[0].data; + PipelinedAuxiliaryImageConfig aux_config; + aux_config.load_masks = use_masking; + aux_config.use_alpha_as_mask = use_masking && _params.optimization.use_alpha_as_mask; + aux_config.invert_masks = _params.optimization.invert_masks; + aux_config.mask_threshold = _params.optimization.mask_threshold; + + auto val_dataloader = create_pipelined_dataloader(val_dataset, {}, aux_config); + + while (auto example_opt = val_dataloader->next()) { + auto& example = *example_opt; + auto camera_with_image = std::move(example.data); lfs::core::Camera* cam = camera_with_image.camera; lfs::core::Tensor gt_image = std::move(camera_with_image.image); @@ -484,16 +403,9 @@ namespace lfs::training { lfs::core::Tensor mask; if (use_masking) { - const bool cam_alpha = _params.optimization.use_alpha_as_mask && cam->has_alpha(); - try { - mask = load_eval_mask(cam, gt_image, cam_alpha); - } catch (const std::exception& e) { - LOG_WARN("Eval: skipping camera '{}' (failed to load mask: {})", cam->image_name(), e.what()); - skipped_images++; - continue; - } - - if (!mask.is_valid()) { + if (example.mask && example.mask->is_valid()) { + mask = std::move(*example.mask); + } else { LOG_DEBUG("Eval: camera '{}' has no mask, proceeding unmasked", cam->image_name()); mask = lfs::core::Tensor(); } diff --git a/src/training/metrics/metrics.hpp b/src/training/metrics/metrics.hpp index 1054adb11..cba637fa2 100644 --- a/src/training/metrics/metrics.hpp +++ b/src/training/metrics/metrics.hpp @@ -137,11 +137,6 @@ namespace lfs::training { std::unique_ptr _reporter; // Helper functions - lfs::core::Tensor load_eval_mask(lfs::core::Camera* cam, lfs::core::Tensor& gt_image, - bool alpha_as_mask) const; lfs::core::Tensor apply_depth_colormap(const lfs::core::Tensor& depth_normalized) const; - - // Create dataloader from dataset - auto make_dataloader(std::shared_ptr dataset, const int workers = 1) const; }; } // namespace lfs::training diff --git a/src/training/strategies/improved_gs_plus.cpp b/src/training/strategies/improved_gs_plus.cpp index 6966483b3..f01e3e3b9 100644 --- a/src/training/strategies/improved_gs_plus.cpp +++ b/src/training/strategies/improved_gs_plus.cpp @@ -286,7 +286,7 @@ namespace lfs::training { if (cam->is_undistort_prepared()) { params.undistort = &cam->undistort_params(); } - lfs::core::Tensor image = _image_loader->load_image_immediate(cam->image_path(), params); + lfs::core::Tensor image = _image_loader->load_camera_image_immediate(*cam, params); const int img_h = image.shape()[1]; const int img_w = image.shape()[2]; diff --git a/src/training/strategies/mrnf.cpp b/src/training/strategies/mrnf.cpp index ad074189c..90dd1020d 100644 --- a/src/training/strategies/mrnf.cpp +++ b/src/training/strategies/mrnf.cpp @@ -1428,7 +1428,7 @@ namespace lfs::training { params.undistort = &cam->undistort_params(); } - lfs::core::Tensor image = _image_loader->load_image_immediate(cam->image_path(), params); + lfs::core::Tensor image = _image_loader->load_camera_image_immediate(*cam, params); const int img_h = static_cast(image.shape()[1]); const int img_w = static_cast(image.shape()[2]); diff --git a/src/training/trainer.cpp b/src/training/trainer.cpp index 0f340730a..dac166a0f 100644 --- a/src/training/trainer.cpp +++ b/src/training/trainer.cpp @@ -351,7 +351,8 @@ namespace lfs::training { const lfs::core::param::OptimizationParameters& opt_params, lfs::io::PipelinedImageLoader& image_loader) { try { - auto mask = image_loader.load_image_immediate( + auto mask = image_loader.load_camera_image_immediate( + camera, camera.mask_path(), make_metrics_load_params(gt_config, camera, false, false)); mask = normalize_mask_tensor(std::move(mask)); @@ -388,64 +389,55 @@ namespace lfs::training { const Trainer::GTLoadConfigSnapshot& gt_config, const lfs::core::param::OptimizationParameters& opt_params) { try { - auto [img_data, width, height, channels] = lfs::core::load_image_with_alpha( - camera.image_path(), gt_config.resize_factor, gt_config.max_width); + lfs::io::PipelinedLoaderConfig loader_config; + loader_config.prefetch_count = 1; + loader_config.output_queue_size = 1; + loader_config.io_threads = 1; + loader_config.cold_process_threads = 1; + loader_config.use_filesystem_cache = false; + + auto params = make_metrics_load_params(gt_config, camera, true, true); + + lfs::io::ImageRequest request; + request.sequence_id = 0; + request.path = camera.image_path(); + request.params = params; + request.extract_alpha_as_mask = true; + request.alpha_mask_params.invert = opt_params.invert_masks; + request.alpha_mask_params.threshold = opt_params.mask_threshold; + request.undistort = params.undistort; + request.cube_face_projection = camera.cube_face_projection(); + + lfs::io::PipelinedImageLoader loader(loader_config); + loader.prefetch({request}); + + while (loader.in_flight_count() > 0) { + auto ready = loader.try_get_for(loader_config.output_wait_timeout); + if (!ready) + continue; - if (!img_data || channels != 4) { - if (img_data) { - lfs::core::free_image(img_data); + if (!ready->tensor.is_valid()) { + return std::unexpected("failed to decode RGBA image"); + } + if (!ready->mask || !ready->mask->is_valid()) { + return std::unexpected("failed to extract alpha mask"); } - return std::unexpected("failed to decode RGBA image"); - } - - const auto H = static_cast(height); - const auto W = static_cast(width); - - auto cpu_tensor = lfs::core::Tensor::from_blob( - img_data, lfs::core::TensorShape({H, W, 4}), - lfs::core::Device::CPU, lfs::core::DataType::UInt8); - auto gpu_uint8 = cpu_tensor.to(lfs::core::Device::CUDA); - lfs::core::free_image(img_data); - - auto rgb = lfs::core::Tensor::zeros( - lfs::core::TensorShape({3, H, W}), - lfs::core::Device::CUDA, lfs::core::DataType::UInt8); - auto mask = lfs::core::Tensor::zeros( - lfs::core::TensorShape({H, W}), - lfs::core::Device::CUDA, lfs::core::DataType::Float32); - - lfs::io::cuda::launch_uint8_rgba_split_to_uint8_rgb_and_float32_alpha( - gpu_uint8.ptr(), rgb.ptr(), mask.ptr(), H, W, nullptr); - if (opt_params.invert_masks) { - lfs::io::cuda::launch_mask_invert(mask.ptr(), H, W, nullptr); - } - if (opt_params.mask_threshold > 0.0f) { - lfs::io::cuda::launch_mask_threshold(mask.ptr(), H, W, opt_params.mask_threshold, nullptr); - } + auto mask = normalize_mask_tensor(std::move(*ready->mask)); + if (mask.dtype() == lfs::core::DataType::UInt8 || + mask.dtype() == lfs::core::DataType::Bool) { + mask = opt_params.mask_threshold > 0.0f ? mask.gt(0) : mask.ge(128); + } else { + mask = mask.ge(0.5f); + } + mask = mask.to(lfs::core::DataType::UInt8).contiguous(); - if (camera.is_undistort_prepared()) { - const auto scaled = lfs::core::scale_undistort_params( - camera.undistort_params(), - static_cast(W), - static_cast(H)); - auto rgb_float = rgb.to(lfs::core::DataType::Float32) / 255.0f; - rgb_float = lfs::core::undistort_image(rgb_float, scaled, nullptr); - auto rgb_uint8 = lfs::core::Tensor::empty( - rgb_float.shape(), lfs::core::Device::CUDA, lfs::core::DataType::UInt8); - lfs::io::cuda::launch_float32_chw_to_uint8_chw( - rgb_float.ptr(), - rgb_uint8.ptr(), - rgb_float.shape()[1], - rgb_float.shape()[2], - rgb_float.shape()[0], - nullptr); - rgb = std::move(rgb_uint8); - mask = lfs::core::undistort_mask(mask, scaled, nullptr); + return LoadedCameraMetricsInputs{ + .gt_image = std::move(ready->tensor), + .mask = std::move(mask)}; } - mask = mask.ge(0.5f).to(lfs::core::DataType::UInt8).contiguous(); - return LoadedCameraMetricsInputs{.gt_image = std::move(rgb), .mask = std::move(mask)}; + return std::unexpected("failed to decode RGBA image"); } catch (const std::exception& e) { return std::unexpected(e.what()); } @@ -480,8 +472,8 @@ namespace lfs::training { LoadedCameraMetricsInputs inputs; try { - inputs.gt_image = loader->load_image_immediate( - camera.image_path(), + inputs.gt_image = loader->load_camera_image_immediate( + camera, make_metrics_load_params(gt_config, camera, true, true)); } catch (const std::exception& e) { return std::unexpected(e.what()); @@ -1121,11 +1113,17 @@ namespace lfs::training { BilateralGrid::Config config; config.lr = params_.optimization.bilateral_grid_lr; - // BilateralGrid is indexed with cam->uid() in the training loop. Those UIDs stay - // in the original camera space even when train/val splits are enabled, so the grid - // must be sized for the full camera set rather than only the training subset. + int image_slots = 0; + for (const auto& cam : train_dataset_->get_cameras()) { + if (cam) + image_slots = std::max(image_slots, cam->uid() + 1); + } + image_slots = std::max(image_slots, 1); + + // BilateralGrid is indexed with cam->uid() in the training loop, so + // allocate by UID range rather than by camera count. bilateral_grid_ = std::make_unique( - static_cast(total_cameras_count_), + image_slots, params_.optimization.bilateral_grid_X, params_.optimization.bilateral_grid_Y, params_.optimization.bilateral_grid_W, @@ -1136,7 +1134,7 @@ namespace lfs::training { params_.optimization.bilateral_grid_X, params_.optimization.bilateral_grid_Y, params_.optimization.bilateral_grid_W, - total_cameras_count_, + image_slots, train_dataset_size_); return {}; @@ -2302,31 +2300,68 @@ namespace lfs::training { if (source_cameras.empty()) { return std::unexpected("Scene has no active cameras enabled for training"); } - - if (params.optimization.enable_eval) { - for (const auto& camera : source_cameras) { - switch (camera->split()) { - case lfs::core::CameraSplit::Train: - train_cameras.push_back(camera); - break; - case lfs::core::CameraSplit::Eval: - val_cameras.push_back(camera); - break; - default: - assert(false && "Camera split must be Train or Eval"); - break; - } - } - assert(train_cameras.size() + val_cameras.size() == source_cameras.size()); - } } else if (base_dataset_) { source_cameras = base_dataset_->get_cameras(); } else { return std::unexpected("No camera source available"); } + if (scene_ && params.optimization.undistort) { + auto expanded_cameras = expandEquirectangularCamerasForUndistort(source_cameras); + if (!expanded_cameras) + return std::unexpected(expanded_cameras.error()); + source_cameras = std::move(*expanded_cameras); + } + + if (scene_ && params.optimization.enable_eval) { + for (const auto& camera : source_cameras) { + switch (camera->split()) { + case lfs::core::CameraSplit::Train: + train_cameras.push_back(camera); + break; + case lfs::core::CameraSplit::Eval: + val_cameras.push_back(camera); + break; + default: + assert(false && "Camera split must be Train or Eval"); + break; + } + } + assert(train_cameras.size() + val_cameras.size() == source_cameras.size()); + } + total_cameras_count_ = source_cameras.size(); + if (params.optimization.gut) { + for (const auto& cam : source_cameras) { + if (cam && cam->camera_model_type() == core::CameraModelType::ORTHO) { + return std::unexpected("Training on cameras with ortho model is not supported yet."); + } + } + } else if (params.optimization.undistort) { + for (const auto& cam : source_cameras) { + if (!cam) + continue; + const auto model = cam->camera_model_type(); + if (model != core::CameraModelType::PINHOLE && + model != core::CameraModelType::EQUIRECTANGULAR) { + return std::unexpected("Unsupported non-pinhole cameras detected. Use --gut for native non-pinhole training."); + } + } + } else { + for (const auto& cam : source_cameras) { + if (!cam) + continue; + if (cam->radial_distortion().numel() != 0 || + cam->tangential_distortion().numel() != 0) { + return std::unexpected("Distorted images detected. Use --undistort or --gut to train on cameras with distortion."); + } + if (cam->camera_model_type() != core::CameraModelType::PINHOLE) { + return std::unexpected("Spherical/non-pinhole cameras detected. Use --undistort for internal spherical tangent-view sampling or --gut for native non-pinhole training."); + } + } + } + if (auto result = initialize_camera_loss_heatmap(source_cameras); !result) { return std::unexpected(result.error()); } @@ -3194,10 +3229,10 @@ namespace lfs::training { } else if (!params_.optimization.undistort || !cam->is_undistort_prepared()) { if (cam->radial_distortion().numel() != 0 || cam->tangential_distortion().numel() != 0) { - return std::unexpected("Distorted images detected. Use --gut or --undistort to train on cameras with distortion."); + return std::unexpected("Distorted images detected. Use --undistort or --gut to train on cameras with distortion."); } if (cam->camera_model_type() != core::CameraModelType::PINHOLE) { - return std::unexpected("Use --gut or --undistort to train on cameras with non-pinhole model."); + return std::unexpected("Spherical/non-pinhole cameras detected. Use --undistort for internal spherical tangent-view sampling or --gut for native non-pinhole training."); } } diff --git a/src/training/training_setup.cpp b/src/training/training_setup.cpp index 012d6a26f..5c97d7c68 100644 --- a/src/training/training_setup.cpp +++ b/src/training/training_setup.cpp @@ -3,8 +3,10 @@ * SPDX-License-Identifier: GPL-3.0-or-later */ #include "training_setup.hpp" +#include "core/camera.hpp" #include "core/cuda/sh_layout.cuh" #include "core/events.hpp" +#include "core/image_io.hpp" #include "core/logger.hpp" #include "core/mesh_data.hpp" #include "core/path_utils.hpp" @@ -15,12 +17,19 @@ #include "dataset.hpp" #include "io/loader.hpp" #include +#include +#include +#include #include +#include #include #include #include #include +#include +#include #include +#include namespace lfs::training { @@ -91,6 +100,256 @@ namespace lfs::training { return lfs::io::CentralizeDataset::Off; } + constexpr int kSphericalUndistortViewCount = 12; + constexpr float kSphericalUndistortFovDegrees = 90.0f; + + struct SphericalViewDefinition { + std::string_view suffix; + std::array forward; + }; + + constexpr float kIcosA = 0.5257311121191336f; + constexpr float kIcosB = 0.8506508083520399f; + + constexpr std::array kSphericalUndistortViews{{ + {"ico_00", {0.f, kIcosA, kIcosB}}, + {"ico_01", {0.f, -kIcosA, kIcosB}}, + {"ico_02", {0.f, kIcosA, -kIcosB}}, + {"ico_03", {0.f, -kIcosA, -kIcosB}}, + {"ico_04", {kIcosA, kIcosB, 0.f}}, + {"ico_05", {-kIcosA, kIcosB, 0.f}}, + {"ico_06", {kIcosA, -kIcosB, 0.f}}, + {"ico_07", {-kIcosA, -kIcosB, 0.f}}, + {"ico_08", {kIcosB, 0.f, kIcosA}}, + {"ico_09", {-kIcosB, 0.f, kIcosA}}, + {"ico_10", {kIcosB, 0.f, -kIcosA}}, + {"ico_11", {-kIcosB, 0.f, -kIcosA}}, + }}; + + using Vec3 = std::array; + + [[nodiscard]] Vec3 cross(const Vec3& a, const Vec3& b) { + return { + a[1] * b[2] - a[2] * b[1], + a[2] * b[0] - a[0] * b[2], + a[0] * b[1] - a[1] * b[0]}; + } + + [[nodiscard]] float length(const Vec3& v) { + return std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); + } + + [[nodiscard]] Vec3 normalize(const Vec3& v) { + const float len = length(v); + if (len <= 0.0f) + return {0.0f, 0.0f, 1.0f}; + return {v[0] / len, v[1] / len, v[2] / len}; + } + + [[nodiscard]] std::array makeViewFromForward(const Vec3& forward_in) { + const Vec3 forward = normalize(forward_in); + constexpr Vec3 world_up{0.0f, 1.0f, 0.0f}; + constexpr Vec3 fallback_up{0.0f, 0.0f, 1.0f}; + + Vec3 right = cross(world_up, forward); + if (length(right) < 1e-5f) + right = cross(fallback_up, forward); + right = normalize(right); + const Vec3 up = cross(forward, right); + + return { + right[0], right[1], right[2], + up[0], up[1], up[2], + forward[0], forward[1], forward[2]}; + } + + [[nodiscard]] float focalFromFov(const int image_size, const float fov_degrees) { + constexpr float kPi = 3.14159265358979323846f; + const float fov = std::clamp(fov_degrees, 1.0f, 179.0f); + return 0.5f * static_cast(image_size) / + std::tan(0.5f * fov * kPi / 180.0f); + } + + [[nodiscard]] std::string sanitizeFilenameComponent(const std::string_view value) { + std::string out; + out.reserve(value.size()); + for (const unsigned char ch : value) { + if (std::isalnum(ch) || ch == '-' || ch == '_') { + out.push_back(static_cast(ch)); + } else { + out.push_back('_'); + } + } + while (!out.empty() && out.back() == '_') + out.pop_back(); + return out.empty() ? "camera" : out; + } + + [[nodiscard]] int resolveSphericalViewSize(const int width, const int height, const float fov_degrees) { + const float fov = std::clamp(fov_degrees, 1.0f, 179.0f); + const int by_width = std::max(1, static_cast(std::round(static_cast(width) * fov / 360.0f))); + const int by_height = std::max(1, static_cast(std::round(static_cast(height) * fov / 180.0f))); + return std::max(1, std::min(by_width, by_height)); + } + + [[nodiscard]] std::shared_ptr makeSphericalViewCamera( + const lfs::core::Camera& source, + const SphericalViewDefinition& view, + const int view_size, + const int source_width, + const int source_height, + const int uid, + const std::string& image_name, + const bool has_alpha) { + const auto pano_to_view = makeViewFromForward(view.forward); + auto R_cpu = source.R().cpu(); + if (!R_cpu.is_contiguous()) + R_cpu = R_cpu.contiguous(); + auto T_cpu = source.T().cpu(); + if (!T_cpu.is_contiguous()) + T_cpu = T_cpu.contiguous(); + + auto R_acc = R_cpu.accessor(); + auto T_acc = T_cpu.accessor(); + + std::vector R_face(9, 0.0f); + std::vector T_face(3, 0.0f); + for (int row = 0; row < 3; ++row) { + for (int col = 0; col < 3; ++col) { + for (int k = 0; k < 3; ++k) + R_face[row * 3 + col] += pano_to_view[row * 3 + k] * R_acc(k, col); + } + for (int k = 0; k < 3; ++k) + T_face[row] += pano_to_view[row * 3 + k] * T_acc(k); + } + + const float focal = focalFromFov(view_size, kSphericalUndistortFovDegrees); + const float center = 0.5f * static_cast(view_size); + auto camera = std::make_shared( + lfs::core::Tensor::from_vector(R_face, {3, 3}, lfs::core::Device::CPU), + lfs::core::Tensor::from_vector(T_face, {3}, lfs::core::Device::CPU), + focal, + focal, + center, + center, + lfs::core::Tensor::empty({0}, lfs::core::Device::CPU), + lfs::core::Tensor::empty({0}, lfs::core::Device::CPU), + lfs::core::CameraModelType::PINHOLE, + image_name, + source.image_path(), + source.has_mask() ? source.mask_path() : std::filesystem::path{}, + view_size, + view_size, + uid, + source.camera_id()); + camera->set_has_alpha(has_alpha); + camera->set_split(source.split()); + camera->set_cube_face_projection(lfs::core::CubeFaceProjection{ + .pano_to_face = pano_to_view, + .face_size = view_size, + .source_width = source_width, + .source_height = source_height, + .fov_degrees = kSphericalUndistortFovDegrees}); + return camera; + } + + [[nodiscard]] std::expected>, std::string> + expandEquirectangularCamerasForUndistortImpl( + const std::vector>& cameras) { + int next_uid = 0; + const bool has_passthrough_cameras = std::any_of(cameras.begin(), cameras.end(), [](const auto& camera) { + return camera && + camera->camera_model_type() != lfs::core::CameraModelType::EQUIRECTANGULAR; + }); + if (has_passthrough_cameras) { + for (const auto& camera : cameras) { + if (camera) + next_uid = std::max(next_uid, camera->uid() + 1); + } + } + + std::vector> expanded; + expanded.reserve(cameras.size() * kSphericalUndistortViews.size()); + + size_t panoramas = 0; + size_t virtual_views = 0; + size_t ignored_depth_maps = 0; + int min_face_size = std::numeric_limits::max(); + int max_face_size = 0; + + for (const auto& camera : cameras) { + if (!camera || + camera->camera_model_type() != lfs::core::CameraModelType::EQUIRECTANGULAR) { + expanded.push_back(camera); + if (!has_passthrough_cameras && camera) + next_uid = std::max(next_uid, camera->uid() + 1); + continue; + } + + ++panoramas; + int source_width = 0; + int source_height = 0; + int source_channels = 0; + try { + std::tie(source_width, source_height, source_channels) = + lfs::core::get_image_info(camera->image_path()); + } catch (const std::exception& e) { + return std::unexpected(std::format("Failed to inspect panorama '{}': {}", + lfs::core::path_to_utf8(camera->image_path()), e.what())); + } + if (source_width <= 0 || source_height <= 0 || source_channels <= 0) { + return std::unexpected(std::format("Failed to inspect panorama '{}'", + lfs::core::path_to_utf8(camera->image_path()))); + } + + const int view_size = resolveSphericalViewSize(source_width, source_height, kSphericalUndistortFovDegrees); + min_face_size = std::min(min_face_size, view_size); + max_face_size = std::max(max_face_size, view_size); + + const std::string stem = sanitizeFilenameComponent( + std::filesystem::path(camera->image_name()).stem().string()); + const std::string base_name = std::format("{:06d}_{}", camera->uid(), stem); + + if (camera->has_depth()) + ++ignored_depth_maps; + + for (size_t view_index = 0; view_index < kSphericalUndistortViews.size(); ++view_index) { + const auto& view = kSphericalUndistortViews[view_index]; + const std::string image_name = + std::format("{}_{}.png", base_name, view.suffix); + expanded.push_back(makeSphericalViewCamera( + *camera, + view, + view_size, + source_width, + source_height, + next_uid++, + image_name, + camera->has_alpha())); + ++virtual_views; + } + } + + if (panoramas > 0) { + LOG_INFO("Undistort expanded {} equirectangular panorama{} into {} internal spherical pinhole views " + "({} views/panorama, {:.0f} deg FOV, face size {} px auto, sampled internally from source panoramas)", + panoramas, + panoramas == 1 ? "" : "s", + virtual_views, + kSphericalUndistortViews.size(), + kSphericalUndistortFovDegrees, + min_face_size == max_face_size ? std::to_string(max_face_size) + : std::format("{}-{}", min_face_size, max_face_size)); + if (ignored_depth_maps > 0) { + LOG_WARN("Equirectangular spherical undistort ignored {} depth map{}; perspective depth conversion is not implemented", + ignored_depth_maps, + ignored_depth_maps == 1 ? "" : "s"); + } + } + + return expanded; + } + void applyTrainingSHDegree(lfs::core::SplatData& splat, const int target_degree) { const int before = splat.get_max_sh_degree(); if (splat.set_sh_degree(target_degree)) { @@ -376,6 +635,12 @@ namespace lfs::training { } } // namespace + std::expected>, std::string> + expandEquirectangularCamerasForUndistort( + const std::vector>& cameras) { + return expandEquirectangularCamerasForUndistortImpl(cameras); + } + std::expected migrateTrainingModelToAllocator( const lfs::core::param::TrainingParameters& params, lfs::core::SplatData& model, @@ -501,22 +766,36 @@ namespace lfs::training { } } - const auto& cameras = data.cameras; const bool enable_eval = params.optimization.enable_eval; const int test_every = params.dataset.test_every; + for (size_t i = 0; i < data.cameras.size(); ++i) { + const bool is_eval = enable_eval && (i % test_every) == 0; + data.cameras[i]->set_split(is_eval ? lfs::core::CameraSplit::Eval : lfs::core::CameraSplit::Train); + } + + std::vector> undistorted_cameras; + const auto* cameras_ptr = &data.cameras; + if (params.optimization.undistort) { + auto expanded_cameras = expandEquirectangularCamerasForUndistort(data.cameras); + if (!expanded_cameras) + return std::unexpected(expanded_cameras.error()); + undistorted_cameras = std::move(*expanded_cameras); + cameras_ptr = &undistorted_cameras; + } + const auto& cameras = *cameras_ptr; + size_t train_count = 0; size_t val_count = 0; size_t mask_count = 0; - for (size_t i = 0; i < cameras.size(); ++i) { - const bool is_eval = enable_eval && (i % test_every) == 0; - cameras[i]->set_split(is_eval ? lfs::core::CameraSplit::Eval : lfs::core::CameraSplit::Train); + for (const auto& camera : cameras) { + const bool is_eval = enable_eval && camera->split() == lfs::core::CameraSplit::Eval; if (is_eval) { val_count++; } else { train_count++; } - if (cameras[i]->has_mask()) { + if (camera->has_mask()) { mask_count++; } } @@ -529,7 +808,7 @@ namespace lfs::training { train_count); for (size_t i = 0; i < cameras.size(); ++i) { - if (!enable_eval || (i % test_every) != 0) { + if (!enable_eval || cameras[i]->split() != lfs::core::CameraSplit::Eval) { scene.addCamera(cameras[i]->image_name(), train_cameras_id, cameras[i]); } } @@ -541,7 +820,7 @@ namespace lfs::training { val_count); for (size_t i = 0; i < cameras.size(); ++i) { - if ((i % test_every) == 0) { + if (cameras[i]->split() == lfs::core::CameraSplit::Eval) { scene.addCamera(cameras[i]->image_name(), val_cameras_id, cameras[i]); } } @@ -870,16 +1149,30 @@ namespace lfs::training { scene.addPointCloud("PointCloud", createRandomPointCloud(), dataset_id); } - const auto& cameras = data.cameras; const bool enable_eval = params.optimization.enable_eval; const int test_every = params.dataset.test_every; - size_t train_count = 0, val_count = 0, mask_count = 0; - for (size_t i = 0; i < cameras.size(); ++i) { + for (size_t i = 0; i < data.cameras.size(); ++i) { const bool is_val = enable_eval && (i % test_every) == 0; - cameras[i]->set_split(is_val ? lfs::core::CameraSplit::Eval : lfs::core::CameraSplit::Train); + data.cameras[i]->set_split(is_val ? lfs::core::CameraSplit::Eval : lfs::core::CameraSplit::Train); + } + + std::vector> undistorted_cameras; + const auto* cameras_ptr = &data.cameras; + if (params.optimization.undistort) { + auto expanded_cameras = expandEquirectangularCamerasForUndistort(data.cameras); + if (!expanded_cameras) + return std::unexpected(expanded_cameras.error()); + undistorted_cameras = std::move(*expanded_cameras); + cameras_ptr = &undistorted_cameras; + } + const auto& cameras = *cameras_ptr; + + size_t train_count = 0, val_count = 0, mask_count = 0; + for (const auto& camera : cameras) { + const bool is_val = enable_eval && camera->split() == lfs::core::CameraSplit::Eval; is_val ? ++val_count : ++train_count; - if (cameras[i]->has_mask()) + if (camera->has_mask()) ++mask_count; } @@ -888,7 +1181,7 @@ namespace lfs::training { std::format("Training ({})", train_count), cameras_group_id, train_count); for (size_t i = 0; i < cameras.size(); ++i) { - if (!enable_eval || (i % test_every) != 0) { + if (!enable_eval || cameras[i]->split() != lfs::core::CameraSplit::Eval) { scene.addCamera(cameras[i]->image_name(), train_cameras_id, cameras[i]); } } @@ -897,7 +1190,7 @@ namespace lfs::training { const auto val_cameras_id = scene.addCameraGroup( std::format("Validation ({})", val_count), cameras_group_id, val_count); for (size_t i = 0; i < cameras.size(); ++i) { - if ((i % test_every) == 0) { + if (cameras[i]->split() == lfs::core::CameraSplit::Eval) { scene.addCamera(cameras[i]->image_name(), val_cameras_id, cameras[i]); } } diff --git a/src/training/training_setup.hpp b/src/training/training_setup.hpp index 7a7da00bf..5620bd2db 100644 --- a/src/training/training_setup.hpp +++ b/src/training/training_setup.hpp @@ -8,11 +8,14 @@ #include "core/splat_data.hpp" #include "io/loader.hpp" #include +#include #include +#include namespace lfs::core { + class Camera; class Scene; -} +} // namespace lfs::core namespace lfs::training { /** @@ -36,6 +39,17 @@ namespace lfs::training { const lfs::core::param::TrainingParameters& params, lfs::core::Scene& scene); + /** + * @brief Expand equirectangular cameras into internal spherical pinhole views. + * + * Non-equirectangular cameras are preserved. The returned virtual cameras keep + * the source image path and sample the requested tangent view through the native + * image loader cache. + */ + std::expected>, std::string> + expandEquirectangularCamerasForUndistort( + const std::vector>& cameras); + /** * @brief Initialize training model from point cloud * diff --git a/tests/test_undistort.cpp b/tests/test_undistort.cpp index 0b8098889..8e042bccd 100644 --- a/tests/test_undistort.cpp +++ b/tests/test_undistort.cpp @@ -5,6 +5,10 @@ #include "core/cuda/undistort/undistort.hpp" #include "core/image_io.hpp" #include "io/formats/colmap.hpp" +#include "training/training_setup.hpp" +#include +#include +#include #include #include @@ -18,6 +22,61 @@ namespace { constexpr float TEST_CY = 240.0f; constexpr int TEST_W = 640; constexpr int TEST_H = 480; + constexpr float SPHERICAL_UNDISTORT_FOV = 90.0f; + constexpr size_t SPHERICAL_UNDISTORT_VIEW_COUNT = 12; + constexpr float ICOS_A = 0.5257311121191336f; + constexpr float ICOS_B = 0.8506508083520399f; + constexpr std::array, SPHERICAL_UNDISTORT_VIEW_COUNT> SPHERICAL_UNDISTORT_FORWARDS{{ + {0.f, ICOS_A, ICOS_B}, + {0.f, -ICOS_A, ICOS_B}, + {0.f, ICOS_A, -ICOS_B}, + {0.f, -ICOS_A, -ICOS_B}, + {ICOS_A, ICOS_B, 0.f}, + {-ICOS_A, ICOS_B, 0.f}, + {ICOS_A, -ICOS_B, 0.f}, + {-ICOS_A, -ICOS_B, 0.f}, + {ICOS_B, 0.f, ICOS_A}, + {-ICOS_B, 0.f, ICOS_A}, + {ICOS_B, 0.f, -ICOS_A}, + {-ICOS_B, 0.f, -ICOS_A}, + }}; + + using Vec3 = std::array; + + Vec3 cross(const Vec3& a, const Vec3& b) { + return { + a[1] * b[2] - a[2] * b[1], + a[2] * b[0] - a[0] * b[2], + a[0] * b[1] - a[1] * b[0]}; + } + + float length(const Vec3& v) { + return std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); + } + + Vec3 normalize(const Vec3& v) { + const float len = length(v); + if (len <= 0.0f) + return {0.0f, 0.0f, 1.0f}; + return {v[0] / len, v[1] / len, v[2] / len}; + } + + std::array expected_spherical_matrix(const Vec3& forward_in) { + const Vec3 forward = normalize(forward_in); + constexpr Vec3 world_up{0.0f, 1.0f, 0.0f}; + constexpr Vec3 fallback_up{0.0f, 0.0f, 1.0f}; + + Vec3 right = cross(world_up, forward); + if (length(right) < 1e-5f) + right = cross(fallback_up, forward); + right = normalize(right); + const Vec3 up = cross(forward, right); + + return { + right[0], right[1], right[2], + up[0], up[1], up[2], + forward[0], forward[1], forward[2]}; + } void validate_params(const UndistortParams& p, int src_w, int src_h) { EXPECT_GT(p.dst_width, 0); @@ -588,6 +647,111 @@ TEST_F(UndistortCameraTest, EquirectangularModelDoesNotUseUndistortion) { EXPECT_FALSE(equirect_cam.is_undistort_prepared()); } +TEST_F(UndistortCameraTest, EquirectangularSphericalExpansionUsesVirtualPinholeCameras) { + auto& source = cameras_[0]; + auto equirect_cam = std::make_shared( + source->R(), source->T(), + source->focal_x(), source->focal_y(), + source->center_x(), source->center_y(), + Tensor(), Tensor(), + CameraModelType::EQUIRECTANGULAR, + "pano.jpg", source->image_path(), "", + source->camera_width(), source->camera_height(), 7, 99); + equirect_cam->set_has_alpha(true); + equirect_cam->set_split(CameraSplit::Eval); + + auto expanded = lfs::training::expandEquirectangularCamerasForUndistort({equirect_cam}); + + ASSERT_TRUE(expanded.has_value()) << expanded.error(); + ASSERT_EQ(expanded->size(), SPHERICAL_UNDISTORT_VIEW_COUNT); + + const auto [source_width, source_height, source_channels] = get_image_info(source->image_path()); + ASSERT_GT(source_channels, 0); + const int expected_face_size = std::max(1, std::min( + static_cast(std::round(static_cast(source_width) * SPHERICAL_UNDISTORT_FOV / 360.0f)), + static_cast(std::round(static_cast(source_height) * SPHERICAL_UNDISTORT_FOV / 180.0f)))); + const float expected_focal = 0.5f * static_cast(expected_face_size) / + std::tan(0.5f * SPHERICAL_UNDISTORT_FOV * 3.14159265358979323846f / 180.0f); + + for (size_t i = 0; i < expanded->size(); ++i) { + const auto& face = (*expanded)[i]; + ASSERT_TRUE(face); + EXPECT_EQ(face->camera_model_type(), CameraModelType::PINHOLE); + EXPECT_TRUE(face->has_cube_face_projection()); + EXPECT_EQ(face->image_path(), source->image_path()); + EXPECT_EQ(face->split(), CameraSplit::Eval); + EXPECT_TRUE(face->has_alpha()); + EXPECT_EQ(face->camera_id(), 99); + EXPECT_EQ(face->uid(), static_cast(i)); + EXPECT_EQ(face->camera_width(), expected_face_size); + EXPECT_EQ(face->camera_height(), expected_face_size); + EXPECT_NEAR(face->focal_x(), expected_focal, 1e-4f); + EXPECT_NEAR(face->focal_y(), expected_focal, 1e-4f); + ASSERT_TRUE(face->cube_face_projection().has_value()); + EXPECT_EQ(face->cube_face_projection()->face_size, expected_face_size); + EXPECT_EQ(face->cube_face_projection()->source_width, source_width); + EXPECT_EQ(face->cube_face_projection()->source_height, source_height); + EXPECT_NEAR(face->cube_face_projection()->fov_degrees, SPHERICAL_UNDISTORT_FOV, 1e-4f); + const auto& pano_to_face = face->cube_face_projection()->pano_to_face; + const auto expected_matrix = expected_spherical_matrix(SPHERICAL_UNDISTORT_FORWARDS[i]); + for (size_t j = 0; j < pano_to_face.size(); ++j) + EXPECT_NEAR(pano_to_face[j], expected_matrix[j], 1e-6f); + for (int row = 0; row < 3; ++row) { + const float row_len = + pano_to_face[row * 3 + 0] * pano_to_face[row * 3 + 0] + + pano_to_face[row * 3 + 1] * pano_to_face[row * 3 + 1] + + pano_to_face[row * 3 + 2] * pano_to_face[row * 3 + 2]; + EXPECT_NEAR(row_len, 1.0f, 1e-5f); + } + const float dot01 = pano_to_face[0] * pano_to_face[3] + + pano_to_face[1] * pano_to_face[4] + + pano_to_face[2] * pano_to_face[5]; + const float dot02 = pano_to_face[0] * pano_to_face[6] + + pano_to_face[1] * pano_to_face[7] + + pano_to_face[2] * pano_to_face[8]; + const float dot12 = pano_to_face[3] * pano_to_face[6] + + pano_to_face[4] * pano_to_face[7] + + pano_to_face[5] * pano_to_face[8]; + EXPECT_NEAR(dot01, 0.0f, 1e-5f); + EXPECT_NEAR(dot02, 0.0f, 1e-5f); + EXPECT_NEAR(dot12, 0.0f, 1e-5f); + } +} + +TEST_F(UndistortCameraTest, SphericalExpansionKeepsPassthroughCamerasAndAvoidsUidCollisions) { + auto& source = cameras_[0]; + auto equirect_cam = std::make_shared( + source->R(), source->T(), + source->focal_x(), source->focal_y(), + source->center_x(), source->center_y(), + Tensor(), Tensor(), + CameraModelType::EQUIRECTANGULAR, + "pano.jpg", source->image_path(), "", + source->camera_width(), source->camera_height(), 2, 20); + auto pinhole_cam = std::make_shared( + source->R(), source->T(), + source->focal_x(), source->focal_y(), + source->center_x(), source->center_y(), + Tensor(), Tensor(), + CameraModelType::PINHOLE, + "pinhole.jpg", source->image_path(), "", + source->camera_width(), source->camera_height(), 12, 12); + + auto expanded = lfs::training::expandEquirectangularCamerasForUndistort({equirect_cam, pinhole_cam}); + + ASSERT_TRUE(expanded.has_value()) << expanded.error(); + ASSERT_EQ(expanded->size(), SPHERICAL_UNDISTORT_VIEW_COUNT + 1); + for (size_t i = 0; i < SPHERICAL_UNDISTORT_VIEW_COUNT; ++i) { + ASSERT_TRUE((*expanded)[i]); + EXPECT_EQ((*expanded)[i]->camera_model_type(), CameraModelType::PINHOLE); + EXPECT_TRUE((*expanded)[i]->has_cube_face_projection()); + EXPECT_EQ((*expanded)[i]->uid(), static_cast(13 + i)); + } + EXPECT_EQ((*expanded)[SPHERICAL_UNDISTORT_VIEW_COUNT], pinhole_cam); + EXPECT_FALSE((*expanded)[SPHERICAL_UNDISTORT_VIEW_COUNT]->has_cube_face_projection()); + EXPECT_EQ((*expanded)[SPHERICAL_UNDISTORT_VIEW_COUNT]->uid(), 12); +} + TEST(UndistortScale, ScaleUndistortParams) { const auto radial = Tensor::from_vector({-0.1f, 0.02f}, TensorShape({2}), Device::CPU); const auto params = compute_undistort_params(