diff --git a/cv_utils.cpp b/cv_utils.cpp index 11b0191..35d53e4 100644 --- a/cv_utils.cpp +++ b/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(scaledTensor.data_ptr()); + std::copy(dataPtr, dataPtr + (w * h), mask.data); + + return mask; +} + diff --git a/cv_utils.hpp b/cv_utils.hpp index 87c7c2a..c107ae6 100644 --- a/cv_utils.hpp +++ b/cv_utils.hpp @@ -7,10 +7,13 @@ #include 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 \ No newline at end of file diff --git a/input_data.cpp b/input_data.cpp index 53daebd..f32204b 100644 --- a/input_data.cpp +++ b/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; } diff --git a/input_data.hpp b/input_data.hpp index a05442b..6fd0d81 100644 --- a/input_data.hpp +++ b/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 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 imagePyramids; + std::unordered_map maskPyramids; }; struct Points{ diff --git a/model.cpp b/model.cpp index a88afa7..f4c7b9f 100644 --- a/model.cpp +++ b/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 >, 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 >, 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; } diff --git a/model.hpp b/model.hpp index fbf3edc..64ed25f 100644 --- a/model.hpp +++ b/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 >, float ssimWeight); + torch::Tensor mainLoss(torch::Tensor &rgb, torch::Tensor >, 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); diff --git a/opensplat.cpp b/opensplat.cpp index b2826d1..212466d 100644 --- a/opensplat.cpp +++ b/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 seen(n, 0); + std::vector keep(n, 0); + auto xyzAcc = xyz.accessor(); + + 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(); + auto TinvAcc = Tinv.accessor(); + + torch::Tensor maskT = cam.getMask(1).contiguous(); + auto maskAcc = maskT.accessor(); + + 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(std::round(cam.fx * cx_ / cz_ + cam.cx)); + int v = static_cast(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 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()->default_value("4000")) ("split-screen-size", "Split gaussians that are larger than this percentage of screen space", cxxopts::value()->default_value("0.05")) ("colmap-image-path", "Override the default image path for COLMAP-based input", cxxopts::value()->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()->default_value("")) #ifdef USE_VISUALIZATION ("has-visualization", "Show the visualization steps of training", cxxopts::value()->default_value("0")) #endif @@ -95,6 +153,7 @@ int main(int argc, char *argv[]){ const int stopScreenSizeAt = result["stop-screen-size-at"].as(); const float splitScreenSize = result["split-screen-size"].as(); const std::string colmapImageSourcePath = result["colmap-image-path"].as(); + const std::string masksPath = result["masks-path"].as(); #ifdef USE_VISUALIZATION const bool hasVisualization = result["has-visualization"].as(); #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 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() << 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() << std::endl; } }catch(const std::exception &e){ std::cerr << e.what() << std::endl; diff --git a/ssim.cpp b/ssim.cpp index f8a152e..a5dd029 100644 --- a/ssim.cpp +++ b/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;