Input masks to suppress background/uncertain regions during splat optimization

cv_utils.cpp

@@ -13,6 +13,17 @@ cv::Mat imreadRGB(const std::string &filename){
     return cImg;
 }
 
+cv::Mat imreadMask(const std::string &filename){
+    cv::Mat mask = cv::imread(filename, cv::IMREAD_GRAYSCALE);
+
+    if (mask.empty()){
+        std::cerr << "Cannot read mask " << filename << std::endl;
+        exit(1);
+    }
+
+    return mask;
+}
+
 void imwriteRGB(const std::string &filename, const cv::Mat &image){
     cv::Mat rgb;
     cv::cvtColor(image, rgb, cv::COLOR_RGB2BGR);
@@ -48,3 +59,23 @@ torch::Tensor imageToTensor(const cv::Mat &image){
     return (img.toType(torch::kFloat32) / 255.0f);
 }
 
+torch::Tensor maskToTensor(const cv::Mat &mask){
+    torch::Tensor m = torch::from_blob(mask.data, { mask.rows, mask.cols, 1 }, torch::kU8).clone();
+
+    // Binarize at the midpoint so anti-aliased mask edges (e.g. from a
+    // resize) don't leave the loss weighting on a partial value.
+    return (m.toType(torch::kFloat32) / 255.0f).ge(0.5f).toType(torch::kFloat32);
+}
+
+cv::Mat tensorToMask(const torch::Tensor &t){
+    int h = t.size(0);
+    int w = t.size(1);
+
+    cv::Mat mask(h, w, CV_8UC1);
+    torch::Tensor scaledTensor = (t.squeeze(-1) * 255.0).toType(torch::kU8).contiguous();
+    uint8_t* dataPtr = static_cast<uint8_t*>(scaledTensor.data_ptr());
+    std::copy(dataPtr, dataPtr + (w * h), mask.data);
+
+    return mask;
+}
+

cv_utils.hpp

@@ -7,10 +7,13 @@
 #include <opencv2/imgproc.hpp>
 
 cv::Mat imreadRGB(const std::string &filename);
+cv::Mat imreadMask(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);
+torch::Tensor maskToTensor(const cv::Mat &mask);
+cv::Mat tensorToMask(const torch::Tensor &t);
 
 #endif

input_data.cpp

@@ -116,6 +116,40 @@ torch::Tensor Camera::getImage(int downscaleFactor){
     }
 }
 
+void Camera::loadMask(float downscaleFactor){
+    // Populates mask, resized to match the (already downscaled/undistorted)
+    // image dimensions set by loadImage. Must be called after loadImage.
+    if (maskPath.empty()) return;
+    if (mask.numel()) std::runtime_error("loadMask already called");
+    std::cout << "Loading mask " << maskPath << std::endl;
+
+    cv::Mat cMask = imreadMask(maskPath);
+    if (cMask.rows != height || cMask.cols != width){
+        cv::resize(cMask, cMask, cv::Size(width, height), 0.0, 0.0, cv::INTER_NEAREST);
+    }
+    mask = maskToTensor(cMask);
+}
+
+torch::Tensor Camera::getMask(int downscaleFactor){
+    if (!mask.numel()) return torch::Tensor();
+    if (downscaleFactor <= 1) return mask;
+
+    if (maskPyramids.find(downscaleFactor) != maskPyramids.end()){
+        return maskPyramids[downscaleFactor];
+    }
+
+    // Nearest-neighbor to keep the mask binary (no blended edge values).
+    cv::Mat cMask = tensorToMask(mask);
+    cv::resize(cMask, cMask, cv::Size(cMask.cols / downscaleFactor, cMask.rows / downscaleFactor), 0.0, 0.0, cv::INTER_NEAREST);
+    torch::Tensor t = maskToTensor(cMask);
+    maskPyramids[downscaleFactor] = t;
+    return t;
+}
+
+bool Camera::hasMask() const {
+    return mask.numel() > 0;
+}
+
 bool Camera::hasDistortionParameters(){
     return k1 != 0.0f || k2 != 0.0f || k3 != 0.0f || p1 != 0.0f || p2 != 0.0f;
 }

input_data.hpp

