From df47a765100265b580cc21ccc77d373f1fd3cb90 Mon Sep 17 00:00:00 2001 From: Piero Toffanin Date: Wed, 16 Sep 2026 12:32:13 -0400 Subject: [PATCH 1/2] Fix tile-bin overflow crash and stream images via a compressed mem cache --- CMakeLists.txt | 4 +- cv_utils.cpp | 13 +- cv_utils.hpp | 6 +- image_pipeline.cpp | 275 ++++++++++++++++++++ image_pipeline.hpp | 68 +++++ image_store.cpp | 276 +++++++++++++++++++++ image_store.hpp | 64 +++++ input_data.cpp | 205 +-------------- input_data.hpp | 22 +- model.cpp | 16 +- model.hpp | 4 +- opensplat.cpp | 54 ++-- rasterize_gaussians.cpp | 2 +- rasterizer/gsplat-cpu/bindings.h | 5 +- rasterizer/gsplat-cpu/gsplat_cpu.cpp | 58 ++++- rasterizer/gsplat-metal/bindings.h | 9 +- rasterizer/gsplat-metal/gsplat_metal.metal | 29 ++- rasterizer/gsplat-metal/gsplat_metal.mm | 28 ++- rasterizer/gsplat/bindings.cu | 90 +++++-- rasterizer/gsplat/bindings.h | 9 +- ssim.hpp | 3 +- sysinfo.cpp | 109 ++++++++ sysinfo.hpp | 13 + utils.hpp | 33 ++- visualizer.cpp | 4 +- zip_utils.cpp | 4 +- zip_utils.hpp | 4 +- 27 files changed, 1087 insertions(+), 320 deletions(-) create mode 100644 image_pipeline.cpp create mode 100644 image_pipeline.hpp create mode 100644 image_store.cpp create mode 100644 image_store.hpp create mode 100644 sysinfo.cpp create mode 100644 sysinfo.hpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 479d97df..f7ade660 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -213,7 +213,7 @@ endif() if (NOT WIN32 AND NOT APPLE) set(CMAKE_CUDA_COMPILER "${CUDA_TOOLKIT_ROOT_DIR}/bin/nvcc") endif() -set(OpenCV_LIBS opencv_core opencv_imgproc opencv_highgui) +set(OpenCV_LIBS opencv_core opencv_imgproc opencv_imgcodecs opencv_highgui) set(GSPLAT_LIBS gsplat_cpu) if((GPU_RUNTIME STREQUAL "CUDA") OR (GPU_RUNTIME STREQUAL "HIP")) @@ -265,7 +265,7 @@ 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 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) +tensor_math.cpp rad.cpp zip_utils.cpp undistort.cpp sysinfo.cpp image_store.cpp image_pipeline.cpp) if (OPENSPLAT_BUILD_VISUALIZER) if (Pangolin_FOUND) diff --git a/cv_utils.cpp b/cv_utils.cpp index 5a7f80b6..43b9387a 100644 --- a/cv_utils.cpp +++ b/cv_utils.cpp @@ -22,15 +22,20 @@ cv::Mat tensorToImage(const torch::Tensor &t){ if (c != 3) throw std::runtime_error("Only images with 3 channels are supported"); cv::Mat image(h, w, type); - torch::Tensor scaledTensor = (t * 255.0).toType(torch::kU8); - uint8_t* dataPtr = static_cast(scaledTensor.data_ptr()); + torch::Tensor u8 = t.scalar_type() == torch::kU8 + ? t.contiguous() + : (t * 255.0).toType(torch::kU8); + uint8_t* dataPtr = static_cast(u8.data_ptr()); std::copy(dataPtr, dataPtr + (w * h * c), image.data); return image; } torch::Tensor imageToTensor(const cv::Mat &image){ - torch::Tensor img = torch::from_blob(image.data, { image.rows, image.cols, image.dims + 1 }, torch::kU8); - return (img.toType(torch::kFloat32) / 255.0f); + return torch::from_blob(image.data, { image.rows, image.cols, image.dims + 1 }, torch::kU8).clone(); } +torch::Tensor toUnitFloat(const torch::Tensor &img){ + if (img.scalar_type() == torch::kU8) return img.to(torch::kFloat32).mul_(1.0f / 255.0f); + return img; +} diff --git a/cv_utils.hpp b/cv_utils.hpp index 67fb95d1..e219c70b 100644 --- a/cv_utils.hpp +++ b/cv_utils.hpp @@ -7,7 +7,11 @@ #include cv::Mat imreadRGB(const std::string &filename); +// Accepts uint8 [H,W,3] in [0,255] or float [H,W,3] in [0,1] cv::Mat tensorToImage(const torch::Tensor &t); +// [H,W,3] uint8 tensor owning a copy of the image torch::Tensor imageToTensor(const cv::Mat &image); +// float [0,1] view of an image tensor (no-op for float input) +torch::Tensor toUnitFloat(const torch::Tensor &img); -#endif \ No newline at end of file +#endif diff --git a/image_pipeline.cpp b/image_pipeline.cpp new file mode 100644 index 00000000..47228ddf --- /dev/null +++ b/image_pipeline.cpp @@ -0,0 +1,275 @@ +#include +#include +#include +#include +#include +#include +#include +#include "image_pipeline.hpp" +#include "sysinfo.hpp" + +namespace { + +const int MAX_SLOTS = 100000; +const int MIN_PREFETCH = 2; +const int MAX_PREFETCH = 12; +const int RETUNE_EVERY = 64; + +double seconds(std::chrono::steady_clock::time_point a, std::chrono::steady_clock::time_point b){ + return std::chrono::duration(b - a).count(); +} + +} + +struct ImagePipeline::Slot{ + enum State { FREE, DECODING, READY }; + State state = FREE; + int imageId = -1; + int level = 0; + int leases = 0; + uint64_t lastUse = 0; + torch::Tensor staging; // host uint8 [H,W,3] + torch::Tensor maskStaging; // host uint8 [H,W] + std::shared_ptr frame; +}; + +ImagePipeline::ImagePipeline(ImageStore &store, const torch::Device &device, int numWorkers) + : store(store), device(device){ + target = MIN_PREFETCH; + capacity = target.load() + 2; + numWorkers = (std::max)(1, numWorkers); + for (int i = 0; i < numWorkers; i++){ + workers.emplace_back([this](){ workerLoop(); }); + } +} + +ImagePipeline::~ImagePipeline(){ + { + std::lock_guard lock(m); + stopping = true; + } + queueCv.notify_all(); + for (std::thread &t : workers) t.join(); +} + +void ImagePipeline::request(const Camera &cam, int level){ + std::lock_guard lock(m); + for (const auto &s : slots){ + if (s->state != Slot::FREE && s->imageId == cam.imageId && s->level == level) return; + } + for (const Request &r : queue){ + if (r.imageId == cam.imageId && r.level == level) return; + } + queue.push_back({ cam.imageId, level, cam.width, cam.height, cam.hasMask }); + queueCv.notify_one(); +} + +// Picks a slot for a new decode: a free one, a new one while under capacity, +// or the least recently used frame that nobody holds +ImagePipeline::Slot *ImagePipeline::reserveSlotLocked(const Request &r){ + for (auto &s : slots){ + if (s->state == Slot::FREE) return s.get(); + } + if (static_cast(slots.size()) < capacity){ + slots.push_back(std::make_unique()); + return slots.back().get(); + } + Slot *victim = nullptr; + for (auto &s : slots){ + if (s->state == Slot::READY && s->leases == 0 && (!victim || s->lastUse < victim->lastUse)) victim = s.get(); + } + return victim; +} + +void ImagePipeline::workerLoop(){ + while (true){ + Request r; + Slot *slot = nullptr; + { + std::unique_lock lock(m); + queueCv.wait(lock, [&](){ return stopping || !queue.empty(); }); + if (stopping) return; + r = queue.front(); + queue.pop_front(); + + bool resident = false; + for (auto &s : slots){ + if (s->state != Slot::FREE && s->imageId == r.imageId && s->level == r.level) resident = true; + } + if (resident) continue; + + while (!(slot = reserveSlotLocked(r))){ + queueCv.wait(lock); + if (stopping) return; + } + slot->state = Slot::DECODING; + slot->imageId = r.imageId; + slot->level = r.level; + slot->frame.reset(); + } + + auto t0 = std::chrono::steady_clock::now(); + std::string error; + try{ + decodeInto(*slot, r); + }catch (const std::exception &e){ + error = e.what(); + } + double dt = seconds(t0, std::chrono::steady_clock::now()); + + { + std::lock_guard lock(m); + if (!error.empty()){ + std::cerr << "Image decode failed: " << error << std::endl; + slot->state = Slot::FREE; + slot->imageId = -1; + }else{ + slot->state = Slot::READY; + slot->lastUse = ++useCounter; + decodeEma = decodeEma == 0.0 ? dt : 0.9 * decodeEma + 0.1 * dt; + } + } + readyCv.notify_all(); + queueCv.notify_all(); + } +} + +void ImagePipeline::decodeInto(Slot &slot, const Request &r){ + Blob bytes = store.cache.get(ImageStore::imageKey(r.imageId, r.level)); + cv::Mat bgr = cv::imdecode(cv::Mat(1, static_cast(bytes->size()), CV_8UC1, const_cast(bytes->data())), + cv::IMREAD_COLOR); + if (bgr.empty()) throw std::runtime_error("Cannot decode image " + std::to_string(r.imageId)); + const int H = bgr.rows; + const int W = bgr.cols; + + const bool pinned = device.is_cuda(); + auto hostOpts = torch::TensorOptions().dtype(torch::kU8).pinned_memory(pinned); + if (!slot.staging.defined() || slot.staging.size(0) != H || slot.staging.size(1) != W){ + slot.staging = torch::empty({H, W, 3}, hostOpts); + } + cv::Mat rgb(H, W, CV_8UC3, slot.staging.data_ptr()); + cv::cvtColor(bgr, rgb, cv::COLOR_BGR2RGB); + + auto frame = std::make_shared(); + frame->imageId = r.imageId; + frame->level = r.level; + frame->image = device.is_cpu() ? slot.staging.clone() : slot.staging.to(device); + + if (r.hasMask){ + Blob maskBytes = store.cache.get(ImageStore::maskKey(r.imageId, r.level)); + cv::Mat mask = cv::imdecode(cv::Mat(1, static_cast(maskBytes->size()), CV_8UC1, const_cast(maskBytes->data())), + cv::IMREAD_GRAYSCALE); + if (mask.empty() || mask.rows != H || mask.cols != W){ + throw std::runtime_error("Cannot decode mask for image " + std::to_string(r.imageId)); + } + if (!slot.maskStaging.defined() || slot.maskStaging.size(0) != H || slot.maskStaging.size(1) != W){ + slot.maskStaging = torch::empty({H, W}, hostOpts); + } + std::memcpy(slot.maskStaging.data_ptr(), mask.data, static_cast(H) * W); + frame->mask = slot.maskStaging.to(device).to(torch::kFloat32).div_(255.0f); + } + + uint64_t slotBytes = static_cast(H) * W * (3 + (r.hasMask ? 4 : 0)); + slot.frame = frame; + { + std::lock_guard lock(m); + maxSlotBytes = (std::max)(maxSlotBytes, slotBytes); + } +} + +FrameLease ImagePipeline::acquire(const Camera &cam, int level, bool wantEdges){ + std::unique_lock lock(m); + auto t0 = std::chrono::steady_clock::now(); + bool waited = false; + Slot *slot = nullptr; + while (true){ + slot = nullptr; + bool decoding = false; + for (auto &s : slots){ + if (s->state == Slot::FREE || s->imageId != cam.imageId || s->level != level) continue; + if (s->state == Slot::READY){ slot = s.get(); break; } + decoding = true; + } + if (slot) break; + + if (!decoding){ + bool queued = false; + for (const Request &r : queue){ + if (r.imageId == cam.imageId && r.level == level){ queued = true; break; } + } + if (!queued) queue.push_front({ cam.imageId, level, cam.width, cam.height, cam.hasMask }); + queueCv.notify_all(); + } + waited = true; + readyCv.wait(lock); + } + + if (waited){ + waitAccum += seconds(t0, std::chrono::steady_clock::now()); + windowWaits++; + } + slot->leases++; + slot->lastUse = ++useCounter; + std::shared_ptr frame = slot->frame; + lock.unlock(); + + if (wantEdges && !frame->edges.defined()){ + torch::Tensor host = frame->image.to(torch::kCPU).contiguous(); + cv::Mat rgb(host.size(0), host.size(1), CV_8UC3, host.data_ptr()); + cv::Mat gray, edges; + cv::cvtColor(rgb, gray, cv::COLOR_RGB2GRAY); + cv::Canny(gray, edges, 50, 150); + frame->edges = torch::from_blob(edges.data, {edges.rows, edges.cols}, torch::kU8) + .to(device).to(torch::kFloat32).div_(255.0f); + } + + return FrameLease(frame.get(), [this, slot](const Frame *){ + { + std::lock_guard lock(m); + slot->leases--; + } + queueCv.notify_all(); + }); +} + +void ImagePipeline::noteTrainStep(double sec){ + std::lock_guard lock(m); + double compute = (std::max)(sec - waitAccum, 1e-6); + waitAccum = 0.0; + trainEma = trainEma == 0.0 ? compute : 0.9 * trainEma + 0.1 * compute; + stepsTotal++; + // Adapt quickly at the start, then settle into longer windows + const int retuneEvery = stepsTotal < 256 ? 16 : RETUNE_EVERY; + if (++stepsSinceRetune >= retuneEvery){ + stepsSinceRetune = 0; + retuneLocked(windowWaits > 0); + windowWaits = 0; + } +} + +// Prefetch depth follows the decode/train latency ratio: it grows whenever the +// loop had to wait for a frame and shrinks only after several quiet windows +void ImagePipeline::retuneLocked(bool waited){ + int base = MIN_PREFETCH; + if (trainEma > 0.0 && decodeEma > 0.0){ + base = std::clamp(static_cast(std::ceil(decodeEma / trainEma)) + 2, MIN_PREFETCH, MAX_PREFETCH); + } + int t = target.load(); + if (waited){ + t = (std::max)(base, t + 1); + quietWindows = 0; + }else if (++quietWindows >= 4){ + t = base; + quietWindows = 0; + }else{ + t = (std::max)(base, t); + } + t = std::clamp(t, MIN_PREFETCH, MAX_PREFETCH); + target = t; + + // Frames stay resident while memory allows: the pool holds at least the + // prefetch window and grows up to half of the device memory free at start + if (frameBudget == 0) frameBudget = freeDeviceMemoryBytes(device.is_cuda()) / 2; + int maxSlots = maxSlotBytes > 0 ? static_cast((std::min)(frameBudget / maxSlotBytes, MAX_SLOTS)) : 0; + capacity = (std::max)(t + 2, maxSlots); +} diff --git a/image_pipeline.hpp b/image_pipeline.hpp new file mode 100644 index 00000000..df273827 --- /dev/null +++ b/image_pipeline.hpp @@ -0,0 +1,68 @@ +#ifndef IMAGE_PIPELINE_H +#define IMAGE_PIPELINE_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include "image_store.hpp" + +// A decoded camera image resident on the training device +struct Frame{ + torch::Tensor image; // uint8 [H,W,3] + torch::Tensor mask; // float [H,W] in {0,1}, or undefined + torch::Tensor edges; // float [H,W] in {0,1}, or undefined + int imageId = -1; + int level = 1; +}; +// Keeps the frame buffer reserved while the frame is in use +using FrameLease = std::shared_ptr; + +class ImagePipeline{ +public: + ImagePipeline(ImageStore &store, const torch::Device &device, int numWorkers); + ~ImagePipeline(); + + void request(const Camera &cam, int level); + FrameLease acquire(const Camera &cam, int level, bool wantEdges = false); + void noteTrainStep(double seconds); + int prefetchTarget() const { return target.load(); } + +private: + struct Slot; + struct Request{ int imageId; int level; int width; int height; bool hasMask; }; + + void workerLoop(); + Slot *reserveSlotLocked(const Request &r); + void decodeInto(Slot &slot, const Request &r); + void retuneLocked(bool waited); + + ImageStore &store; + torch::Device device; + std::vector> slots; + std::deque queue; + std::vector workers; + std::mutex m; + std::condition_variable queueCv; + std::condition_variable readyCv; + bool stopping = false; + + std::atomic target{2}; + int capacity = 4; + double decodeEma = 0.0; + double trainEma = 0.0; // step time excluding time spent waiting for frames + double waitAccum = 0.0; // wait time since the last noteTrainStep + int windowWaits = 0; // acquire() calls that had to wait in this window + int quietWindows = 0; + int stepsTotal = 0; + int stepsSinceRetune = 0; + uint64_t useCounter = 0; + uint64_t maxSlotBytes = 0; + uint64_t frameBudget = 0; // device bytes the frame pool may hold, sampled at the first retune +}; + +#endif diff --git a/image_store.cpp b/image_store.cpp new file mode 100644 index 00000000..98080f03 --- /dev/null +++ b/image_store.cpp @@ -0,0 +1,276 @@ +#include +#include +#include +#include +#include +#include +#include "image_store.hpp" +#include "sysinfo.hpp" +#include "undistort.hpp" +#include "utils.hpp" + +namespace fs = std::filesystem; + +static const std::string savePrefix = "opensplat-cache-"; + +BlobCache::BlobCache(uint64_t capBytes, const fs::path &baseDir) : cap(capBytes){ + // Remove save directories left behind by crashed runs + std::error_code ec; + for (const auto &entry : fs::directory_iterator(baseDir, ec)){ + std::string name = entry.path().filename().string(); + if (name.rfind(savePrefix, 0) != 0) continue; + int pid = std::atoi(name.substr(savePrefix.size()).c_str()); + if (pid > 0 && pid != currentPid() && !processAlive(pid)){ + fs::remove_all(entry.path(), ec); + } + } + + saveDir = baseDir / (savePrefix + std::to_string(currentPid())); +} + +BlobCache::~BlobCache(){ + std::error_code ec; + fs::remove_all(saveDir, ec); +} + +fs::path BlobCache::savePath(const std::string &key) const{ + std::string name = key; + std::replace(name.begin(), name.end(), ':', '_'); + return saveDir / (name + ".bin"); +} + +void BlobCache::touchLocked(Entry &e, const std::string &key){ + lru.erase(e.lru); + lru.push_front(key); + e.lru = lru.begin(); +} + +// Saves least recently used blobs to disk until resident bytes fit the cap +void BlobCache::evictLocked(const std::string &keep){ + if (resident <= cap) return; + + if (!lowRamGuardTripped && availableRamBytes() < physicalRamBytes() / 10){ + cap = (std::max)(cap / 2, static_cast(64) << 20); + lowRamGuardTripped = true; + } + + auto it = lru.end(); + while (resident > cap && it != lru.begin()){ + --it; + if (*it == keep) continue; + Entry &e = map[*it]; + if (!e.bytes) continue; + if (!e.saved){ + std::error_code ec; + fs::create_directories(saveDir, ec); + std::ofstream f(savePath(*it), std::ios::binary); + f.write(reinterpret_cast(e.bytes->data()), e.bytes->size()); + if (!f) throw std::runtime_error("Cannot write image cache file in " + saveDir.string()); + e.saved = true; + } + resident -= e.size; + e.bytes.reset(); + } +} + +void BlobCache::put(const std::string &key, std::vector &&bytes){ + std::lock_guard lock(m); + auto it = map.find(key); + if (it != map.end()){ + if (it->second.bytes) resident -= it->second.size; + lru.erase(it->second.lru); + map.erase(it); + } + Entry e; + e.size = bytes.size(); + e.bytes = std::make_shared>(std::move(bytes)); + lru.push_front(key); + e.lru = lru.begin(); + resident += e.size; + map[key] = e; + evictLocked(key); +} + +Blob BlobCache::get(const std::string &key){ + std::lock_guard lock(m); + auto it = map.find(key); + if (it == map.end()) throw std::runtime_error("Image blob not found: " + key); + Entry &e = it->second; + if (!e.bytes){ + std::ifstream f(savePath(key), std::ios::binary); + if (!f) throw std::runtime_error("Cannot reload saved image " + key); + auto data = std::make_shared>(e.size); + f.read(reinterpret_cast(data->data()), e.size); + e.bytes = data; + resident += e.size; + } + touchLocked(e, key); + Blob b = e.bytes; + evictLocked(key); + return b; +} + +bool BlobCache::has(const std::string &key) const{ + std::lock_guard lock(m); + return map.find(key) != map.end(); +} + +ImageStore::ImageStore(uint64_t capBytes, const fs::path &baseDir) : cache(capBytes, baseDir){} + +std::string ImageStore::imageKey(int imageId, int level){ + return "i:" + std::to_string(imageId) + ":" + std::to_string(level); +} + +std::string ImageStore::maskKey(int imageId, int level){ + return "m:" + std::to_string(imageId) + ":" + std::to_string(level); +} + +static bool isJpegOrPng(const std::string &path){ + std::string ext = fs::path(path).extension().string(); + std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c){ return std::tolower(c); }); + return ext == ".jpg" || ext == ".jpeg" || ext == ".png"; +} + +static std::vector readFileBytes(const std::string &path){ + std::ifstream f(path, std::ios::binary | std::ios::ate); + std::vector bytes(static_cast(f.tellg())); + f.seekg(0); + f.read(reinterpret_cast(bytes.data()), bytes.size()); + return bytes; +} + +static std::vector encodeJpeg(const cv::Mat &bgr){ + std::vector out; + cv::imencode(".jpg", bgr, out, { cv::IMWRITE_JPEG_QUALITY, 95 }); + return out; +} + +static std::vector encodePng(const cv::Mat &gray){ + std::vector out; + cv::imencode(".png", gray, out, { cv::IMWRITE_PNG_COMPRESSION, 1 }); + return out; +} + +void ImageStore::prepare(std::vector &cameras, float downscaleFactor, int numDownscales){ + numLevels = numDownscales + 1; + for (size_t i = 0; i < cameras.size(); i++) cameras[i].imageId = static_cast(i); + + // Bound the number of images decoded at once by the available RAM + uint64_t maxPixels = 1; + for (const Camera &cam : cameras){ + maxPixels = (std::max)(maxPixels, static_cast(cam.width) * cam.height); + } + const size_t chunk = static_cast(std::clamp(availableRamBytes() / 2 / (maxPixels * 17), 1, 32)); + + const size_t total = cameras.size(); + std::atomic done{0}; + std::mutex logMutex; + std::string firstError; + + for (size_t start = 0; start < total; start += chunk){ + size_t end = (std::min)(total, start + chunk); + parallel_for(cameras.begin() + start, cameras.begin() + end, [&](Camera &cam){ + try{ + prepareCamera(cam, downscaleFactor); + }catch (const std::exception &e){ + std::lock_guard lock(logMutex); + if (firstError.empty()) firstError = e.what(); + return; + } + size_t n = ++done; + if (n % 100 == 0 || n == total){ + std::lock_guard lock(logMutex); + std::cout << "Preprocessing images " << n << "/" << total << std::endl; + } + }); + if (!firstError.empty()) throw std::runtime_error(firstError); + } +} + +// Decodes, rescales and undistorts one camera image, then stores a compressed +// copy per downscale level and updates the camera intrinsics +void ImageStore::prepareCamera(Camera &cam, float downscaleFactor){ + cv::Mat img = cv::imread(cam.filePath); + if (img.empty()){ + throw std::runtime_error("Cannot read " + cam.filePath + + "\nMake sure the path to your images is correct"); + } + + cv::Mat mask; + if (!cam.maskPath.empty()){ + mask = cv::imread(cam.maskPath, cv::IMREAD_GRAYSCALE); + if (mask.empty()) throw std::runtime_error("Cannot read mask " + cam.maskPath); + } + + // If camera intrinsics don't match the image dimensions + if (img.rows != cam.height || img.cols != cam.width){ + float rescaleF = static_cast(img.rows) / static_cast(cam.height); + cam.fx *= rescaleF; + cam.fy *= rescaleF; + cam.cx *= rescaleF; + cam.cy *= rescaleF; + } + + bool pixelsUntouched = true; + if (downscaleFactor > 1.0f){ + float scaleFactor = 1.0f / downscaleFactor; + cv::resize(img, img, cv::Size(), scaleFactor, scaleFactor, cv::INTER_AREA); + cam.fx *= scaleFactor; + cam.fy *= scaleFactor; + cam.cx *= scaleFactor; + cam.cy *= scaleFactor; + pixelsUntouched = false; + } + + if (!mask.empty()){ + cv::threshold(mask, mask, 127, 255, cv::THRESH_BINARY); + if (mask.rows != img.rows || mask.cols != img.cols){ + cv::resize(mask, mask, cv::Size(img.cols, img.rows), 0.0, 0.0, cv::INTER_LINEAR); + } + } + + if (cam.hasDistortionParameters()){ + UndistortParams p = computeUndistortParams(cam.fx, cam.fy, cam.cx, cam.cy, img.cols, img.rows, + cam.k1, cam.k2, cam.k3, cam.k4, cam.k5, cam.k6, cam.p1, cam.p2); + cv::Mat mapx, mapy; + buildUndistortMaps(p, mapx, mapy); + cv::Mat undistorted; + cv::remap(img, undistorted, mapx, mapy, cv::INTER_LINEAR, cv::BORDER_CONSTANT); + img = undistorted; + if (!mask.empty()){ + cv::Mat remapped; + cv::remap(mask, remapped, mapx, mapy, cv::INTER_LINEAR, cv::BORDER_CONSTANT); + mask = remapped; + } + cam.fx = p.dstFx; + cam.fy = p.dstFy; + cam.cx = p.dstCx; + cam.cy = p.dstCy; + pixelsUntouched = false; + } + + cam.width = img.cols; + cam.height = img.rows; + cam.K = cam.getIntrinsicsMatrix(); + cam.hasMask = !mask.empty(); + + for (int l = 0; l < numLevels; l++){ + const int level = 1 << l; + cv::Mat lvlImg = img, lvlMask = mask; + if (level > 1){ + cv::resize(img, lvlImg, cv::Size(img.cols / level, img.rows / level), 0.0, 0.0, cv::INTER_AREA); + if (!mask.empty()){ + cv::resize(mask, lvlMask, cv::Size(img.cols / level, img.rows / level), 0.0, 0.0, cv::INTER_LINEAR); + cv::threshold(lvlMask, lvlMask, 127, 255, cv::THRESH_BINARY); + } + } + // Source bytes are kept verbatim when nothing changed the pixels + std::vector bytes = (level == 1 && pixelsUntouched && isJpegOrPng(cam.filePath)) + ? readFileBytes(cam.filePath) + : encodeJpeg(lvlImg); + cache.put(imageKey(cam.imageId, level), std::move(bytes)); + if (!lvlMask.empty()){ + cache.put(maskKey(cam.imageId, level), encodePng(lvlMask)); + } + } +} diff --git a/image_store.hpp b/image_store.hpp new file mode 100644 index 00000000..33bc3245 --- /dev/null +++ b/image_store.hpp @@ -0,0 +1,64 @@ +#ifndef IMAGE_STORE_H +#define IMAGE_STORE_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include "input_data.hpp" + +using Blob = std::shared_ptr>; + +// Thread-safe cache of compressed image bytes with a RAM cap. Least recently +// used entries are saved to a per-process directory under baseDir and reloaded on demand +class BlobCache{ +public: + BlobCache(uint64_t capBytes, const std::filesystem::path &baseDir); + ~BlobCache(); + + void put(const std::string &key, std::vector &&bytes); + Blob get(const std::string &key); + bool has(const std::string &key) const; + uint64_t residentBytes() const { return resident; } + uint64_t capacityBytes() const { return cap; } + +private: + struct Entry{ + Blob bytes; // null when saved + uint64_t size = 0; + bool saved = false; + std::list::iterator lru; + }; + void touchLocked(Entry &e, const std::string &key); + void evictLocked(const std::string &keep); + std::filesystem::path savePath(const std::string &key) const; + + mutable std::mutex m; + std::unordered_map map; + std::list lru; // front = most recently used + uint64_t resident = 0; + uint64_t cap; + bool lowRamGuardTripped = false; + std::filesystem::path saveDir; +}; + +// Preprocesses every camera image once (resize, undistort, downscale levels) +// into compressed blobs and updates the camera intrinsics and dimensions +struct ImageStore{ + ImageStore(uint64_t capBytes, const std::filesystem::path &baseDir); + + void prepare(std::vector &cameras, float downscaleFactor, int numDownscales); + void prepareCamera(Camera &cam, float downscaleFactor); + + static std::string imageKey(int imageId, int level); + static std::string maskKey(int imageId, int level); + + BlobCache cache; + int numLevels = 1; +}; + +#endif diff --git a/input_data.cpp b/input_data.cpp index 90c91dec..c1661c5d 100644 --- a/input_data.cpp +++ b/input_data.cpp @@ -1,18 +1,6 @@ #include -#include -#include -#ifdef USE_CUDA -#include -#elif defined(USE_HIP) -#include -#endif -#ifdef __APPLE__ -#include -#endif #include #include "input_data.hpp" -#include "cv_utils.hpp" -#include "undistort.hpp" namespace fs = std::filesystem; using namespace torch::indexing; @@ -48,195 +36,10 @@ torch::Tensor Camera::getIntrinsicsMatrix(){ {0.0f, 0.0f, 1.0f}}, torch::kFloat32); } -void Camera::loadImage(float downscaleFactor){ - // Populates image and K, then updates the camera parameters - // Caution: this function has destructive behaviors - // and should be called only once - if (image.numel()) std::runtime_error("loadImage already called"); - - { - static std::mutex logMutex; - std::lock_guard lock(logMutex); - std::cout << "Loading " << fs::path(filePath).filename().string() << std::endl; - } - - 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 (cImg.rows != height || cImg.cols != width){ - rescaleF = static_cast(cImg.rows) / static_cast(height); - } - fx *= rescaleF; - fy *= rescaleF; - cx *= rescaleF; - cy *= rescaleF; - - if (downscaleFactor > 1.0f){ - float scaleFactor = 1.0f / downscaleFactor; - cv::resize(cImg, cImg, cv::Size(), scaleFactor, scaleFactor, cv::INTER_AREA); - fx *= scaleFactor; - fy *= scaleFactor; - cx *= scaleFactor; - cy *= scaleFactor; - } - - 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()){ - 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); - 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{ - image = imageToTensor(cImg); - } - - height = image.size(0); - width = image.size(1); - 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){ - if (downscaleFactor <= 1) return image; - else{ - - // torch::jit::script::Module container = torch::jit::load("gt.pt"); - // return container.attr("val").toTensor(); - - if (imagePyramids.find(downscaleFactor) != imagePyramids.end()){ - return imagePyramids[downscaleFactor]; - } - - // Rescale, store and return - cv::Mat cImg = tensorToImage(image); - cv::resize(cImg, cImg, cv::Size(cImg.cols / downscaleFactor, cImg.rows / downscaleFactor), 0.0, 0.0, cv::INTER_AREA); - torch::Tensor t = imageToTensor(cImg); - imagePyramids[downscaleFactor] = t; - return t; - } -} - 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; } -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{ 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(freeB / 2); - } -#elif defined(USE_HIP) - size_t freeB = 0, totalB = 0; - if (hipMemGetInfo(&freeB, &totalB) == hipSuccess){ - return static_cast(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 &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 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" }; @@ -291,7 +94,7 @@ std::tuple, Camera *> InputData::getCameras(bool validate, c void InputData::saveCameras(const std::string &filename, bool keepCrs){ json j = json::array(); - + for (size_t i = 0; i < cameras.size(); i++){ Camera &cam = cameras[i]; @@ -305,7 +108,7 @@ void InputData::saveCameras(const std::string &filename, bool keepCrs){ torch::Tensor R = cam.camToWorld.index({Slice(None, 3), Slice(None, 3)}); torch::Tensor T = cam.camToWorld.index({Slice(None, 3), Slice(3,4)}).squeeze(); - + // Flip z and y R = torch::matmul(R, torch::diag(torch::tensor({1.0f, -1.0f, -1.0f}))); @@ -324,10 +127,10 @@ void InputData::saveCameras(const std::string &filename, bool keepCrs){ camera["rotation"] = rotation; j.push_back(camera); } - + std::ofstream of(filename); of << j; of.close(); std::cout << "Wrote " << filename << std::endl; -} \ No newline at end of file +} diff --git a/input_data.hpp b/input_data.hpp index eea4363e..9eb36b74 100644 --- a/input_data.hpp +++ b/input_data.hpp @@ -38,27 +38,11 @@ struct Camera{ camToWorld(camToWorld), filePath(filePath) {} torch::Tensor getIntrinsicsMatrix(); bool hasDistortionParameters(); - torch::Tensor getImage(int downscaleFactor); - torch::Tensor getMask(int downscaleFactor); - torch::Tensor getEdgeMap(int downscaleFactor); - torch::Tensor getImageGpu(int downscaleFactor, const torch::Device &device); - torch::Tensor getMaskGpu(int downscaleFactor, const torch::Device &device); - torch::Tensor getEdgeMapGpu(int downscaleFactor, const torch::Device &device); - bool hasMask() const { return mask.numel() > 0; } - void loadImage(float downscaleFactor); torch::Tensor K; - torch::Tensor image; - torch::Tensor mask; // [H,W] float 0/1, aligned with image - - std::unordered_map imagePyramids; - std::unordered_map maskPyramids; - std::unordered_map edgePyramids; - std::unordered_map gpuImageCache; - std::unordered_map gpuMaskCache; - std::unordered_map gpuEdgeCache; - - static bool gpuCacheEnabled; + // Pixels are not kept here; the ImageStore decodes them on demand by imageId + int imageId = -1; + bool hasMask = false; }; struct Points{ diff --git a/model.cpp b/model.cpp index 2a446e38..465dc8ff 100644 --- a/model.cpp +++ b/model.cpp @@ -12,6 +12,7 @@ #include "tensor_math.hpp" #include "gsplat.hpp" #include "utils.hpp" +#include "cv_utils.hpp" #include "rad.hpp" #ifdef USE_MPS @@ -73,7 +74,7 @@ void Model::releaseOptimizers(){ } -torch::Tensor Model::forward(Camera& cam, int step){ +torch::Tensor Model::forward(Camera& cam, int step, const torch::Tensor &edgeMap){ const float scaleFactor = getDownscaleFactor(step); const float fx = cam.fx / scaleFactor; @@ -194,8 +195,8 @@ torch::Tensor Model::forward(Camera& cam, int step){ densificationInfo = torch::empty({0}, fOpts); xyAbsGrad = step <= densifyUntilIter ? torch::zeros({means.size(0), 2}, fOpts) : torch::empty({0}, fOpts); - }else if (edgeGuidance){ - camEdgeMap = cam.getEdgeMapGpu(getDownscaleFactor(step), device); + }else if (edgeGuidance && edgeMap.defined()){ + camEdgeMap = edgeMap; } if (device == torch::kCPU){ @@ -375,17 +376,18 @@ std::tuple Model::computeMultiViewScores(int step, for (size_t v = 0; v < numViews; v++){ Camera &cam = (*trainCams)[indices[v]]; int ds = getDownscaleFactor(step); - torch::Tensor gt = cam.getImageGpu(ds, device); + FrameLease frame = images->acquire(cam, ds, edgeGuidance); + torch::Tensor gt = frame->image; errorMap = torch::zeros({gt.size(0), gt.size(1)}, fOpts); densificationInfo = torch::zeros({4, N}, fOpts); xyAbsGrad = torch::empty({0}, fOpts); - torch::Tensor rgb = forward(cam, step); + torch::Tensor rgb = forward(cam, step, frame->edges); { torch::NoGradGuard noGrad; - torch::Tensor l1Map = (rgb.detach() - gt).abs().mean(-1); + torch::Tensor l1Map = (rgb.detach() - toUnitFloat(gt)).abs().mean(-1); float lo = l1Map.min().item(); float hi = l1Map.max().item(); torch::Tensor norm = (l1Map - lo) / (std::max)(hi - lo, 1e-8f); @@ -1028,7 +1030,7 @@ torch::Tensor Model::mainLoss(torch::Tensor &rgb, torch::Tensor >, torch::Tens torch::Tensor m = hasMask ? mask : torch::empty({0}, rgb.options()); loss = fusedL1SsimLoss(rgb, gt, m, ssimWeight, !hasMask); }else{ - torch::Tensor absDiff = torch::abs(gt - rgb); + torch::Tensor absDiff = torch::abs(toUnitFloat(gt) - rgb); loss = hasMask ? (mask.unsqueeze(-1) * absDiff).sum() / (mask.sum() * gt.size(2) + 1e-8f) : absDiff.sum() / (static_cast(gt.numel()) + 1e-8f); diff --git a/model.hpp b/model.hpp index 415d3191..40d0195c 100644 --- a/model.hpp +++ b/model.hpp @@ -9,6 +9,7 @@ #include "spherical_harmonics.hpp" #include "ssim.hpp" #include "input_data.hpp" +#include "image_pipeline.hpp" using namespace torch::indexing; using namespace torch::autograd; @@ -74,7 +75,7 @@ struct Model{ void setupOptimizers(); void releaseOptimizers(); - torch::Tensor forward(Camera& cam, int step); + torch::Tensor forward(Camera& cam, int step, const torch::Tensor &edgeMap = torch::Tensor()); void optimizerStepCadence(int step); // FastGS stepping schedule with gradient accumulation void schedulersStep(int step); int getDownscaleFactor(int step); @@ -110,6 +111,7 @@ struct Model{ float spatialLrScale = 1.0f; std::vector *trainCams = nullptr; // set by the trainer, used for multi-view scoring + ImagePipeline *images = nullptr; // set by the trainer, decodes ground truth on demand torch::Tensor radii; // set in forward() torch::Tensor xys; // set in forward() diff --git a/opensplat.cpp b/opensplat.cpp index ecde4f18..0d2676d1 100644 --- a/opensplat.cpp +++ b/opensplat.cpp @@ -6,6 +6,10 @@ #include "cv_utils.hpp" #include "constants.hpp" #include "zip_utils.hpp" +#include "image_store.hpp" +#include "image_pipeline.hpp" +#include "sysinfo.hpp" +#include #include #ifdef USE_VISUALIZATION @@ -43,7 +47,6 @@ int main(int argc, char *argv[]){ ("no-edge-guidance", "Disable Canny edge weighting of the densification importance", cxxopts::value()->default_value("false")) ("max-gaussians", "Maximum number of gaussians (0 = unlimited)", cxxopts::value()->default_value("5000000")) ("no-masks", "Ignore image masks even when present", cxxopts::value()->default_value("false")) - ("no-gpu-cache", "Do not cache images/masks on the GPU (reduces VRAM usage, slower)", cxxopts::value()->default_value("false")) #ifdef USE_VISUALIZATION ("has-visualization", "Show the visualization steps of training", cxxopts::value()->default_value("0")) #endif @@ -73,12 +76,15 @@ int main(int argc, char *argv[]){ const std::string projectRoot = result["input"].as(); + // Default outputs and temporary directories go next to the input scene + // (for a .zip, next to the archive) + fs::path sceneDir = fs::absolute(fs::path(projectRoot)); + if (sceneDir.filename().empty()) sceneDir = sceneDir.parent_path(); + sceneDir = sceneDir.parent_path(); + std::string outputScene = result["output"].as(); if (result.count("output") == 0){ - // Output next to the input scene, for a .zip, next to the archive - fs::path in = fs::absolute(fs::path(projectRoot)); - if (in.filename().empty()) in = in.parent_path(); - outputScene = (in.parent_path() / outputScene).string(); + outputScene = (sceneDir / outputScene).string(); } const std::string outputCameras = result["output-cameras"].as(); const int saveEvery = result["save-every"].as(); @@ -128,7 +134,7 @@ int main(int argc, char *argv[]){ try{ std::string projectPath = projectRoot; - if (isZipArchive(projectRoot)) projectPath = extractZipToCache(projectRoot); + if (isZipArchive(projectRoot)) projectPath = extractZipToCache(projectRoot, sceneDir.string()); InputData inputData = inputDataFromX(projectPath); int numMasks = 0; @@ -140,9 +146,10 @@ int main(int argc, char *argv[]){ } if (numMasks > 0) std::cout << "Found " << numMasks << " masks" << std::endl; - parallel_for(inputData.cameras.begin(), inputData.cameras.end(), [&downScaleFactor](Camera &cam){ - cam.loadImage(downScaleFactor); - }); + // Images are preprocessed once into a compressed cache (90% of RAM, + // saving to disk next to the scene) and decoded on demand while training + ImageStore imageStore(physicalRamBytes() / 10 * 9, sceneDir); + imageStore.prepare(inputData.cameras, downScaleFactor, numDownscales); // Withhold a validation camera if necessary auto t = inputData.getCameras(validate, valImage); @@ -183,7 +190,11 @@ int main(int argc, char *argv[]){ device); model.trainCams = &cams; model.edgeGuidance = !result["no-edge-guidance"].as(); - Camera::gpuCacheEnabled = !result["no-gpu-cache"].as(); + + const int hw = (std::max)(1u, std::thread::hardware_concurrency()); + const int decodeThreads = device == torch::kCPU ? (std::max)(1, hw / 4) : (std::max)(1, (std::min)(8, hw / 2)); + ImagePipeline images(imageStore, device, decodeThreads); + model.images = &images; std::vector< size_t > camIndices( cams.size() ); std::iota( camIndices.begin(), camIndices.end(), 0 ); @@ -197,14 +208,22 @@ int main(int argc, char *argv[]){ } for (; step <= numIters; step++){ + auto stepStart = std::chrono::steady_clock::now(); Camera& cam = cams[ camsIter.next() ]; + // Keep the decoders ahead of the training loop + for (int k = 0; k < images.prefetchTarget(); k++){ + images.request(cams[camsIter.peek(k)], model.getDownscaleFactor(step + 1 + k)); + } + torch::Tensor rgb = model.forward(cam, step); - torch::Tensor gt = cam.getImageGpu(model.getDownscaleFactor(step), device); - torch::Tensor mask = cam.getMaskGpu(model.getDownscaleFactor(step), device); + FrameLease frame = images.acquire(cam, model.getDownscaleFactor(step)); + torch::Tensor gt = frame->image; + torch::Tensor mask = frame->mask; torch::Tensor mainLoss = model.mainLoss(rgb, gt, mask, ssimWeight); mainLoss.backward(); + frame.reset(); if (step % displayStep == 0) { const float percentage = static_cast(step) / numIters; @@ -214,6 +233,7 @@ int main(int argc, char *argv[]){ model.afterTrain(step); model.optimizerStepCadence(step); model.schedulersStep(step); + images.noteTrainStep(std::chrono::duration(std::chrono::steady_clock::now() - stepStart).count()); if (saveEvery > 0 && step % saveEvery == 0){ fs::path p(outputScene); @@ -248,15 +268,17 @@ int main(int argc, char *argv[]){ // Validate if (valCam != nullptr){ torch::Tensor rgb = model.forward(*valCam, numIters); - torch::Tensor gt = valCam->getImageGpu(model.getDownscaleFactor(numIters), device); - torch::Tensor valMask = valCam->getMaskGpu(model.getDownscaleFactor(numIters), device); + FrameLease frame = images.acquire(*valCam, model.getDownscaleFactor(numIters)); + torch::Tensor gt = frame->image; + torch::Tensor valMask = frame->mask; std::cout << valCam->filePath << " validation loss: " << model.mainLoss(rgb, gt, valMask, ssimWeight).item() << std::endl; + torch::Tensor gtF = toUnitFloat(gt); torch::Tensor mse; if (valMask.defined() && valMask.numel() > 0){ - mse = (valMask.unsqueeze(-1) * (rgb - gt).pow(2)).sum() / (valMask.sum() * gt.size(2) + 1e-8f); + mse = (valMask.unsqueeze(-1) * (rgb - gtF).pow(2)).sum() / (valMask.sum() * gtF.size(2) + 1e-8f); }else{ - mse = (rgb - gt).pow(2).mean(); + mse = (rgb - gtF).pow(2).mean(); } std::cout << valCam->filePath << " validation PSNR: " << (10.0f * torch::log10(1.0f / mse)).item() << std::endl; } diff --git a/rasterize_gaussians.cpp b/rasterize_gaussians.cpp index cb728dfa..87b1c14a 100644 --- a/rasterize_gaussians.cpp +++ b/rasterize_gaussians.cpp @@ -32,7 +32,7 @@ std::tuple fused_loss_forward_tensor_cpu( const torch::Tensor &rendered, const torch::Tensor >, diff --git a/rasterizer/gsplat-cpu/gsplat_cpu.cpp b/rasterizer/gsplat-cpu/gsplat_cpu.cpp index 368bed2d..1d4600f6 100644 --- a/rasterizer/gsplat-cpu/gsplat_cpu.cpp +++ b/rasterizer/gsplat-cpu/gsplat_cpu.cpp @@ -590,10 +590,8 @@ torch::Tensor compute_sh_forward_tensor_cpu( return (result.index({"...", None}) * coeffs).sum(-2); } -// Fused L1 + DSSIM loss over [H,W,C] images: same method as the GPU -// backends (two-pass separable 11-tap blur computing all five moments in one -// sweep, closed-form SSIM partials saved for a direct backward), parallelized -// over row bands. +// Fused L1 + DSSIM loss over [H,W,C] images, same method as the GPU backends +// (tiled 11-tap blur, closed-form SSIM partials), parallelized over row bands namespace { @@ -611,9 +609,12 @@ inline bool lossValid(int y, int x, int H, int W, bool validPad){ return x >= LOSS_HALO && x < W - LOSS_HALO && y >= LOSS_HALO && y < H - LOSS_HALO; } -} +// Ground truth may be float [0,1] or uint8 [0,255] +inline float gtUnit(float v){ return v; } +inline float gtUnit(uint8_t v){ return static_cast(v) * (1.0f / 255.0f); } -std::tuple fused_loss_forward_tensor_cpu( +template +std::tuple fusedLossForwardCpuImpl( const torch::Tensor &rendered, const torch::Tensor >, const torch::Tensor &mask, @@ -630,7 +631,7 @@ std::tuple fused_loss_forward_tensor_cpu( torch::Tensor m = hasMask ? mask.contiguous() : torch::Tensor(); const float *rp = r.data_ptr(); - const float *gp = g.data_ptr(); + const GtT *gp = g.data_ptr(); const float *mp = hasMask ? m.data_ptr() : nullptr; auto fOpts = r.options(); @@ -664,7 +665,7 @@ std::tuple fused_loss_forward_tensor_cpu( if (xx < 0 || xx >= W) continue; const float w = lossGauss[LOSS_HALO + d]; const float X = rp[(y * W + xx) * C + c]; - const float Y = gp[(y * W + xx) * C + c]; + const float Y = gtUnit(gp[(y * W + xx) * C + c]); sX += X * w; sX2 += X * X * w; sY += Y * w; @@ -737,7 +738,7 @@ std::tuple fused_loss_forward_tensor_cpu( if (gate == 0.0f) continue; float l1 = 0.0f; for (int c = 0; c < C; c++){ - l1 += std::fabs(rp[p * C + c] - gp[p * C + c]); + l1 += std::fabs(rp[p * C + c] - gtUnit(gp[p * C + c])); } lossSum += gate * ((1.0f - ssim_weight) * l1 + static_cast(C) * ssim_weight * (1.0f - sp[p])); @@ -757,7 +758,8 @@ std::tuple fused_loss_forward_tensor_cpu( return std::make_tuple(stats, partials); } -torch::Tensor fused_loss_backward_tensor_cpu( +template +torch::Tensor fusedLossBackwardCpuImpl( const torch::Tensor &rendered, const torch::Tensor >, const torch::Tensor &mask, @@ -776,7 +778,7 @@ torch::Tensor fused_loss_backward_tensor_cpu( torch::Tensor m = hasMask ? mask.contiguous() : torch::Tensor(); const float *rp = r.data_ptr(); - const float *gp = g.data_ptr(); + const GtT *gp = g.data_ptr(); const float *mp = hasMask ? m.data_ptr() : nullptr; const at::Half *pBase = partials.data_ptr(); const long long planeSize = static_cast(H) * W * C; @@ -839,7 +841,7 @@ torch::Tensor fused_loss_backward_tensor_cpu( } const long long p = static_cast(y) * W + x; const float p1 = rp[p * C + c]; - const float p2 = gp[p * C + c]; + const float p2 = gtUnit(gp[p * C + c]); const float gradSsim = s0 + 2.f * p1 * s1 + p2 * s2; const float gate = gateAt(y, x); @@ -851,4 +853,36 @@ torch::Tensor fused_loss_backward_tensor_cpu( } return vRendered; +} + +} + +std::tuple fused_loss_forward_tensor_cpu( + const torch::Tensor &rendered, + const torch::Tensor >, + const torch::Tensor &mask, + const float ssim_weight, + const bool valid_padding, + const bool want_grad +){ + if (gt.scalar_type() == torch::kU8){ + return fusedLossForwardCpuImpl(rendered, gt, mask, ssim_weight, valid_padding, want_grad); + } + return fusedLossForwardCpuImpl(rendered, gt, mask, ssim_weight, valid_padding, want_grad); +} + +torch::Tensor fused_loss_backward_tensor_cpu( + const torch::Tensor &rendered, + const torch::Tensor >, + const torch::Tensor &mask, + const torch::Tensor &partials, + const torch::Tensor &stats, + const torch::Tensor &v_loss, + const float ssim_weight, + const bool valid_padding +){ + if (gt.scalar_type() == torch::kU8){ + return fusedLossBackwardCpuImpl(rendered, gt, mask, partials, stats, v_loss, ssim_weight, valid_padding); + } + return fusedLossBackwardCpuImpl(rendered, gt, mask, partials, stats, v_loss, ssim_weight, valid_padding); } \ No newline at end of file diff --git a/rasterizer/gsplat-metal/bindings.h b/rasterizer/gsplat-metal/bindings.h index 7180142d..6ac9bcaa 100644 --- a/rasterizer/gsplat-metal/bindings.h +++ b/rasterizer/gsplat-metal/bindings.h @@ -16,10 +16,8 @@ std::tuple< torch::Tensor> // output radii compute_cov2d_bounds_tensor(const int num_pts, torch::Tensor &A); -// Fused L1 + DSSIM loss over [H,W,C] float images. -// Returns {stats, partials}: stats[0] = loss, stats[1] = normalization -// denominator (both on device); partials holds the SSIM derivative maps -// needed by the backward pass (empty when want_grad is false). +// Fused L1 + DSSIM loss; gt is float [0,1] or uint8 [0,255]. Returns {stats, partials}: +// stats = {loss, normalization denominator}, partials = SSIM derivative maps for the backward std::tuple fused_loss_forward_tensor( const torch::Tensor &rendered, const torch::Tensor >, @@ -120,7 +118,8 @@ std::tuple map_gaussian_to_intersects_tensor( torch::Tensor get_tile_bin_edges_tensor( int num_intersects, - const torch::Tensor &isect_ids_sorted + const torch::Tensor &isect_ids_sorted, + const std::tuple tile_bounds ); std::tuple< diff --git a/rasterizer/gsplat-metal/gsplat_metal.metal b/rasterizer/gsplat-metal/gsplat_metal.metal index e1bd563e..9b6fc1bc 100644 --- a/rasterizer/gsplat-metal/gsplat_metal.metal +++ b/rasterizer/gsplat-metal/gsplat_metal.metal @@ -1578,10 +1578,8 @@ kernel void compute_cov2d_bounds_kernel( radii[row] = radius; } -// Fused L1 + DSSIM loss over [H,W,C] images. One kernel computes the -// SSIM map and the closed-form partials in a threadgroup-memory tile -// (two-pass separable 11-tap blur), one reduces to the scalar loss -// on-device, and one produces dL/dimage directly. +// Fused L1 + DSSIM loss over [H,W,C] images in three kernels: +// SSIM map + partials (tiled 11-tap blur), loss reduction, and dL/dimage #define LOSS_BX 16 #define LOSS_BY 16 @@ -1608,6 +1606,17 @@ inline bool loss_valid(int y, int x, int H, int W, bool validPad){ return x >= LOSS_HALO && x < W - LOSS_HALO && y >= LOSS_HALO && y < H - LOSS_HALO; } +// Ground truth is float [0,1] (gt) or uint8 [0,255] (gtU8), selected by gtIsU8 +inline float gt_at(device const float* gt, device const uchar* gtU8, int gtIsU8, int idx){ + return gtIsU8 ? (float)gtU8[idx] * (1.0f / 255.0f) : gt[idx]; +} + +inline float loss_pix_gt(device const float* gt, device const uchar* gtU8, int gtIsU8, + int y, int x, int c, int H, int W, int C){ + if (x < 0 || x >= W || y < 0 || y >= H) return 0.0f; + return gt_at(gt, gtU8, gtIsU8, (y * W + x) * C + c); +} + kernel void fused_loss_fwd_kernel( constant int& H, constant int& W, @@ -1615,6 +1624,8 @@ kernel void fused_loss_fwd_kernel( constant int& wantGrad, device const float* rendered, device const float* gt, + constant int& gtIsU8, + device const uchar* gtU8, device float* ssimMap, // [H,W] channel mean device half* pMu, // [H,W,C] each (dummy when !wantGrad) device half* pS1, @@ -1643,7 +1654,7 @@ kernel void fused_loss_fwd_kernel( const int gy = tileY + ly - LOSS_HALO; const int gx = tileX + lx - LOSS_HALO; sTile[ly][lx][0] = loss_pix(rendered, gy, gx, c, H, W, C); - sTile[ly][lx][1] = loss_pix(gt, gy, gx, c, H, W, C); + sTile[ly][lx][1] = loss_pix_gt(gt, gtU8, gtIsU8, gy, gx, c, H, W, C); } } threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1726,6 +1737,8 @@ kernel void fused_loss_reduce_kernel( constant int& validPad, device const float* rendered, device const float* gt, + constant int& gtIsU8, + device const uchar* gtU8, device const float* ssimMap, device const float* mask, device atomic_float* out, @@ -1748,7 +1761,7 @@ kernel void fused_loss_reduce_kernel( if (gate != 0.0f){ float l1 = 0.0f; for (int c = 0; c < C; c++){ - l1 += fabs(rendered[p * C + c] - gt[p * C + c]); + l1 += fabs(rendered[p * C + c] - gt_at(gt, gtU8, gtIsU8, p * C + c)); } const float contrib = (1.0f - ssimWeight) * l1 + (float)C * ssimWeight * (1.0f - ssimMap[p]); @@ -1795,6 +1808,8 @@ kernel void fused_loss_bwd_kernel( constant int& validPad, device const float* rendered, device const float* gt, + constant int& gtIsU8, + device const uchar* gtU8, device const float* mask, device const half* pMu, device const half* pS1, @@ -1818,7 +1833,7 @@ kernel void fused_loss_bwd_kernel( float p1 = 0.f, p2 = 0.f; if (px < W && py < H){ p1 = rendered[(py * W + px) * C + c]; - p2 = gt[(py * W + px) * C + c]; + p2 = gt_at(gt, gtU8, gtIsU8, (py * W + px) * C + c); } // Load the chain-weighted partials for the tile + halo diff --git a/rasterizer/gsplat-metal/gsplat_metal.mm b/rasterizer/gsplat-metal/gsplat_metal.mm index bd3cd858..e3d2da14 100644 --- a/rasterizer/gsplat-metal/gsplat_metal.mm +++ b/rasterizer/gsplat-metal/gsplat_metal.mm @@ -542,11 +542,15 @@ void dispatchKernel(MetalContext* ctx, id cpso, MTLSize torch::Tensor get_tile_bin_edges_tensor( int num_intersects, - const torch::Tensor &isect_ids_sorted + const torch::Tensor &isect_ids_sorted, + const std::tuple tile_bounds ) { CHECK_INPUT(isect_ids_sorted); + // Indexed by tile id, so it must cover every tile even when few + // gaussians intersect the image + const int num_tiles = std::get<0>(tile_bounds) * std::get<1>(tile_bounds); torch::Tensor tile_bins = torch::zeros( - {num_intersects, 2}, isect_ids_sorted.options().dtype(torch::kInt32) + {num_tiles, 2}, isect_ids_sorted.options().dtype(torch::kInt32) ); MetalContext* ctx = get_global_context(); @@ -940,6 +944,10 @@ void dispatchKernel(MetalContext* ctx, id cpso, MTLSize const int W = rendered.size(1); const int C = rendered.size(2); const bool hasMask = mask.defined() && mask.numel() > 0; + // Ground truth is float [0,1] or uint8 [0,255]; the unused variant is a dummy buffer + const bool gtU8 = gt.scalar_type() == torch::kU8; + torch::Tensor gtF = gtU8 ? torch::empty({1}, rendered.options()) : gt; + torch::Tensor gtU = gtU8 ? gt : torch::empty({1}, rendered.options().dtype(torch::kU8)); if (hasMask){ CHECK_INPUT(mask); } auto opts = rendered.options(); @@ -963,7 +971,9 @@ void dispatchKernel(MetalContext* ctx, id cpso, MTLSize EncodeArg::scalar((int32_t)C), EncodeArg::scalar((int32_t)(want_grad ? 1 : 0)), EncodeArg::tensor(rendered), - EncodeArg::tensor(gt), + EncodeArg::tensor(gtF), + EncodeArg::scalar((int32_t)(gtU8 ? 1 : 0)), + EncodeArg::tensor(gtU), EncodeArg::tensor(ssimMap), EncodeArg::tensor(pMu), EncodeArg::tensor(pS1), @@ -983,7 +993,9 @@ void dispatchKernel(MetalContext* ctx, id cpso, MTLSize EncodeArg::scalar(ssim_weight), EncodeArg::scalar((int32_t)(valid_padding ? 1 : 0)), EncodeArg::tensor(rendered), - EncodeArg::tensor(gt), + EncodeArg::tensor(gtF), + EncodeArg::scalar((int32_t)(gtU8 ? 1 : 0)), + EncodeArg::tensor(gtU), EncodeArg::tensor(ssimMap), EncodeArg::tensor(maskBuf), EncodeArg::tensor(stats) @@ -1010,6 +1022,10 @@ void dispatchKernel(MetalContext* ctx, id cpso, MTLSize const int W = rendered.size(1); const int C = rendered.size(2); const bool hasMask = mask.defined() && mask.numel() > 0; + // Ground truth is float [0,1] or uint8 [0,255]; the unused variant is a dummy buffer + const bool gtU8 = gt.scalar_type() == torch::kU8; + torch::Tensor gtF = gtU8 ? torch::empty({1}, rendered.options()) : gt; + torch::Tensor gtU = gtU8 ? gt : torch::empty({1}, rendered.options().dtype(torch::kU8)); torch::Tensor vRendered = torch::empty_like(rendered); torch::Tensor maskBuf = hasMask ? mask : torch::empty({1}, rendered.options()); @@ -1028,7 +1044,9 @@ void dispatchKernel(MetalContext* ctx, id cpso, MTLSize EncodeArg::scalar(ssim_weight), EncodeArg::scalar((int32_t)(valid_padding ? 1 : 0)), EncodeArg::tensor(rendered), - EncodeArg::tensor(gt), + EncodeArg::tensor(gtF), + EncodeArg::scalar((int32_t)(gtU8 ? 1 : 0)), + EncodeArg::tensor(gtU), EncodeArg::tensor(maskBuf), EncodeArg::tensor(pMu), EncodeArg::tensor(pS1), diff --git a/rasterizer/gsplat/bindings.cu b/rasterizer/gsplat/bindings.cu index cdb92f81..19bb50f2 100644 --- a/rasterizer/gsplat/bindings.cu +++ b/rasterizer/gsplat/bindings.cu @@ -25,10 +25,8 @@ namespace cg = cooperative_groups; -// Fused L1 + DSSIM loss over [H,W,C] images. One kernel computes the -// SSIM map and the closed-form partials in a shared-memory tile (two-pass -// separable 11-tap blur), one reduces to the scalar loss on-device, and one -// produces dL/dimage directly. +// Fused L1 + DSSIM loss over [H,W,C] images in three kernels: +// SSIM map + partials (tiled 11-tap blur), loss reduction, and dL/dimage #define LOSS_BX 16 #define LOSS_BY 16 @@ -45,11 +43,16 @@ __device__ __constant__ float lossGauss[11] = { 0.21300552785396576f, 0.10936068743467331f, 0.036000773310661316f, 0.0075987582094967365f, 0.001028380123898387f}; +// Ground truth may be float [0,1] or uint8 [0,255] +__device__ __forceinline__ float gt_unit(float v){ return v; } +__device__ __forceinline__ float gt_unit(uint8_t v){ return static_cast(v) * (1.0f / 255.0f); } + +template __device__ __forceinline__ float loss_pix( - const float* __restrict__ img, int y, int x, int c, int H, int W, int C + const T* __restrict__ img, int y, int x, int c, int H, int W, int C ){ if (x < 0 || x >= W || y < 0 || y >= H) return 0.0f; - return img[(y * W + x) * C + c]; + return gt_unit(img[(y * W + x) * C + c]); } // A pixel participates in the unmasked loss only away from the blur border @@ -58,10 +61,11 @@ __device__ __forceinline__ bool loss_valid(int y, int x, int H, int W, bool vali return x >= LOSS_HALO && x < W - LOSS_HALO && y >= LOSS_HALO && y < H - LOSS_HALO; } +template __global__ void fused_loss_fwd_kernel( const int H, const int W, const int C, const float* __restrict__ rendered, - const float* __restrict__ gt, + const GtT* __restrict__ gt, float* __restrict__ ssimMap, // [H,W] channel mean __half* __restrict__ pMu, // [H,W,C] each, or nullptr __half* __restrict__ pS1, @@ -167,10 +171,11 @@ __global__ void fused_loss_fwd_kernel( // Reduces the combined loss over pixels: out[0] += sum of gate * ((1-w)*sum_c|d_c| + C*w*(1-ssim)), // out[1] += sum of gate. The gate is the mask value or the valid-padding indicator. +template __global__ void fused_loss_reduce_kernel( const int H, const int W, const int C, const float* __restrict__ rendered, - const float* __restrict__ gt, + const GtT* __restrict__ gt, const float* __restrict__ ssimMap, const float* __restrict__ mask, // nullptr when unmasked const float ssimWeight, @@ -192,7 +197,7 @@ __global__ void fused_loss_reduce_kernel( if (gate != 0.0f){ float l1 = 0.0f; for (int c = 0; c < C; c++){ - l1 += fabsf(rendered[p * C + c] - gt[p * C + c]); + l1 += fabsf(rendered[p * C + c] - gt_unit(gt[p * C + c])); } const float contrib = (1.0f - ssimWeight) * l1 + static_cast(C) * ssimWeight * (1.0f - ssimMap[p]); @@ -226,12 +231,13 @@ __global__ void fused_loss_finalize_kernel(const int C, float* __restrict__ out) out[1] = denom; } +template __global__ void fused_loss_bwd_kernel( const int H, const int W, const int C, const float ssimWeight, const bool validPad, const float* __restrict__ rendered, - const float* __restrict__ gt, + const GtT* __restrict__ gt, const float* __restrict__ mask, // nullptr when unmasked const __half* __restrict__ pMu, const __half* __restrict__ pS1, @@ -253,7 +259,7 @@ __global__ void fused_loss_bwd_kernel( float p1 = 0.f, p2 = 0.f; if (px < W && py < H){ p1 = rendered[(py * W + px) * C + c]; - p2 = gt[(py * W + px) * C + c]; + p2 = gt_unit(gt[(py * W + px) * C + c]); } // Load the chain-weighted partials for the tile + halo @@ -328,16 +334,15 @@ __global__ void fused_loss_bwd_kernel( } } -std::tuple fused_loss_forward_tensor( +template +static std::tuple fused_loss_forward_impl( const torch::Tensor &rendered, - const torch::Tensor >, + const GtT* gt, const torch::Tensor &mask, const float ssim_weight, const bool valid_padding, const bool want_grad ){ - CHECK_INPUT(rendered); - CHECK_INPUT(gt); const int H = rendered.size(0); const int W = rendered.size(1); const int C = rendered.size(2); @@ -354,9 +359,9 @@ std::tuple fused_loss_forward_tensor( const dim3 block(LOSS_BX, LOSS_BY); const dim3 grid((W + LOSS_BX - 1) / LOSS_BX, (H + LOSS_BY - 1) / LOSS_BY); - fused_loss_fwd_kernel<<>>( + fused_loss_fwd_kernel<<>>( H, W, C, - rendered.data_ptr(), gt.data_ptr(), + rendered.data_ptr(), gt, ssimMap.data_ptr(), pBase, pBase ? pBase + planeSize : nullptr, pBase ? pBase + 2 * planeSize : nullptr ); @@ -364,9 +369,9 @@ std::tuple fused_loss_forward_tensor( torch::Tensor stats = torch::zeros({2}, opts); const int numPix = H * W; const int reduceBlocks = (std::min)(1024, (numPix + 255) / 256); - fused_loss_reduce_kernel<<>>( + fused_loss_reduce_kernel<<>>( H, W, C, - rendered.data_ptr(), gt.data_ptr(), + rendered.data_ptr(), gt, ssimMap.data_ptr(), hasMask ? mask.data_ptr() : nullptr, ssim_weight, valid_padding, @@ -377,10 +382,27 @@ std::tuple fused_loss_forward_tensor( return std::make_tuple(stats, partials); } -torch::Tensor fused_loss_backward_tensor( +std::tuple fused_loss_forward_tensor( const torch::Tensor &rendered, const torch::Tensor >, const torch::Tensor &mask, + const float ssim_weight, + const bool valid_padding, + const bool want_grad +){ + CHECK_INPUT(rendered); + CHECK_INPUT(gt); + if (gt.scalar_type() == torch::kU8){ + return fused_loss_forward_impl(rendered, gt.data_ptr(), mask, ssim_weight, valid_padding, want_grad); + } + return fused_loss_forward_impl(rendered, gt.data_ptr(), mask, ssim_weight, valid_padding, want_grad); +} + +template +static torch::Tensor fused_loss_backward_impl( + const torch::Tensor &rendered, + const GtT* gt, + const torch::Tensor &mask, const torch::Tensor &partials, const torch::Tensor &stats, const torch::Tensor &v_loss, @@ -398,9 +420,9 @@ torch::Tensor fused_loss_backward_tensor( const dim3 block(LOSS_BX, LOSS_BY); const dim3 grid((W + LOSS_BX - 1) / LOSS_BX, (H + LOSS_BY - 1) / LOSS_BY); - fused_loss_bwd_kernel<<>>( + fused_loss_bwd_kernel<<>>( H, W, C, ssim_weight, valid_padding, - rendered.data_ptr(), gt.data_ptr(), + rendered.data_ptr(), gt, hasMask ? mask.data_ptr() : nullptr, pBase, pBase + planeSize, pBase + 2 * planeSize, stats.data_ptr(), @@ -410,6 +432,22 @@ torch::Tensor fused_loss_backward_tensor( return vRendered; } +torch::Tensor fused_loss_backward_tensor( + const torch::Tensor &rendered, + const torch::Tensor >, + const torch::Tensor &mask, + const torch::Tensor &partials, + const torch::Tensor &stats, + const torch::Tensor &v_loss, + const float ssim_weight, + const bool valid_padding +){ + if (gt.scalar_type() == torch::kU8){ + return fused_loss_backward_impl(rendered, gt.data_ptr(), mask, partials, stats, v_loss, ssim_weight, valid_padding); + } + return fused_loss_backward_impl(rendered, gt.data_ptr(), mask, partials, stats, v_loss, ssim_weight, valid_padding); +} + __global__ void compute_cov2d_bounds_kernel( const unsigned num_pts, const float* __restrict__ covs2d, float* __restrict__ conics, float* __restrict__ radii ) { @@ -705,11 +743,15 @@ std::tuple map_gaussian_to_intersects_tensor( } torch::Tensor get_tile_bin_edges_tensor( - int num_intersects, const torch::Tensor &isect_ids_sorted + int num_intersects, const torch::Tensor &isect_ids_sorted, + const std::tuple tile_bounds ) { CHECK_INPUT(isect_ids_sorted); + // Indexed by tile id, so it must cover every tile even when few + // gaussians intersect the image + const int num_tiles = std::get<0>(tile_bounds) * std::get<1>(tile_bounds); torch::Tensor tile_bins = torch::zeros( - {num_intersects, 2}, isect_ids_sorted.options().dtype(torch::kInt32) + {num_tiles, 2}, isect_ids_sorted.options().dtype(torch::kInt32) ); get_tile_bin_edges<<< (num_intersects + N_THREADS - 1) / N_THREADS, diff --git a/rasterizer/gsplat/bindings.h b/rasterizer/gsplat/bindings.h index a1babb44..d0df5f25 100644 --- a/rasterizer/gsplat/bindings.h +++ b/rasterizer/gsplat/bindings.h @@ -23,10 +23,8 @@ std::tuple< torch::Tensor> // output radii compute_cov2d_bounds_tensor(const int num_pts, torch::Tensor &A); -// Fused L1 + DSSIM loss over [H,W,C] float images. -// Returns {stats, partials}: stats[0] = loss, stats[1] = normalization -// denominator (both on device); partials holds the SSIM derivative maps -// needed by the backward pass (empty when want_grad is false). +// Fused L1 + DSSIM loss; gt is float [0,1] or uint8 [0,255]. Returns {stats, partials}: +// stats = {loss, normalization denominator}, partials = SSIM derivative maps for the backward std::tuple fused_loss_forward_tensor( const torch::Tensor &rendered, const torch::Tensor >, @@ -127,7 +125,8 @@ std::tuple map_gaussian_to_intersects_tensor( torch::Tensor get_tile_bin_edges_tensor( int num_intersects, - const torch::Tensor &isect_ids_sorted + const torch::Tensor &isect_ids_sorted, + const std::tuple tile_bounds ); std::tuple< diff --git a/ssim.hpp b/ssim.hpp index e7a99e8a..56e35026 100644 --- a/ssim.hpp +++ b/ssim.hpp @@ -4,8 +4,7 @@ #include // Fused (1-ssimWeight)*L1 + ssimWeight*DSSIM loss with autograd support. -// mask may be an empty tensor; validPadding crops the blur border from the -// unmasked loss +// mask may be empty; validPadding excludes the blur border from the unmasked loss torch::Tensor fusedL1SsimLoss(const torch::Tensor &rendered, const torch::Tensor >, const torch::Tensor &mask, float ssimWeight, bool validPadding); diff --git a/sysinfo.cpp b/sysinfo.cpp new file mode 100644 index 00000000..a0f380bd --- /dev/null +++ b/sysinfo.cpp @@ -0,0 +1,109 @@ +#include "sysinfo.hpp" + +#ifdef _WIN32 +#ifndef NOMINMAX +#define NOMINMAX +#endif +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#include +#else +#include +#include +#include +#endif + +#ifdef __APPLE__ +#include +#include +#elif defined(__linux__) +#include +#include +#include +#endif + +#ifdef USE_CUDA +#include +#elif defined(USE_HIP) +#include +#endif + +uint64_t physicalRamBytes(){ +#ifdef _WIN32 + MEMORYSTATUSEX st; + st.dwLength = sizeof(st); + if (GlobalMemoryStatusEx(&st)) return st.ullTotalPhys; +#elif defined(__APPLE__) + int64_t ram = 0; + size_t size = sizeof(ram); + if (sysctlbyname("hw.memsize", &ram, &size, nullptr, 0) == 0) return static_cast(ram); +#elif defined(__linux__) + struct sysinfo si; + if (sysinfo(&si) == 0) return static_cast(si.totalram) * si.mem_unit; +#endif + return 8ull << 30; +} + +uint64_t availableRamBytes(){ +#ifdef _WIN32 + MEMORYSTATUSEX st; + st.dwLength = sizeof(st); + if (GlobalMemoryStatusEx(&st)) return st.ullAvailPhys; +#elif defined(__APPLE__) + vm_size_t pageSize = 0; + vm_statistics64_data_t vm; + mach_msg_type_number_t count = HOST_VM_INFO64_COUNT; + mach_port_t host = mach_host_self(); + if (host_page_size(host, &pageSize) == KERN_SUCCESS && + host_statistics64(host, HOST_VM_INFO64, reinterpret_cast(&vm), &count) == KERN_SUCCESS){ + return static_cast(vm.free_count + vm.inactive_count) * pageSize; + } +#elif defined(__linux__) + std::ifstream f("/proc/meminfo"); + std::string key; + uint64_t kb = 0; + while (f >> key >> kb){ + if (key == "MemAvailable:") return kb * 1024ull; + f.ignore(256, '\n'); + } + struct sysinfo si; + if (sysinfo(&si) == 0) return static_cast(si.freeram) * si.mem_unit; +#endif + return physicalRamBytes() / 2; +} + +uint64_t freeDeviceMemoryBytes(bool gpu){ + if (gpu){ +#ifdef USE_CUDA + size_t freeB = 0, totalB = 0; + if (cudaMemGetInfo(&freeB, &totalB) == cudaSuccess) return freeB; +#elif defined(USE_HIP) + size_t freeB = 0, totalB = 0; + if (hipMemGetInfo(&freeB, &totalB) == hipSuccess) return freeB; +#endif + } + return availableRamBytes(); +} + +int currentPid(){ +#ifdef _WIN32 + return static_cast(GetCurrentProcessId()); +#else + return static_cast(getpid()); +#endif +} + +bool processAlive(int pid){ +#ifdef _WIN32 + HANDLE h = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, static_cast(pid)); + if (h == nullptr) return GetLastError() == ERROR_ACCESS_DENIED; + DWORD code = 0; + bool alive = GetExitCodeProcess(h, &code) && code == STILL_ACTIVE; + CloseHandle(h); + return alive; +#else + if (kill(pid, 0) == 0) return true; + return errno == EPERM; +#endif +} diff --git a/sysinfo.hpp b/sysinfo.hpp new file mode 100644 index 00000000..57a0f29a --- /dev/null +++ b/sysinfo.hpp @@ -0,0 +1,13 @@ +#ifndef SYSINFO_H +#define SYSINFO_H + +#include + +uint64_t physicalRamBytes(); +uint64_t availableRamBytes(); +// Free memory on the training device: GPU memory for CUDA/HIP, RAM otherwise +uint64_t freeDeviceMemoryBytes(bool gpu); +int currentPid(); +bool processAlive(int pid); + +#endif diff --git a/utils.hpp b/utils.hpp index 5aa5400e..fe5b7df8 100644 --- a/utils.hpp +++ b/utils.hpp @@ -2,6 +2,7 @@ #define UTILS_H #include +#include #include #include #include @@ -28,11 +29,41 @@ class InfiniteRandomIterator T next(){ T ret = v[i++]; - if (i >= v.size()) shuffleV(); + if (i >= v.size()){ + if (upcoming.empty()){ + shuffleV(); + }else{ + v = std::move(upcoming.front()); + upcoming.pop_front(); + i = 0; + } + } return ret; } + + // Element that next() will return k calls from now (k = 0 is the next one). + // Looking ahead does not change the sequence + T peek(size_t k){ + size_t idx = i + k; + size_t available = v.size(); + for (const VecType &p : upcoming) available += p.size(); + while (idx >= available){ + VecType p = upcoming.empty() ? v : upcoming.back(); + std::shuffle(std::begin(p), std::end(p), engine); + available += p.size(); + upcoming.push_back(std::move(p)); + } + if (idx < v.size()) return v[idx]; + idx -= v.size(); + for (const VecType &p : upcoming){ + if (idx < p.size()) return p[idx]; + idx -= p.size(); + } + return v[0]; + } private: VecType v; + std::deque upcoming; // permutations generated ahead by peek() size_t i; std::default_random_engine engine; }; diff --git a/visualizer.cpp b/visualizer.cpp index a1946f7a..ce35cbca 100644 --- a/visualizer.cpp +++ b/visualizer.cpp @@ -96,7 +96,9 @@ void Visualizer::SetGaussians(const torch::Tensor& means, void Visualizer::SetImage(const torch::Tensor& rendered_img, const torch::Tensor& gt_img) { rendered_img_ = (rendered_img.cpu() * 255).to(torch::kUInt8); - gt_img_ = (gt_img.cpu() * 255).to(torch::kUInt8); + gt_img_ = gt_img.scalar_type() == torch::kUInt8 + ? gt_img.cpu().contiguous() + : (gt_img.cpu() * 255).to(torch::kUInt8); } void Visualizer::DrawInern() { diff --git a/zip_utils.cpp b/zip_utils.cpp index 175975c5..9838b514 100644 --- a/zip_utils.cpp +++ b/zip_utils.cpp @@ -117,9 +117,9 @@ static fs::path descendSingleDir(fs::path dir){ return dir; } -std::string extractZipToCache(const std::string &zipPath){ +std::string extractZipToCache(const std::string &zipPath, const std::string &cacheDir){ fs::path zp = fs::absolute(zipPath); - fs::path dest = fs::temp_directory_path() / ("opensplat-" + crc32Hex(zp)); + fs::path dest = fs::path(cacheDir) / ("opensplat-" + crc32Hex(zp)); if (!fs::exists(dest)){ fs::path staging = dest; staging += ".partial"; diff --git a/zip_utils.hpp b/zip_utils.hpp index c77515cc..d96e1946 100644 --- a/zip_utils.hpp +++ b/zip_utils.hpp @@ -4,6 +4,8 @@ #include bool isZipArchive(const std::string &path); -std::string extractZipToCache(const std::string &zipPath); +// Extracts the archive into cacheDir (or reuses a previous extraction) +// and returns the project root inside it +std::string extractZipToCache(const std::string &zipPath, const std::string &cacheDir); #endif From f2f6d42f148f371295aa62ed933f8523bb1ea964 Mon Sep 17 00:00:00 2001 From: Piero Toffanin Date: Wed, 16 Sep 2026 12:32:31 -0400 Subject: [PATCH 2/2] Bump version --- VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/VERSION b/VERSION index 6085e946..23aa8390 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -1.2.1 +1.2.2