Feat: masks, zip files, leaner, faster, OPENCV_FULL undistortion

CMakeLists.txt

@@ -2,7 +2,12 @@ cmake_minimum_required(VERSION 3.21)
 project(opensplat)
 
 set(OPENSPLAT_BUILD_SIMPLE_TRAINER OFF CACHE BOOL "Build simple trainer applications")
-set(GPU_RUNTIME "CUDA" CACHE STRING "HIP or CUDA or MPS")
+if(APPLE)
+    set(GPU_RUNTIME_DEFAULT "MPS")
+else()
+    set(GPU_RUNTIME_DEFAULT "CUDA")
+endif()
+set(GPU_RUNTIME "${GPU_RUNTIME_DEFAULT}" CACHE STRING "HIP or CUDA or MPS or CPU")
 set(OPENCV_DIR "OPENCV_DIR-NOTFOUND" CACHE PATH "Path to the OPENCV installation directory")
 set(OPENSPLAT_MAX_CUDA_COMPATIBILITY OFF CACHE BOOL "Build for maximum CUDA device compatibility")
 set(OPENSPLAT_BUILD_VISUALIZER OFF CACHE BOOL "Build visualizer application")
@@ -13,6 +18,7 @@ set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
 
 # Read version
 file(READ "VERSION" APP_VERSION)
+string(STRIP "${APP_VERSION}" APP_VERSION)
 
 # Read git commit
 set(GIT_REV "")
@@ -146,18 +152,28 @@ elseif(GPU_RUNTIME STREQUAL "HIP")
     endif()
     list(APPEND CMAKE_PREFIX_PATH "${ROCM_ROOT}")
 elseif(GPU_RUNTIME STREQUAL "MPS")
-    find_library(FOUNDATION_LIBRARY Foundation REQUIRED)
-    find_library(METAL_FRAMEWORK    Metal      REQUIRED)
-    find_library(METALKIT_FRAMEWORK MetalKit   REQUIRED)
-    message(STATUS "Metal framework found")
+    execute_process(COMMAND xcrun -sdk macosx metal --version
+                    RESULT_VARIABLE METAL_COMPILER_RESULT
+                    OUTPUT_QUIET ERROR_QUIET)
+    if(NOT METAL_COMPILER_RESULT EQUAL 0)
+        message(WARNING "Metal compiler not found, building with CPU support only. "
+                        "Install Xcode and the Metal toolchain "
+                        "(xcodebuild -downloadComponent MetalToolchain), then re-run cmake.")
+        set(GPU_RUNTIME "CPU")
+    else()
+        find_library(FOUNDATION_LIBRARY Foundation REQUIRED)
+        find_library(METAL_FRAMEWORK    Metal      REQUIRED)
+        find_library(METALKIT_FRAMEWORK MetalKit   REQUIRED)
+        message(STATUS "Metal framework found")
 
-    set(XC_FLAGS -O3)
-    if(OPENSPLAT_USE_FAST_MATH)
-        message(STATUS "Fast math optimizations enabled for Metal")
-
-        set(XC_FLAGS ${XC_FLAGS} -ffast-math)
+        set(XC_FLAGS -O3)
+        if(OPENSPLAT_USE_FAST_MATH)
+            message(STATUS "Fast math optimizations enabled for Metal")
+
+            set(XC_FLAGS ${XC_FLAGS} -ffast-math)
+        endif()
+        set(USE_MPS ON CACHE BOOL "Use MPS for GPU acceleration")
     endif()
-    set(USE_MPS ON CACHE BOOL "Use MPS for GPU acceleration")
 else()
     set(GPU_RUNTIME "CPU")
 endif()