@@ -24,25 +24,31 @@ struct Camera{
     float p2 = 0;
     torch::Tensor camToWorld;
     std::string filePath = "";
+    std::string maskPath = ""; // Optional path to a foreground mask image
     CameraType cameraType = CameraType::Perspective;
 
     Camera(){};
-    Camera(int width, int height, float fx, float fy, float cx, float cy, 
+    Camera(int width, int height, float fx, float fy, float cx, float cy,
         float k1, float k2, float k3, float p1, float p2,
-        const torch::Tensor &camToWorld, const std::string &filePath) : 
-        width(width), height(height), fx(fx), fy(fy), cx(cx), cy(cy), 
+        const torch::Tensor &camToWorld, const std::string &filePath) :
+        width(width), height(height), fx(fx), fy(fy), cx(cx), cy(cy),
         k1(k1), k2(k2), k3(k3), p1(p1), p2(p2),
         camToWorld(camToWorld), filePath(filePath) {}
     torch::Tensor getIntrinsicsMatrix();
     bool hasDistortionParameters();
     std::vector<float> undistortionParameters();
     torch::Tensor getImage(int downscaleFactor);
+    torch::Tensor getMask(int downscaleFactor);
+    bool hasMask() const;
 
     void loadImage(float downscaleFactor);
+    void loadMask(float downscaleFactor);
     torch::Tensor K;
     torch::Tensor image;
+    torch::Tensor mask; // (H, W, 1) in {0, 1}, same size as image once loaded
 
     std::unordered_map<int, torch::Tensor> imagePyramids;
+    std::unordered_map<int, torch::Tensor> maskPyramids;
 };
 
 struct Points{

model.cpp

@@ -777,8 +777,20 @@ int Model::loadPly(const std::string &filename){
     throw std::runtime_error("Invalid PLY file");
 }
 
-torch::Tensor Model::mainLoss(torch::Tensor &rgb, torch::Tensor &gt, float ssimWeight){
-    torch::Tensor ssimLoss = 1.0f - ssim.eval(rgb, gt);
-    torch::Tensor l1Loss = l1(rgb, gt);
+torch::Tensor Model::mainLoss(torch::Tensor &rgb, torch::Tensor &gt, float ssimWeight, const torch::Tensor &mask){
+    torch::Tensor target = gt;
+    if (mask.defined()){
+        // Replace the masked-out (background) pixels of the ground truth with
+        // the same flat backgroundColor the rasterizer composites onto, then
+        // run the ordinary full-frame loss below. Unlike excluding those
+        // pixels from the loss entirely, this actively penalizes any splat
+        // that bleeds past the mask silhouette, since its color no longer
+        // matches the known-flat background there -- which is what keeps
+        // mask contours sharp instead of leaking.
+        torch::Tensor m = (mask.dim() == 3 ? mask.index({"...", 0}) : mask).to(gt.device()).unsqueeze(-1);
+        target = gt * m + backgroundColor.detach() * (1.0f - m);
+    }
+    torch::Tensor ssimLoss = 1.0f - ssim.eval(rgb, target);
+    torch::Tensor l1Loss = l1(rgb, target);
     return (1.0f - ssimWeight) * l1Loss + ssimWeight * ssimLoss;
 }

model.hpp

@@ -74,7 +74,7 @@ struct Model{
   void saveSplat(const std::string &filename);
   void saveDebugPly(const std::string &filename, int step);
   int loadPly(const std::string &filename);
-  torch::Tensor mainLoss(torch::Tensor &rgb, torch::Tensor &gt, float ssimWeight);
+  torch::Tensor mainLoss(torch::Tensor &rgb, torch::Tensor &gt, float ssimWeight, const torch::Tensor &mask = torch::Tensor());
 
   void addToOptimizer(torch::optim::Adam *optimizer, const torch::Tensor &newParam, const torch::Tensor &idcs, int nSamples);
   void removeFromOptimizer(torch::optim::Adam *optimizer, const torch::Tensor &newParam, const torch::Tensor &deletedMask);

opensplat.cpp

@@ -14,6 +14,63 @@
 namespace fs = std::filesystem;
 using namespace torch::indexing;
 
+// Drops points from inputData.points that are never seen projecting into a
+// masked-in (white) region of any camera that has a mask loaded. Points not
+// observed by any masked camera (e.g. behind every frustum) are kept, since
+// no camera actually voted to exclude them. This runs once on the initial
+// sparse cloud, before Model seeds gaussians from it, so gaussians are only
+// ever created where the masks say the subject is.
+static void filterPointsByMasks(InputData &inputData){
+    torch::Tensor xyz = inputData.points.xyz.to(torch::kCPU).contiguous();
+    long long n = xyz.size(0);
+    if (n == 0) return;
+
+    std::vector<uint8_t> seen(n, 0);
+    std::vector<uint8_t> keep(n, 0);
+    auto xyzAcc = xyz.accessor<float, 2>();
+
+    for (Camera &cam : inputData.cameras){
+        if (!cam.hasMask()) continue;
+
+        torch::Tensor R = cam.camToWorld.index({Slice(None, 3), Slice(None, 3)});
+        torch::Tensor T = cam.camToWorld.index({Slice(None, 3), Slice(3, 4)});
+        R = torch::matmul(R, torch::diag(torch::tensor({1.0f, -1.0f, -1.0f})));
+        torch::Tensor Rinv = R.transpose(0, 1).contiguous();
+        torch::Tensor Tinv = torch::matmul(-Rinv, T).contiguous();
+
+        auto RinvAcc = Rinv.accessor<float, 2>();
+        auto TinvAcc = Tinv.accessor<float, 2>();
+
+        torch::Tensor maskT = cam.getMask(1).contiguous();
+        auto maskAcc = maskT.accessor<float, 3>();
+
+        for (long long i = 0; i < n; i++){
+            float px = xyzAcc[i][0], py = xyzAcc[i][1], pz = xyzAcc[i][2];
+            float cx_ = RinvAcc[0][0] * px + RinvAcc[0][1] * py + RinvAcc[0][2] * pz + TinvAcc[0][0];
+            float cy_ = RinvAcc[1][0] * px + RinvAcc[1][1] * py + RinvAcc[1][2] * pz + TinvAcc[1][0];
+            float cz_ = RinvAcc[2][0] * px + RinvAcc[2][1] * py + RinvAcc[2][2] * pz + TinvAcc[2][0];
+            if (cz_ <= 1e-6f) continue;
+
+            int u = static_cast<int>(std::round(cam.fx * cx_ / cz_ + cam.cx));
+            int v = static_cast<int>(std::round(cam.fy * cy_ / cz_ + cam.cy));
+            if (u < 0 || u >= cam.width || v < 0 || v >= cam.height) continue;
+
+            seen[i] = 1;
+            if (maskAcc[v][u][0] >= 0.5f) keep[i] = 1;
+        }
+    }
+
+    std::vector<int64_t> idxs;
+    idxs.reserve(n);
+    for (long long i = 0; i < n; i++){
+        if (keep[i] || !seen[i]) idxs.push_back(i);
+    }
+
+    torch::Tensor idx = torch::tensor(idxs, torch::kInt64);
+    inputData.points.xyz = inputData.points.xyz.index_select(0, idx);
+    inputData.points.rgb = inputData.points.rgb.index_select(0, idx);
+}
+
 int main(int argc, char *argv[]){
     cxxopts::Options options("opensplat", "Open Source 3D Gaussian Splats generator - " APP_VERSION);
     options.add_options()
@@ -42,6 +99,7 @@ int main(int argc, char *argv[]){
         ("stop-screen-size-at", "Stop splitting gaussians that are larger than [split-screen-size] after these many steps", cxxopts::value<int>()->default_value("4000"))
         ("split-screen-size", "Split gaussians that are larger than this percentage of screen space", cxxopts::value<float>()->default_value("0.05"))
         ("colmap-image-path", "Override the default image path for COLMAP-based input", cxxopts::value<std::string>()->default_value(""))
+        ("masks-path", "Path to a directory of foreground mask images (one per input image, matched by filename), used to only fit gaussians within the masked-in (white) regions", cxxopts::value<std::string>()->default_value(""))
 #ifdef USE_VISUALIZATION
         ("has-visualization", "Show the visualization steps of training", cxxopts::value<bool>()->default_value("0"))
 #endif
@@ -95,6 +153,7 @@ int main(int argc, char *argv[]){
     const int stopScreenSizeAt = result["stop-screen-size-at"].as<int>();
     const float splitScreenSize = result["split-screen-size"].as<float>();
     const std::string colmapImageSourcePath = result["colmap-image-path"].as<std::string>();
+    const std::string masksPath = result["masks-path"].as<std::string>();
 #ifdef USE_VISUALIZATION
     const bool hasVisualization = result["has-visualization"].as<bool>();
 #endif
@@ -121,10 +180,42 @@ int main(int argc, char *argv[]){
     try{
         InputData inputData = inputDataFromX(projectRoot, colmapImageSourcePath);
 
+        if (!masksPath.empty()){
+            fs::path masksDir(masksPath);
+            if (!fs::exists(masksDir)){
+                std::cerr << "Masks path does not exist: " << masksPath << std::endl;
+                exit(1);
+            }
+
+            for (Camera &cam : inputData.cameras){
+                fs::path imagePath(cam.filePath);
+                for (const std::string &name : { imagePath.filename().string(), imagePath.stem().string() }){
+                    for (const std::string &ext : { ".png", ".jpg", ".jpeg", ".PNG", ".JPG", ".JPEG" }){
+                        fs::path maskPath = masksDir / (name + ext);
+                        if (fs::exists(maskPath)){
+                            cam.maskPath = maskPath.string();
+                            break;
+                        }
+                    }
+                    if (!cam.maskPath.empty()) break;
+                }
+                if (cam.maskPath.empty()){
+                    std::cerr << "Warning: no mask found for " << cam.filePath << std::endl;
+                }
+            }
+        }
+
         parallel_for(inputData.cameras.begin(), inputData.cameras.end(), [&downScaleFactor](Camera &cam){
             cam.loadImage(downScaleFactor);
+            cam.loadMask(downScaleFactor);
         });
 
+        if (!masksPath.empty()){
+            size_t before = inputData.points.xyz.size(0);
+            filterPointsByMasks(inputData);
+            std::cout << "Masked point filtering: " << before << " -> " << inputData.points.xyz.size(0) << " points" << std::endl;
+        }
+
         // Withhold a validation camera if necessary
         auto t = inputData.getCameras(validate, valImage);
         std::vector<Camera> cams = std::get<0>(t);
@@ -157,7 +248,10 @@ int main(int argc, char *argv[]){
             torch::Tensor gt = cam.getImage(model.getDownscaleFactor(step));
             gt = gt.to(device);
 
-            torch::Tensor mainLoss = model.mainLoss(rgb, gt, ssimWeight);
+            torch::Tensor mask;
+            if (cam.hasMask()) mask = cam.getMask(model.getDownscaleFactor(step)).to(device);
+
+            torch::Tensor mainLoss = model.mainLoss(rgb, gt, ssimWeight, mask);
             mainLoss.backward();
 
             if (step % displayStep == 0) {
@@ -203,7 +297,9 @@ int main(int argc, char *argv[]){
         if (valCam != nullptr){
             torch::Tensor rgb = model.forward(*valCam, numIters);
             torch::Tensor gt = valCam->getImage(model.getDownscaleFactor(numIters)).to(device);
-            std::cout << valCam->filePath << " validation loss: " << model.mainLoss(rgb, gt, ssimWeight).item<float>() << std::endl; 
+            torch::Tensor valMask;
+            if (valCam->hasMask()) valMask = valCam->getMask(model.getDownscaleFactor(numIters)).to(device);
+            std::cout << valCam->filePath << " validation loss: " << model.mainLoss(rgb, gt, ssimWeight, valMask).item<float>() << std::endl;
         }
     }catch(const std::exception &e){
         std::cerr << e.what() << std::endl;

ssim.cpp

@@ -8,7 +8,7 @@ using namespace torch::indexing;
 torch::Tensor SSIM::eval(const torch::Tensor& rendered, const torch::Tensor& gt) {
     torch::Tensor img1 = gt.permute({2, 0, 1}).index({None, "..."});
     torch::Tensor img2 = rendered.permute({2, 0, 1}).index({None, "..."});
-    
+
     if (img1.device() != window.device()){
         window = window.to(img1.device());
     }
@@ -22,7 +22,7 @@ torch::Tensor SSIM::eval(const torch::Tensor& rendered, const torch::Tensor& gt)
     torch::Tensor sigma1Sq = torch::nn::functional::conv2d(img1 * img1, window, torch::nn::functional::Conv2dFuncOptions().padding(windowSize / 2).groups(channel)) - mu1Sq;
     torch::Tensor sigma2Sq = torch::nn::functional::conv2d(img2 * img2, window, torch::nn::functional::Conv2dFuncOptions().padding(windowSize / 2).groups(channel)) - mu2Sq;
     torch::Tensor sigma12 = torch::nn::functional::conv2d(img1 * img2, window, torch::nn::functional::Conv2dFuncOptions().padding(windowSize / 2).groups(channel)) - mu1mu2;
-    
+
     const float C1 = 0.01 * 0.01;
     const float C2 = 0.03 * 0.03;