diff --git a/CMakeLists.txt b/CMakeLists.txt index c194e5aa..85956b72 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -35,7 +35,7 @@ target_include_directories(gsplat PRIVATE set_target_properties(gsplat PROPERTIES LINKER_LANGUAGE CXX) set_target_properties(gsplat PROPERTIES CUDA_ARCHITECTURES "70;75") -add_executable(opensplat opensplat.cpp point_io.cpp nerfstudio.cpp model.cpp kdtree_tensor.cpp spherical_harmonics.cpp cv_utils.cpp utils.cpp project_gaussians.cpp rasterize_gaussians.cpp ssim.cpp optim_scheduler.cpp) +add_executable(opensplat opensplat.cpp point_io.cpp nerfstudio.cpp model.cpp kdtree_tensor.cpp spherical_harmonics.cpp cv_utils.cpp utils.cpp project_gaussians.cpp rasterize_gaussians.cpp ssim.cpp optim_scheduler.cpp colmap.cpp input_data.cpp tensor_math.cpp) target_include_directories(opensplat PRIVATE ${PROJECT_SOURCE_DIR}/vendor/glm) target_link_libraries(opensplat PUBLIC ${STDPPFS_LIBRARY} cuda gsplat ${TORCH_LIBRARIES} ${OpenCV_LIBS}) diff --git a/README.md b/README.md index 23974175..2a9fb917 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ A free and open source implementation of 3D gaussian splatting written in C++, f ![OpenSplat](https://github.com/pierotofy/OpenSplat/assets/1951843/3461e0e4-e134-4d6a-8a56-d89d00258e41) -OpenSplat takes camera poses + sparse points and computes a [scene file](https://drive.google.com/file/d/1w-CBxyWNXF3omA8B_IeOsRmSJel3iwyr/view?usp=sharing) (.ply) that can be later imported for viewing, editing and rendering in other [software](https://github.com/MrNeRF/awesome-3D-gaussian-splatting?tab=readme-ov-file#open-source-implementations). +OpenSplat takes camera poses + sparse points (in [COLMAP](https://colmap.github.io/) or [nerfstudio](https://docs.nerf.studio/quickstart/custom_dataset.html) project format) and computes a [scene file](https://drive.google.com/file/d/1w-CBxyWNXF3omA8B_IeOsRmSJel3iwyr/view?usp=sharing) (.ply) that can be later imported for viewing, editing and rendering in other [software](https://github.com/MrNeRF/awesome-3D-gaussian-splatting?tab=readme-ov-file#open-source-implementations). Commercial use allowed and encouraged under the terms of the [AGPLv3](https://www.tldrlegal.com/license/gnu-affero-general-public-license-v3-agpl-3-0). ✅ @@ -28,7 +28,9 @@ Requirements: The software has been tested on Ubuntu 20.04 and Windows. With some changes it could run on macOS (help us by opening a PR?). ## Build Docker Image + Navigate to the root directory of OpenSplat repo that has Dockerfile and run the following command to build the Docker image: + ```bash docker build -t opensplat . ``` @@ -60,10 +62,7 @@ Wrote splat.ply The output `splat.ply` can then be dragged and dropped in one of the many [viewers](https://github.com/MrNeRF/awesome-3D-gaussian-splatting?tab=readme-ov-file#viewers) such as https://playcanvas.com/viewer. You can also edit/cleanup the scene using https://playcanvas.com/supersplat/editor -To run on your own data, choose the path to an existing [nerfstudio](https://docs.nerf.studio/) project. The project must have sparse points included (random initialization is not supported, see https://github.com/pierotofy/OpenSplat/issues/7). You can generate nerfstudio projects from [COLMAP](https://github.com/colmap/colmap/) by using nerfstudio's `ns-process-data` command: https://docs.nerf.studio/quickstart/custom_dataset.html - - -We have plans to add support for reading COLMAP projects directly in the near future. See https://github.com/pierotofy/OpenSplat/issues/1 +To run on your own data, choose the path to an existing [COLMAP](https://colmap.github.io/) or [nerfstudio](https://docs.nerf.studio/) project. The project must have sparse points included (random initialization is not supported, see https://github.com/pierotofy/OpenSplat/issues/7). There's several parameters you can tune. To view the full list: diff --git a/colmap.cpp b/colmap.cpp new file mode 100644 index 00000000..6368841f --- /dev/null +++ b/colmap.cpp @@ -0,0 +1,154 @@ +#include +#include "colmap.hpp" +#include "point_io.hpp" +#include "tensor_math.hpp" + +namespace fs = std::filesystem; +using namespace torch::indexing; + +namespace cm{ + +InputData inputDataFromColmap(const std::string &projectRoot){ + InputData ret; + fs::path cmRoot(projectRoot); + + if (!fs::exists(cmRoot / "cameras.bin") && fs::exists(cmRoot / "sparse" / "0" / "cameras.bin")){ + cmRoot = cmRoot / "sparse" / "0"; + } + + fs::path camerasPath = cmRoot / "cameras.bin"; + fs::path imagesPath = cmRoot / "images.bin"; + fs::path pointsPath = cmRoot / "points3D.bin"; + + if (!fs::exists(camerasPath)) throw std::runtime_error(camerasPath.string() + " does not exist"); + if (!fs::exists(imagesPath)) throw std::runtime_error(imagesPath.string() + " does not exist"); + if (!fs::exists(pointsPath)) throw std::runtime_error(pointsPath.string() + " does not exist"); + + std::ifstream camf(camerasPath.string(), std::ios::binary); + if (!camf.is_open()) throw std::runtime_error("Cannot open " + camerasPath.string()); + std::ifstream imgf(imagesPath.string(), std::ios::binary); + if (!imgf.is_open()) throw std::runtime_error("Cannot open " + imagesPath.string()); + + size_t numCameras = readBinary(camf); + std::vector cameras(numCameras); + + std::unordered_map camMap; + + for (size_t i = 0; i < numCameras; i++) { + Camera *cam = &cameras[i]; + + cam->id = readBinary(camf); + + CameraModel model = static_cast(readBinary(camf)); // model ID + cam->width = readBinary(camf); + cam->height = readBinary(camf); + + if (model == SimplePinhole){ + cam->fx = readBinary(camf); + cam->fy = cam->fx; + cam->cx = readBinary(camf); + cam->cy = readBinary(camf); + }else if (model == Pinhole){ + cam->fx = readBinary(camf); + cam->fy = readBinary(camf); + cam->cx = readBinary(camf); + cam->cy = readBinary(camf); + }else if (model == OpenCV){ + cam->fx = readBinary(camf); + cam->fy = readBinary(camf); + cam->cx = readBinary(camf); + cam->cy = readBinary(camf); + cam->k1 = readBinary(camf); + cam->k2 = readBinary(camf); + cam->p1 = readBinary(camf); + cam->p2 = readBinary(camf); + }else{ + throw std::runtime_error("Unsupported camera model: " + std::to_string(model)); + } + + camMap[cam->id] = cam; + } + + camf.close(); + + + size_t numImages = readBinary(imgf); + torch::Tensor unorientedPoses = torch::zeros({static_cast(numImages), 4, 4}, torch::kFloat32); + + for (size_t i = 0; i < numImages; i++){ + readBinary(imgf); // imageId + + torch::Tensor qVec = torch::tensor({ + readBinary(imgf), + readBinary(imgf), + readBinary(imgf), + readBinary(imgf) + }, torch::kFloat32); + torch::Tensor R = quatToRotMat(qVec); + torch::Tensor T = torch::tensor({ + { readBinary(imgf) }, + { readBinary(imgf) }, + { readBinary(imgf) } + }, torch::kFloat32); + + torch::Tensor Rinv = R.transpose(0, 1); + torch::Tensor Tinv = torch::matmul(-Rinv, T); + + uint32_t camId = readBinary(imgf); + + Camera cam = *camMap[camId]; + + char ch = '\0'; + std::string filePath = ""; + while(true){ + imgf.read(&ch, 1); + if (ch == '\0') break; + filePath += ch; + } + + // TODO: should "images" be an option? + cam.filePath = (fs::path(projectRoot) / "images" / filePath).string(); + + unorientedPoses[i].index_put_({Slice(None, 3), Slice(None, 3)}, Rinv); + unorientedPoses[i].index_put_({Slice(None, 3), Slice(3, 4)}, Tinv); + unorientedPoses[i][3][3] = 1.0f; + + // Convert COLMAP's camera CRS (OpenCV) to OpenGL + unorientedPoses[i].index_put_({Slice(0, 3), Slice(1,3)}, unorientedPoses[i].index({Slice(0, 3), Slice(1,3)}) * -1.0f); + + size_t numPoints2D = readBinary(imgf); + for (size_t j = 0; j < numPoints2D; j++){ + readBinary(imgf); // x + readBinary(imgf); // y + readBinary(imgf); // point3D ID + } + + ret.cameras.push_back(cam); + } + + imgf.close(); + + auto r = autoOrientAndCenterPoses(unorientedPoses); + torch::Tensor poses = std::get<0>(r); + ret.transformMatrix = std::get<1>(r); + ret.scaleFactor = 1.0f / torch::max(torch::abs(poses.index({Slice(), Slice(None, 3), 3}))).item(); + poses.index({Slice(), Slice(None, 3), 3}) *= ret.scaleFactor; + + for (size_t i = 0; i < ret.cameras.size(); i++){ + ret.cameras[i].camToWorld = poses[i]; + } + + PointSet *pSet = readPointSet(pointsPath.string()); + torch::Tensor points = pSet->pointsTensor().clone(); + + ret.points.xyz = torch::matmul(torch::cat({points, torch::ones_like(points.index({"...", Slice(None, 1)}))}, -1), + ret.transformMatrix.transpose(0, 1)); + ret.points.xyz *= ret.scaleFactor; + ret.points.rgb = pSet->colorsTensor().clone(); + + RELEASE_POINTSET(pSet); + + return ret; +} + +} \ No newline at end of file diff --git a/colmap.hpp b/colmap.hpp new file mode 100644 index 00000000..1e4e0d7f --- /dev/null +++ b/colmap.hpp @@ -0,0 +1,17 @@ +#ifndef COLMAP_H +#define COLMAP_H + +#include +#include "input_data.hpp" + +namespace cm{ + InputData inputDataFromColmap(const std::string &projectRoot); + + enum CameraModel{ + SimplePinhole = 0, Pinhole, SimpleRadial, Radial, + OpenCV, OpenCVFisheye, FullOpenCV, FOV, + SimpleRadialFisheye, RadialFisheye, ThinPrismFisheye + }; +} + +#endif \ No newline at end of file diff --git a/input_data.cpp b/input_data.cpp new file mode 100644 index 00000000..bf8e045a --- /dev/null +++ b/input_data.cpp @@ -0,0 +1,151 @@ +#include +#include "input_data.hpp" +#include "cv_utils.hpp" + +namespace fs = std::filesystem; +using namespace torch::indexing; + +namespace ns{ InputData inputDataFromNerfStudio(const std::string &projectRoot); } +namespace cm{ InputData inputDataFromColmap(const std::string &projectRoot); } + +InputData inputDataFromX(const std::string &projectRoot){ + fs::path root(projectRoot); + + if (fs::exists(root / "transforms.json")){ + return ns::inputDataFromNerfStudio(projectRoot); + }else if (fs::exists(root / "sparse") || fs::exists(root / "cameras.bin")){ + return cm::inputDataFromColmap(projectRoot); + }else{ + throw std::runtime_error("Invalid project folder (must be either a colmap or nerfstudio project folder)"); + } +} + +torch::Tensor Camera::getIntrinsicsMatrix(){ + return torch::tensor({{fx, 0.0f, cx}, + {0.0f, fy, cy}, + {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"); + std::cout << "Loading " << filePath << std::endl; + + float scaleFactor = 1.0f / downscaleFactor; + cv::Mat cImg = imreadRGB(filePath); + + 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 *= scaleFactor * rescaleF; + fy *= scaleFactor * rescaleF; + cx *= scaleFactor * rescaleF; + cy *= scaleFactor * rescaleF; + + if (downscaleFactor > 1.0f){ + float f = 1.0f / downscaleFactor; + cv::resize(cImg, cImg, cv::Size(), f, f, cv::INTER_AREA); + } + + K = getIntrinsicsMatrix(); + cv::Rect roi; + + if (hasDistortionParameters()){ + // Undistort + std::vector distCoeffs = undistortionParameters(); + cv::Mat cK = floatNxNtensorToMat(K); + cv::Mat newK = cv::getOptimalNewCameraMatrix(cK, distCoeffs, cv::Size(cImg.cols, cImg.rows), 0, cv::Size(), &roi); + + cv::Mat undistorted = cv::Mat::zeros(cImg.rows, cImg.cols, cImg.type()); + cv::undistort(cImg, undistorted, cK, distCoeffs, newK); + + image = imageToTensor(undistorted); + K = floatNxNMatToTensor(newK); + }else{ + roi = cv::Rect(0, 0, cImg.cols, cImg.rows); + image = imageToTensor(cImg); + } + + // Crop to ROI + image = image.index({Slice(roi.y, roi.y + roi.height), Slice(roi.x, roi.x + roi.width), Slice()}); + + // Update parameters + height = image.size(0); + width = image.size(1); + fx = K[0][0].item(); + fy = K[1][1].item(); + cx = K[0][2].item(); + cy = K[1][2].item(); +} + +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 || p1 != 0.0f || p2 != 0.0f; +} + +std::vector Camera::undistortionParameters(){ + std::vector p = { k1, k2, p1, p2, k3, 0.0f, 0.0f, 0.0f }; + return p; +} + +void Camera::scaleOutputResolution(float scaleFactor){ + fx = fx * scaleFactor; + fy = fy * scaleFactor; + cx = cx * scaleFactor; + cy = cy * scaleFactor; + height = static_cast(static_cast(height) * scaleFactor); + width = static_cast(static_cast(width) * scaleFactor); +} + +std::tuple, Camera *> InputData::getCameras(bool validate, const std::string &valImage){ + if (!validate) return std::make_tuple(cameras, nullptr); + else{ + size_t valIdx = -1; + std::srand(42); + + if (valImage == "random"){ + valIdx = std::rand() % cameras.size(); + }else{ + for (size_t i = 0; i < cameras.size(); i++){ + if (fs::path(cameras[i].filePath).filename().string() == valImage){ + valIdx = i; + break; + } + } + if (valIdx == -1) throw std::runtime_error(valImage + " not in the list of cameras"); + } + + std::vector cams; + Camera *valCam = nullptr; + + for (size_t i = 0; i < cameras.size(); i++){ + if (i != valIdx) cams.push_back(cameras[i]); + else valCam = &cameras[i]; + } + + return std::make_tuple(cams, valCam); + } +} diff --git a/input_data.hpp b/input_data.hpp new file mode 100644 index 00000000..2505c39b --- /dev/null +++ b/input_data.hpp @@ -0,0 +1,64 @@ +#ifndef INPUTDATA_H +#define INPUTDATA_H + +#include +#include +#include +#include +#include +#include + +enum CameraType { Perspective }; +struct Camera{ + int id = -1; + int width = 0; + int height = 0; + float fx = 0; + float fy = 0; + float cx = 0; + float cy = 0; + float k1 = 0; + float k2 = 0; + float k3 = 0; + float p1 = 0; + float p2 = 0; + torch::Tensor camToWorld; + std::string filePath = ""; + CameraType cameraType = CameraType::Perspective; + + Camera(){}; + 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), + k1(k1), k2(k2), k3(k3), p1(p1), p2(p2), + camToWorld(camToWorld), filePath(filePath) {} + + torch::Tensor getIntrinsicsMatrix(); + bool hasDistortionParameters(); + std::vector undistortionParameters(); + void scaleOutputResolution(float scaleFactor); + torch::Tensor getImage(int downscaleFactor); + + void loadImage(float downscaleFactor); + torch::Tensor K; + torch::Tensor image; + + std::unordered_map imagePyramids; +}; + +struct Points{ + torch::Tensor xyz; + torch::Tensor rgb; +}; +struct InputData{ + std::vector cameras; + float scaleFactor; + torch::Tensor transformMatrix; + Points points; + + std::tuple, Camera *> getCameras(bool validate, const std::string &valImage = "random"); +}; +InputData inputDataFromX(const std::string &projectRoot); + +#endif \ No newline at end of file diff --git a/model.cpp b/model.cpp index c9ff3023..915c334e 100644 --- a/model.cpp +++ b/model.cpp @@ -3,10 +3,9 @@ #include "tile_bounds.hpp" #include "project_gaussians.hpp" #include "rasterize_gaussians.hpp" +#include "tensor_math.hpp" #include "vendor/gsplat/config.h" -namespace ns{ - torch::Tensor randomQuatTensor(long long n){ torch::Tensor u = torch::rand(n); torch::Tensor v = torch::rand(n); @@ -19,32 +18,6 @@ torch::Tensor randomQuatTensor(long long n){ }, -1); } -torch::Tensor quatToRotMat(const torch::Tensor &quat){ - auto u = torch::unbind(torch::nn::functional::normalize(quat, torch::nn::functional::NormalizeFuncOptions().dim(-1)), -1); - torch::Tensor w = u[0]; - torch::Tensor x = u[1]; - torch::Tensor y = u[2]; - torch::Tensor z = u[3]; - return torch::stack({ - torch::stack({ - 1.0 - 2.0 * (y.pow(2) + z.pow(2)), - 2.0 * (x * y - w * z), - 2.0 * (x * z + w * y) - }, -1), - torch::stack({ - 2.0 * (x * y + w * z), - 1.0 - 2.0 * (x.pow(2) + z.pow(2)), - 2.0 * (y * z - w * x) - }, -1), - torch::stack({ - 2.0 * (x * z - w * y), - 2.0 * (y * z + w * x), - 1.0 - 2.0 * (x.pow(2) + y.pow(2)) - }, -1) - }, -2); - -} - torch::Tensor projectionMatrix(float zNear, float zFar, float fovX, float fovY, const torch::Device &device){ // OpenGL perspective projection matrix float t = zNear * std::tan(0.5f * fovY); @@ -482,6 +455,3 @@ torch::Tensor Model::mainLoss(torch::Tensor &rgb, torch::Tensor >, float ssimW torch::Tensor l1Loss = l1(rgb, gt); return (1.0f - ssimWeight) * l1Loss + ssimWeight * ssimLoss; } - - -} \ No newline at end of file diff --git a/model.hpp b/model.hpp index 29e373ad..f947231d 100644 --- a/model.hpp +++ b/model.hpp @@ -8,15 +8,13 @@ #include "kdtree_tensor.hpp" #include "spherical_harmonics.hpp" #include "ssim.hpp" +#include "input_data.hpp" #include "optim_scheduler.hpp" using namespace torch::indexing; using namespace torch::autograd; -namespace ns{ - torch::Tensor randomQuatTensor(long long n); -torch::Tensor quatToRotMat(const torch::Tensor &quat); 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); @@ -131,6 +129,4 @@ struct Model{ }; -} - #endif \ No newline at end of file diff --git a/nerfstudio.cpp b/nerfstudio.cpp index bbbe2adf..ce65a0c0 100644 --- a/nerfstudio.cpp +++ b/nerfstudio.cpp @@ -4,6 +4,7 @@ #include "nerfstudio.hpp" #include "point_io.hpp" #include "cv_utils.hpp" +#include "tensor_math.hpp" namespace fs = std::filesystem; @@ -123,44 +124,6 @@ torch::Tensor posesFromTransforms(const Transforms &t){ return poses; } -std::tuple autoOrientAndCenterPoses(const torch::Tensor &poses){ - // Center at mean and orient up - torch::Tensor origins = poses.index({"...", Slice(None, 3), 3}); - torch::Tensor translation = torch::mean(origins, 0); - torch::Tensor up = torch::mean(poses.index({Slice(), Slice(None, 3), 1}), 0); - up = up / up.norm(); - - torch::Tensor rotation = rotationMatrix(up, torch::tensor({0, 0, 1}, torch::kFloat32)); - torch::Tensor transform = torch::cat({rotation, torch::matmul(rotation, -translation.index({"...", None}))}, -1); - torch::Tensor orientedPoses = torch::matmul(transform, poses); - return std::make_tuple(orientedPoses, transform); -} - - -torch::Tensor rotationMatrix(const torch::Tensor &a, const torch::Tensor &b){ - // Rotation matrix that rotates vector a to vector b - torch::Tensor a1 = a / a.norm(); - torch::Tensor b1 = b / b.norm(); - torch::Tensor v = torch::cross(a1, b1); - torch::Tensor c = torch::dot(a1, b1); - const float EPS = 1e-8; - if (c.item() < -1 + EPS){ - torch::Tensor eps = (torch::rand(3) - 0.5f) * 0.01f; - return rotationMatrix(a1 + eps, b1); - } - torch::Tensor s = v.norm(); - torch::Tensor skew = torch::zeros({3, 3}, torch::kFloat32); - skew[0][1] = -v[2]; - skew[0][2] = v[1]; - skew[1][0] = v[2]; - skew[1][2] = -v[0]; - skew[2][0] = -v[1]; - skew[2][1] = v[0]; - - return torch::eye(3) + skew + torch::matmul(skew, skew * ((1 - c) / (s.pow(2) + EPS))); -} - - InputData inputDataFromNerfStudio(const std::string &projectRoot){ InputData ret; fs::path nsRoot(projectRoot); @@ -206,129 +169,4 @@ InputData inputDataFromNerfStudio(const std::string &projectRoot){ return ret; } -torch::Tensor Camera::getIntrinsicsMatrix(){ - return torch::tensor({{fx, 0.0f, cx}, - {0.0f, fy, cy}, - {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"); - std::cout << "Loading " << filePath << std::endl; - - float scaleFactor = 1.0f / downscaleFactor; - fx *= scaleFactor; - fy *= scaleFactor; - cx *= scaleFactor; - cy *= scaleFactor; - - cv::Mat cImg = imreadRGB(filePath); - - if (downscaleFactor > 1.0f){ - float f = 1.0f / downscaleFactor; - cv::resize(cImg, cImg, cv::Size(), f, f, cv::INTER_AREA); - } - - K = getIntrinsicsMatrix(); - cv::Rect roi; - - if (hasDistortionParameters()){ - // Undistort - std::vector distCoeffs = undistortionParameters(); - cv::Mat cK = floatNxNtensorToMat(K); - cv::Mat newK = cv::getOptimalNewCameraMatrix(cK, distCoeffs, cv::Size(cImg.cols, cImg.rows), 0, cv::Size(), &roi); - - cv::Mat undistorted = cv::Mat::zeros(cImg.rows, cImg.cols, cImg.type()); - cv::undistort(cImg, undistorted, cK, distCoeffs, newK); - - image = imageToTensor(undistorted); - K = floatNxNMatToTensor(newK); - }else{ - roi = cv::Rect(0, 0, cImg.cols, cImg.rows); - image = imageToTensor(cImg); - } - - // Crop to ROI - image = image.index({Slice(roi.y, roi.y + roi.height), Slice(roi.x, roi.x + roi.width), Slice()}); - - // Update parameters - height = image.size(0); - width = image.size(1); - fx = K[0][0].item(); - fy = K[1][1].item(); - cx = K[0][2].item(); - cy = K[1][2].item(); -} - -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 || p1 != 0.0f || p2 != 0.0f; -} - -std::vector Camera::undistortionParameters(){ - std::vector p = { k1, k2, p1, p2, k3, 0.0f, 0.0f, 0.0f }; - return p; -} - -void Camera::scaleOutputResolution(float scaleFactor){ - fx = fx * scaleFactor; - fy = fy * scaleFactor; - cx = cx * scaleFactor; - cy = cy * scaleFactor; - height = static_cast(static_cast(height) * scaleFactor); - width = static_cast(static_cast(width) * scaleFactor); -} - -std::tuple, Camera *> InputData::getCameras(bool validate, const std::string &valImage){ - if (!validate) return std::make_tuple(cameras, nullptr); - else{ - size_t valIdx = -1; - std::srand(42); - - if (valImage == "random"){ - valIdx = std::rand() % cameras.size(); - }else{ - for (size_t i = 0; i < cameras.size(); i++){ - if (fs::path(cameras[i].filePath).filename().string() == valImage){ - valIdx = i; - break; - } - } - if (valIdx == -1) throw std::runtime_error(valImage + " not in the list of cameras"); - } - - std::vector cams; - Camera *valCam = nullptr; - - for (size_t i = 0; i < cameras.size(); i++){ - if (i != valIdx) cams.push_back(cameras[i]); - else valCam = &cameras[i]; - } - - return std::make_tuple(cams, valCam); - } -} - } \ No newline at end of file diff --git a/nerfstudio.hpp b/nerfstudio.hpp index 32f43f35..9ceb3524 100644 --- a/nerfstudio.hpp +++ b/nerfstudio.hpp @@ -4,9 +4,8 @@ #include #include #include -#include #include -#include +#include "input_data.hpp" #include "vendor/json/json_fwd.hpp" using json = nlohmann::json; @@ -40,62 +39,9 @@ namespace ns{ void to_json(json &j, const Transforms &t); void from_json(const json& j, Transforms &t); - enum CameraType { Perspective }; - struct Camera{ - int width = 0; - int height = 0; - float fx = 0; - float fy = 0; - float cx = 0; - float cy = 0; - float k1 = 0; - float k2 = 0; - float k3 = 0; - float p1 = 0; - float p2 = 0; - torch::Tensor camToWorld; - std::string filePath = ""; - CameraType cameraType = CameraType::Perspective; - - 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), - k1(k1), k2(k2), k3(k3), p1(p1), p2(p2), - camToWorld(camToWorld), filePath(filePath) {} - - torch::Tensor getIntrinsicsMatrix(); - bool hasDistortionParameters(); - std::vector undistortionParameters(); - void scaleOutputResolution(float scaleFactor); - torch::Tensor getImage(int downscaleFactor); - - void loadImage(float downscaleFactor); - torch::Tensor K; - torch::Tensor image; - - std::unordered_map imagePyramids; - }; - Transforms readTransforms(const std::string &filename); - torch::Tensor posesFromTransforms(const Transforms &t); - std::tuple autoOrientAndCenterPoses(const torch::Tensor &poses); - - torch::Tensor rotationMatrix(const torch::Tensor &a, const torch::Tensor &b); - struct Points{ - torch::Tensor xyz; - torch::Tensor rgb; - }; - struct InputData{ - std::vector cameras; - float scaleFactor; - torch::Tensor transformMatrix; - Points points; - - std::tuple, Camera *> getCameras(bool validate, const std::string &valImage = "random"); - }; InputData inputDataFromNerfStudio(const std::string &projectRoot); } diff --git a/opensplat.cpp b/opensplat.cpp index 539f0a4f..234038a5 100644 --- a/opensplat.cpp +++ b/opensplat.cpp @@ -1,7 +1,7 @@ #include #include "vendor/json/json.hpp" #include "opensplat.hpp" -#include "point_io.hpp" +#include "input_data.hpp" #include "utils.hpp" #include "cv_utils.hpp" #include "vendor/cxxopts.hpp" @@ -85,28 +85,30 @@ int main(int argc, char *argv[]){ } try{ - ns::InputData inputData = ns::inputDataFromNerfStudio(projectRoot); - for (ns::Camera &cam : inputData.cameras){ + InputData inputData = inputDataFromX(projectRoot); + for (Camera &cam : inputData.cameras){ cam.loadImage(downScaleFactor); } + + // Withhold a validation camera if necessary auto t = inputData.getCameras(validate, valImage); - std::vector cams = std::get<0>(t); - ns::Camera *valCam = std::get<1>(t); + std::vector cams = std::get<0>(t); + Camera *valCam = std::get<1>(t); - ns::Model model(inputData.points, + Model model(inputData.points, cams.size(), numDownscales, resolutionSchedule, shDegree, shDegreeInterval, refineEvery, warmupLength, resetAlphaEvery, stopSplitAt, densifyGradThresh, densifySizeThresh, stopScreenSizeAt, splitScreenSize, numIters, device); - InfiniteRandomIterator camsIter(cams); + InfiniteRandomIterator camsIter(cams); int imageSize = -1; for (size_t step = 1; step <= numIters; step++){ - ns::Camera cam = camsIter.next(); + Camera cam = camsIter.next(); model.optimizersZeroGrad(); diff --git a/point_io.cpp b/point_io.cpp index 89078c6b..bf918249 100644 --- a/point_io.cpp +++ b/point_io.cpp @@ -101,6 +101,7 @@ PointSet *readPointSet(const std::string &filename) { PointSet *r; const fs::path p(filename); if (p.extension().string() == ".ply") r = fastPlyReadPointSet(filename); + else if (p.extension().string() == ".bin") r = colmapReadPointSet(filename); else r = pdalReadPointSet(filename); return r; @@ -357,6 +358,41 @@ PointSet *pdalReadPointSet(const std::string &filename) { #endif } +PointSet *colmapReadPointSet(const std::string &filename){ + + std::ifstream reader(filename, std::ios::binary); + if (!reader.is_open()) throw std::runtime_error("Cannot open " + filename); + + auto *r = new PointSet(); + size_t numPoints = readBinary(reader); + std::cout << "Reading " << numPoints << " points" << std::endl; + + r->points.resize(numPoints); + r->colors.resize(numPoints); + + for (size_t i = 0; i < numPoints; i++){ + readBinary(reader); // point ID + + r->points[i][0] = readBinary(reader); + r->points[i][1] = readBinary(reader); + r->points[i][2] = readBinary(reader); + r->colors[i][0] = readBinary(reader); + r->colors[i][1] = readBinary(reader); + r->colors[i][2] = readBinary(reader); + + readBinary(reader); // error + size_t trackLen = readBinary(reader); + for (size_t j = 0; j < trackLen; j++){ + readBinary(reader); // imageId + readBinary(reader); // point2D Idx + } + } + + reader.close(); + + return r; +} + void checkHeader(std::ifstream &reader, const std::string &prop) { std::string line; std::getline(reader, line); diff --git a/point_io.hpp b/point_io.hpp index d63eb60e..a2d79f75 100644 --- a/point_io.hpp +++ b/point_io.hpp @@ -103,8 +103,16 @@ size_t getVertexCount(const std::string &line); inline void checkHeader(std::ifstream &reader, const std::string &prop); inline bool hasHeader(const std::string &line, const std::string &prop); +template +inline T readBinary(std::ifstream &s){ + T data; + s.read(reinterpret_cast(&data), sizeof(T)); + return data; +} + PointSet *fastPlyReadPointSet(const std::string &filename); PointSet *pdalReadPointSet(const std::string &filename); +PointSet *colmapReadPointSet(const std::string &filename); PointSet *readPointSet(const std::string &filename); void fastPlySavePointSet(PointSet &pSet, const std::string &filename); diff --git a/tensor_math.cpp b/tensor_math.cpp new file mode 100644 index 00000000..41cec6ce --- /dev/null +++ b/tensor_math.cpp @@ -0,0 +1,66 @@ +#include "tensor_math.hpp" + +using namespace torch::indexing; + +torch::Tensor quatToRotMat(const torch::Tensor &quat){ + auto u = torch::unbind(torch::nn::functional::normalize(quat, torch::nn::functional::NormalizeFuncOptions().dim(-1)), -1); + torch::Tensor w = u[0]; + torch::Tensor x = u[1]; + torch::Tensor y = u[2]; + torch::Tensor z = u[3]; + return torch::stack({ + torch::stack({ + 1.0 - 2.0 * (y.pow(2) + z.pow(2)), + 2.0 * (x * y - w * z), + 2.0 * (x * z + w * y) + }, -1), + torch::stack({ + 2.0 * (x * y + w * z), + 1.0 - 2.0 * (x.pow(2) + z.pow(2)), + 2.0 * (y * z - w * x) + }, -1), + torch::stack({ + 2.0 * (x * z - w * y), + 2.0 * (y * z + w * x), + 1.0 - 2.0 * (x.pow(2) + y.pow(2)) + }, -1) + }, -2); + +} + +std::tuple autoOrientAndCenterPoses(const torch::Tensor &poses){ + // Center at mean and orient up + torch::Tensor origins = poses.index({"...", Slice(None, 3), 3}); + torch::Tensor translation = torch::mean(origins, 0); + torch::Tensor up = torch::mean(poses.index({Slice(), Slice(None, 3), 1}), 0); + up = up / up.norm(); + + torch::Tensor rotation = rotationMatrix(up, torch::tensor({0, 0, 1}, torch::kFloat32)); + torch::Tensor transform = torch::cat({rotation, torch::matmul(rotation, -translation.index({"...", None}))}, -1); + torch::Tensor orientedPoses = torch::matmul(transform, poses); + return std::make_tuple(orientedPoses, transform); +} + + +torch::Tensor rotationMatrix(const torch::Tensor &a, const torch::Tensor &b){ + // Rotation matrix that rotates vector a to vector b + torch::Tensor a1 = a / a.norm(); + torch::Tensor b1 = b / b.norm(); + torch::Tensor v = torch::cross(a1, b1); + torch::Tensor c = torch::dot(a1, b1); + const float EPS = 1e-8; + if (c.item() < -1 + EPS){ + torch::Tensor eps = (torch::rand(3) - 0.5f) * 0.01f; + return rotationMatrix(a1 + eps, b1); + } + torch::Tensor s = v.norm(); + torch::Tensor skew = torch::zeros({3, 3}, torch::kFloat32); + skew[0][1] = -v[2]; + skew[0][2] = v[1]; + skew[1][0] = v[2]; + skew[1][2] = -v[0]; + skew[2][0] = -v[1]; + skew[2][1] = v[0]; + + return torch::eye(3) + skew + torch::matmul(skew, skew * ((1 - c) / (s.pow(2) + EPS))); +} \ No newline at end of file diff --git a/tensor_math.hpp b/tensor_math.hpp new file mode 100644 index 00000000..12b0849d --- /dev/null +++ b/tensor_math.hpp @@ -0,0 +1,12 @@ +#ifndef TENSOR_MATH_H +#define TENSOR_MATH_H + +#include +#include + +torch::Tensor quatToRotMat(const torch::Tensor &quat); +std::tuple autoOrientAndCenterPoses(const torch::Tensor &poses); +torch::Tensor rotationMatrix(const torch::Tensor &a, const torch::Tensor &b); + + +#endif \ No newline at end of file