@@ -247,9 +263,9 @@ add_library(gsplat_cpu rasterizer/gsplat-cpu/gsplat_cpu.cpp)
 target_include_directories(gsplat_cpu PRIVATE ${TORCH_INCLUDE_DIRS})
 
 set(OPENSPLAT_SRC_FILES opensplat.cpp point_io.cpp nerfstudio.cpp model.cpp
-kdtree_tensor.cpp spherical_harmonics.cpp cv_utils.cpp utils.cpp project_gaussians.cpp
-rasterize_gaussians.cpp ssim.cpp optim_scheduler.cpp colmap.cpp opensfm.cpp openmvg.cpp input_data.cpp
-tensor_math.cpp rad.cpp zip_utils.cpp)
+kdtree_tensor.cpp spherical_harmonics.cpp cv_utils.cpp project_gaussians.cpp
+rasterize_gaussians.cpp ssim.cpp colmap.cpp opensfm.cpp openmvg.cpp input_data.cpp
+tensor_math.cpp rad.cpp zip_utils.cpp undistort.cpp)
 
 if (OPENSPLAT_BUILD_VISUALIZER)
     if (Pangolin_FOUND)

README.md

@@ -11,8 +11,6 @@ Graphics card recommended, but not required! OpenSplat runs the fastest on NVIDI
 
 Commercial use allowed and encouraged under the terms of the [AGPLv3](https://www.tldrlegal.com/license/gnu-affero-general-public-license-v3-agpl-3-0). ✅
 
-We even have a [song](https://youtu.be/1bma7XJkoDM) 🎵
-
 ## Getting Started
 
 If you're on Windows, you can [buy](http://sites.fastspring.com/masseranolabs/product/opensplatforwindows) the pre-built program. This saves you time and helps support the project ❤️. Then jump directly to the [run](#run) section. As an alternative, check the [build](#build) section below.
@@ -125,18 +123,19 @@ You will also need to install Xcode and the Xcode command line tools to compile
 1. Install Xcode from the Apple App Store.
 2. Install the command line tools with `xcode-select --install`. This might do nothing on your machine.
 3. If `xcode-select --print-path` prints `/Library/Developer/CommandLineTools`,then run `sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer`.
+4. On recent Xcode versions the Metal toolchain is a separate download. If `xcrun -sdk macosx metal --version` fails, run `xcodebuild -downloadComponent MetalToolchain`.
 
 Then run:
 
 ```
 git clone https://github.com/pierotofy/OpenSplat OpenSplat
 cd OpenSplat
 mkdir build && cd build
-cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch/ -DGPU_RUNTIME=MPS .. && make -j$(sysctl -n hw.logicalcpu)
+cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch/ .. && make -j$(sysctl -n hw.logicalcpu)
 ./opensplat
 ```
 
-If building CPU-only, remove `-DGPU_RUNTIME=MPS`.
+On macOS `GPU_RUNTIME` defaults to `MPS` (metal acceleration). If the Metal compiler isn't available, the build automatically falls back to CPU. To force a CPU-only build, pass `-DGPU_RUNTIME=CPU`.
 
 :warning: You will probably get a *libc10.dylib can’t be opened because Apple cannot check it for malicious software* error on first run. Open **System Settings** and go to **Privacy & Security** and find the **Allow** button. You might need to repeat this several times until all torch libraries are loaded.
 
@@ -245,6 +244,16 @@ You can resume training of a .PLY file by using the `--resume` option:
 ./opensplat /path/to/banana --resume ./splat.ply
 ```
 
+### Image Masks
+
+You can exclude parts of your images by adding 2D masks. Place them in a `masks` folder (also recognized: `mask`, `segmentation`, `dynamic_masks`) inside your project, named after each image (e.g. `images/IMG_001.JPG` → `masks/IMG_001.png`). Masks are grayscale images matching the input dimensions: white marks pixels to keep, black pixels to ignore.
+
+When masks are found they are applied automatically. Use `--no-masks` to ignore them.
+
+### Coordinate Reference System
+
+By default OpenSplat preserves the input coordinate reference system of the model. If you want to automatically center the result so that it displays nicely in most viewers, use `--center`.
+
 ### AMD GPU Notes
 
 To train a model with AMD GPU using docker container, you can use the following command as a reference:

VERSION

@@ -1 +1 @@
-1.1.6
+1.2.0

cv_utils.cpp

@@ -13,20 +13,6 @@ cv::Mat imreadRGB(const std::string &filename){
     return cImg;
 }
 
-void imwriteRGB(const std::string &filename, const cv::Mat &image){
-    cv::Mat rgb;
-    cv::cvtColor(image, rgb, cv::COLOR_RGB2BGR);
-    cv::imwrite(filename, rgb);
-}
-
-cv::Mat floatNxNtensorToMat(const torch::Tensor &t){
-    return cv::Mat(t.size(0), t.size(1), CV_32F, t.data_ptr());
-}
-
-torch::Tensor floatNxNMatToTensor(const cv::Mat &m){
-    return torch::from_blob(m.data, { m.rows, m.cols }, torch::kFloat32).clone();
-}
-
 cv::Mat tensorToImage(const torch::Tensor &t){
     int h = t.sizes()[0];
     int w = t.sizes()[1];

cv_utils.hpp

@@ -7,9 +7,6 @@
 #include <opencv2/imgproc.hpp>
 
 cv::Mat imreadRGB(const std::string &filename);
-void imwriteRGB(const std::string &filename, const cv::Mat &image);
-cv::Mat floatNxNtensorToMat(const torch::Tensor &t);
-torch::Tensor floatNxNMatToTensor(const cv::Mat &m);
 cv::Mat tensorToImage(const torch::Tensor &t);
 torch::Tensor imageToTensor(const cv::Mat &image);
 

input_data.cpp

@@ -1,8 +1,18 @@
 #include <filesystem>
 #include <mutex>
+#include <atomic>
+#ifdef USE_CUDA
+#include <cuda_runtime_api.h>
+#elif defined(USE_HIP)
+#include <hip/hip_runtime_api.h>
+#endif
+#ifdef __APPLE__
+#include <sys/sysctl.h>
+#endif
 #include <nlohmann/json.hpp>
 #include "input_data.hpp"
 #include "cv_utils.hpp"
+#include "undistort.hpp"
 
 namespace fs = std::filesystem;
 using namespace torch::indexing;
@@ -51,9 +61,15 @@ void Camera::loadImage(float downscaleFactor){
     }
 
     cv::Mat cImg = imreadRGB(filePath);
-
+
+    cv::Mat cMask;
+    if (!maskPath.empty()){
+        cMask = cv::imread(maskPath, cv::IMREAD_GRAYSCALE);
+        if (cMask.empty()) throw std::runtime_error("Cannot read mask " + maskPath);
+    }
+
     float rescaleF = 1.0f;
-    // If camera intrinsics don't match the image dimensions 
+    // If camera intrinsics don't match the image dimensions
     if (cImg.rows != height || cImg.cols != width){
         rescaleF = static_cast<float>(cImg.rows) / static_cast<float>(height);
     }
@@ -71,35 +87,43 @@ void Camera::loadImage(float downscaleFactor){
         cy *= scaleFactor;
     }
 
-    K = getIntrinsicsMatrix();
-    cv::Rect roi;
+    if (!cMask.empty()){
+        cv::threshold(cMask, cMask, 127, 255, cv::THRESH_BINARY);
+        if (cMask.rows != cImg.rows || cMask.cols != cImg.cols){
+            cv::resize(cMask, cMask, cv::Size(cImg.cols, cImg.rows), 0.0, 0.0, cv::INTER_LINEAR);
+        }
+    }
 
     if (hasDistortionParameters()){
-        // Undistort
-        std::vector<float> distCoeffs = undistortionParameters();
-        cv::Mat cK = floatNxNtensorToMat(K);
-        cv::Mat newK = cv::getOptimalNewCameraMatrix(cK, distCoeffs, cv::Size(cImg.cols, cImg.rows), 0, cv::Size(), &roi);
-
-        cv::Mat undistorted = cv::Mat::zeros(cImg.rows, cImg.cols, cImg.type());
-        cv::undistort(cImg, undistorted, cK, distCoeffs, newK);
-
+        UndistortParams p = computeUndistortParams(fx, fy, cx, cy, cImg.cols, cImg.rows,
+                                                   k1, k2, k3, k4, k5, k6, p1, p2);
+        cv::Mat mapx, mapy;
+        buildUndistortMaps(p, mapx, mapy);
+        cv::Mat undistorted;
+        cv::remap(cImg, undistorted, mapx, mapy, cv::INTER_LINEAR, cv::BORDER_CONSTANT);
         image = imageToTensor(undistorted);
-        K = floatNxNMatToTensor(newK);
+        if (!cMask.empty()){
+            cv::Mat remapped;
+            cv::remap(cMask, remapped, mapx, mapy, cv::INTER_LINEAR, cv::BORDER_CONSTANT);
+            cMask = remapped;
+        }
+        fx = p.dstFx;
+        fy = p.dstFy;
+        cx = p.dstCx;
+        cy = p.dstCy;
     }else{
-        roi = cv::Rect(0, 0, cImg.cols, cImg.rows);
         image = imageToTensor(cImg);
     }
 
-    // Crop to ROI
-    image = image.index({Slice(roi.y, roi.y + roi.height), Slice(roi.x, roi.x + roi.width), Slice()});
-
-    // Update parameters
     height = image.size(0);
     width = image.size(1);
-    fx = K[0][0].item<float>();
-    fy = K[1][1].item<float>();
-    cx = K[0][2].item<float>();
-    cy = K[1][2].item<float>();
+    K = getIntrinsicsMatrix();
+
+    if (!cMask.empty()){
+        torch::Tensor m = torch::from_blob(cMask.data, {cMask.rows, cMask.cols}, torch::kU8)
+                            .to(torch::kFloat32).div(255.0f).clone();
+        mask = (m >= 0.5f).to(torch::kFloat32);
+    }
 }
 
 torch::Tensor Camera::getImage(int downscaleFactor){
@@ -126,9 +150,112 @@ bool Camera::hasDistortionParameters(){
     return k1 != 0.0f || k2 != 0.0f || k3 != 0.0f || k4 != 0.0f || k5 != 0.0f || k6 != 0.0f || p1 != 0.0f || p2 != 0.0f;
 }
 
-std::vector<float> Camera::undistortionParameters(){
-    std::vector<float> p = { k1, k2, p1, p2, k3, k4, k5, k6 };
-    return p;
+torch::Tensor Camera::getMask(int downscaleFactor){
+    if (!hasMask()) return mask;
+    if (downscaleFactor <= 1) return mask;
+    if (maskPyramids.find(downscaleFactor) != maskPyramids.end()){
+        return maskPyramids[downscaleFactor];
+    }
+    torch::Tensor m = mask.unsqueeze(0).unsqueeze(0);
+    m = torch::nn::functional::interpolate(m,
+            torch::nn::functional::InterpolateFuncOptions()
+                .size(std::vector<int64_t>{ mask.size(0) / downscaleFactor, mask.size(1) / downscaleFactor })
+                .mode(torch::kBilinear).align_corners(false));
+    m = (m.squeeze(0).squeeze(0) >= 0.5f).to(torch::kFloat32);
+    maskPyramids[downscaleFactor] = m;
+    return m;
+}
+
+bool Camera::gpuCacheEnabled = true;
+
+// Half the free VRAM at first use (CUDA/HIP), a quarter of system RAM on
+// Apple unified memory, 1GB otherwise
+static long long gpuCacheBudget(){
+#ifdef USE_CUDA
+    size_t freeB = 0, totalB = 0;
+    if (cudaMemGetInfo(&freeB, &totalB) == cudaSuccess){
+        return static_cast<long long>(freeB / 2);
+    }
+#elif defined(USE_HIP)
+    size_t freeB = 0, totalB = 0;
+    if (hipMemGetInfo(&freeB, &totalB) == hipSuccess){
+        return static_cast<long long>(freeB / 2);
+    }
+#endif
+#ifdef __APPLE__
+    int64_t ram = 0;
+    size_t size = sizeof(ram);
+    if (sysctlbyname("hw.memsize", &ram, &size, nullptr, 0) == 0){
+        return ram / 4;
+    }
+#endif
+    return 1LL << 30;
+}
+
+// Cache device-side tensors per camera to avoid re-uploading every iteration
+static torch::Tensor gpuCached(std::unordered_map<int, torch::Tensor> &cache, int key,
+                               const torch::Tensor &src, const torch::Device &device){
+    if (device == torch::kCPU || !Camera::gpuCacheEnabled) return src.to(device);
+    auto it = cache.find(key);
+    if (it != cache.end()) return it->second;
+
+    static std::atomic<long long> gpuCacheBytes{0};
+    static const long long budget = gpuCacheBudget();
+    long long bytes = src.numel() * src.element_size();
+    if (gpuCacheBytes.load() + bytes > budget) return src.to(device);
+    gpuCacheBytes += bytes;
+    torch::Tensor t = src.to(device);
+    cache[key] = t;
+    return t;
+}
+
+torch::Tensor Camera::getImageGpu(int downscaleFactor, const torch::Device &device){
+    return gpuCached(gpuImageCache, downscaleFactor, getImage(downscaleFactor), device);
+}
+
+torch::Tensor Camera::getMaskGpu(int downscaleFactor, const torch::Device &device){
+    torch::Tensor m = getMask(downscaleFactor);
+    if (!m.defined() || m.numel() == 0) return m;
+    return gpuCached(gpuMaskCache, downscaleFactor, m, device);
+}
+
+torch::Tensor Camera::getEdgeMapGpu(int downscaleFactor, const torch::Device &device){
+    return gpuCached(gpuEdgeCache, downscaleFactor, getEdgeMap(downscaleFactor).contiguous(), device);
+}
+
+torch::Tensor Camera::getEdgeMap(int downscaleFactor){
+    if (edgePyramids.find(downscaleFactor) != edgePyramids.end()){
+        return edgePyramids[downscaleFactor];
+    }
+    cv::Mat cImg = tensorToImage(getImage(downscaleFactor));
+    cv::Mat gray, edges;
+    cv::cvtColor(cImg, gray, cv::COLOR_RGB2GRAY);
+    cv::Canny(gray, edges, 50, 150);
+    torch::Tensor e = torch::from_blob(edges.data, {edges.rows, edges.cols}, torch::kU8)
+                        .to(torch::kFloat32).div(255.0f).clone();
+    edgePyramids[downscaleFactor] = e;
+    return e;
+}
+
+std::string findMaskPath(const std::string &imagePath, const std::string &projectRoot){
+    static const char *folders[] = { "masks", "mask", "segmentation", "dynamic_masks" };
+    static const char *extensions[] = { ".png", ".jpg", ".jpeg", ".mask.png" };
+
+    fs::path img(imagePath);
+    std::string stem = img.stem().string();
+    std::string name = img.filename().string();
+
+    for (const char *folder : folders){
+        fs::path dir = fs::path(projectRoot) / folder;
+        if (!fs::exists(dir) || !fs::is_directory(dir)) continue;
+        for (const char *ext : extensions){
+            fs::path cand = dir / (stem + ext);
+            if (fs::exists(cand)) return cand.string();
+            cand = dir / (name + ext);
+            if (fs::exists(cand)) return cand.string();
+        }
+    }
+    return "";
 }
 
 std::tuple<std::vector<Camera>, Camera *> InputData::getCameras(bool validate, const std::string &valImage){