From d5a88c2255471b345ebaae058b90735826688834 Mon Sep 17 00:00:00 2001 From: Patricio Gonzalez Vivo Date: Tue, 28 Jul 2026 07:08:17 -0400 Subject: [PATCH 1/2] add masking through images --- cv_utils.cpp | 31 +++++++++++++++ cv_utils.hpp | 3 ++ input_data.cpp | 34 +++++++++++++++++ input_data.hpp | 12 ++++-- model.cpp | 16 +++++--- model.hpp | 4 +- opensplat.cpp | 100 ++++++++++++++++++++++++++++++++++++++++++++++++- ssim.cpp | 17 +++++++-- ssim.hpp | 7 +++- 9 files changed, 208 insertions(+), 16 deletions(-) diff --git a/cv_utils.cpp b/cv_utils.cpp index 11b0191d..35d53e4a 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 87c7c2ad..c107ae6d 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 53daebd9..f32204bf 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 a05442be..6fd0d81e 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 a88afa7d..bb9ea1ce 100644 --- a/model.cpp +++ b/model.cpp @@ -51,8 +51,14 @@ torch::Tensor psnr(const torch::Tensor& rendered, const torch::Tensor& gt){ return (10.f * torch::log10(1.0 / mse)); } -torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt){ - return torch::abs(gt - rendered).mean(); +torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt, const torch::Tensor& mask){ + torch::Tensor diff = torch::abs(gt - rendered); + if (mask.defined()){ + torch::Tensor m = (mask.dim() == 3 ? mask.index({"...", 0}) : mask).to(diff.device()); + m = m.unsqueeze(-1).expand_as(diff); + return (diff * m).sum() / m.sum().clamp_min(1.0f); + } + return diff.mean(); } void Model::setupOptimizers(){ @@ -777,8 +783,8 @@ 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 ssimLoss = 1.0f - ssim.eval(rgb, gt, mask); + torch::Tensor l1Loss = l1(rgb, gt, mask); return (1.0f - ssimWeight) * l1Loss + ssimWeight * ssimLoss; } diff --git a/model.hpp b/model.hpp index fbf3edcc..f7c50748 100644 --- a/model.hpp +++ b/model.hpp @@ -17,7 +17,7 @@ using namespace torch::autograd; torch::Tensor randomQuatTensor(long long n); torch::Tensor projectionMatrix(float zNear, float zFar, float fovX, float fovY, const torch::Device &device); torch::Tensor psnr(const torch::Tensor& rendered, const torch::Tensor& gt); -torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt); +torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt, const torch::Tensor& mask = torch::Tensor()); struct Model{ Model(const InputData &inputData, int numCameras, @@ -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 b2826d13..212466d0 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 f8a152e6..906b0bea 100644 --- a/ssim.cpp +++ b/ssim.cpp @@ -5,10 +5,10 @@ using namespace torch::indexing; -torch::Tensor SSIM::eval(const torch::Tensor& rendered, const torch::Tensor& gt) { +torch::Tensor SSIM::eval(const torch::Tensor& rendered, const torch::Tensor& gt, const torch::Tensor& mask) { 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,12 +22,23 @@ 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; torch::Tensor ssimMap = ((2.0f * mu1mu2 + C1) * (2.0f * sigma12 + C2)) / ((mu1Sq + mu2Sq + C1) * (sigma1Sq + sigma2Sq + C2)); + if (mask.defined()){ + // mask is (H, W) or (H, W, 1) in [0, 1] -- broadcast against + // ssimMap's (1, C, H, W) and take a masked mean instead of a plain + // one, so masked-out (e.g. low VGGT-confidence) regions don't pull + // the score down/up regardless of how well they happen to match. + torch::Tensor m = mask.dim() == 3 ? mask.index({"...", 0}) : mask; + m = m.to(ssimMap.device()).index({None, None, "...", "..."}); + torch::Tensor denom = m.sum().clamp_min(1.0f) * ssimMap.size(1); + return (ssimMap * m).sum() / denom; + } + return ssimMap.mean(); } diff --git a/ssim.hpp b/ssim.hpp index b74a86df..5e5ac80a 100644 --- a/ssim.hpp +++ b/ssim.hpp @@ -12,7 +12,12 @@ class SSIM{ window = createWindow(); }; - torch::Tensor eval(const torch::Tensor& rendered, const torch::Tensor& gt); + // mask: optional (H, W) or (H, W, 1) tensor in [0, 1], same spatial size + // as rendered/gt. Pixels outside the masked-in (white) region don't + // contribute to the returned score -- an undefined (default-constructed) + // mask behaves exactly as before (every pixel counted). + torch::Tensor eval(const torch::Tensor& rendered, const torch::Tensor& gt, + const torch::Tensor& mask = torch::Tensor()); private: torch::Tensor createWindow(); torch::Tensor gaussian(float sigma); From 96565e4535994a3cf6e6204072a3cf3b2fbce966 Mon Sep 17 00:00:00 2001 From: Patricio Gonzalez Vivo Date: Tue, 28 Jul 2026 07:08:18 -0400 Subject: [PATCH 2/2] fixing masked edges leak --- model.cpp | 26 ++++++++++++++++---------- model.hpp | 2 +- ssim.cpp | 13 +------------ ssim.hpp | 7 +------ 4 files changed, 19 insertions(+), 29 deletions(-) diff --git a/model.cpp b/model.cpp index bb9ea1ce..f4c7b9f8 100644 --- a/model.cpp +++ b/model.cpp @@ -51,14 +51,8 @@ torch::Tensor psnr(const torch::Tensor& rendered, const torch::Tensor& gt){ return (10.f * torch::log10(1.0 / mse)); } -torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt, const torch::Tensor& mask){ - torch::Tensor diff = torch::abs(gt - rendered); - if (mask.defined()){ - torch::Tensor m = (mask.dim() == 3 ? mask.index({"...", 0}) : mask).to(diff.device()); - m = m.unsqueeze(-1).expand_as(diff); - return (diff * m).sum() / m.sum().clamp_min(1.0f); - } - return diff.mean(); +torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt){ + return torch::abs(gt - rendered).mean(); } void Model::setupOptimizers(){ @@ -784,7 +778,19 @@ int Model::loadPly(const std::string &filename){ } torch::Tensor Model::mainLoss(torch::Tensor &rgb, torch::Tensor >, float ssimWeight, const torch::Tensor &mask){ - torch::Tensor ssimLoss = 1.0f - ssim.eval(rgb, gt, mask); - torch::Tensor l1Loss = l1(rgb, gt, 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 f7c50748..64ed25f1 100644 --- a/model.hpp +++ b/model.hpp @@ -17,7 +17,7 @@ using namespace torch::autograd; torch::Tensor randomQuatTensor(long long n); torch::Tensor projectionMatrix(float zNear, float zFar, float fovX, float fovY, const torch::Device &device); torch::Tensor psnr(const torch::Tensor& rendered, const torch::Tensor& gt); -torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt, const torch::Tensor& mask = torch::Tensor()); +torch::Tensor l1(const torch::Tensor& rendered, const torch::Tensor& gt); struct Model{ Model(const InputData &inputData, int numCameras, diff --git a/ssim.cpp b/ssim.cpp index 906b0bea..a5dd0295 100644 --- a/ssim.cpp +++ b/ssim.cpp @@ -5,7 +5,7 @@ using namespace torch::indexing; -torch::Tensor SSIM::eval(const torch::Tensor& rendered, const torch::Tensor& gt, const torch::Tensor& mask) { +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, "..."}); @@ -28,17 +28,6 @@ torch::Tensor SSIM::eval(const torch::Tensor& rendered, const torch::Tensor& gt, torch::Tensor ssimMap = ((2.0f * mu1mu2 + C1) * (2.0f * sigma12 + C2)) / ((mu1Sq + mu2Sq + C1) * (sigma1Sq + sigma2Sq + C2)); - if (mask.defined()){ - // mask is (H, W) or (H, W, 1) in [0, 1] -- broadcast against - // ssimMap's (1, C, H, W) and take a masked mean instead of a plain - // one, so masked-out (e.g. low VGGT-confidence) regions don't pull - // the score down/up regardless of how well they happen to match. - torch::Tensor m = mask.dim() == 3 ? mask.index({"...", 0}) : mask; - m = m.to(ssimMap.device()).index({None, None, "...", "..."}); - torch::Tensor denom = m.sum().clamp_min(1.0f) * ssimMap.size(1); - return (ssimMap * m).sum() / denom; - } - return ssimMap.mean(); } diff --git a/ssim.hpp b/ssim.hpp index 5e5ac80a..b74a86df 100644 --- a/ssim.hpp +++ b/ssim.hpp @@ -12,12 +12,7 @@ class SSIM{ window = createWindow(); }; - // mask: optional (H, W) or (H, W, 1) tensor in [0, 1], same spatial size - // as rendered/gt. Pixels outside the masked-in (white) region don't - // contribute to the returned score -- an undefined (default-constructed) - // mask behaves exactly as before (every pixel counted). - torch::Tensor eval(const torch::Tensor& rendered, const torch::Tensor& gt, - const torch::Tensor& mask = torch::Tensor()); + torch::Tensor eval(const torch::Tensor& rendered, const torch::Tensor& gt); private: torch::Tensor createWindow(); torch::Tensor gaussian(float sigma);