-
Notifications
You must be signed in to change notification settings - Fork 90
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #33 from pierotofy/colmap
COLMAP project support
- Loading branch information
Showing
15 changed files
with
527 additions
and
268 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,154 @@ | ||
| #include <filesystem> | ||
| #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<uint64_t>(camf); | ||
| std::vector<Camera> cameras(numCameras); | ||
|
|
||
| std::unordered_map<uint32_t, Camera *> camMap; | ||
|
|
||
| for (size_t i = 0; i < numCameras; i++) { | ||
| Camera *cam = &cameras[i]; | ||
|
|
||
| cam->id = readBinary<uint32_t>(camf); | ||
|
|
||
| CameraModel model = static_cast<CameraModel>(readBinary<int>(camf)); // model ID | ||
| cam->width = readBinary<uint64_t>(camf); | ||
| cam->height = readBinary<uint64_t>(camf); | ||
|
|
||
| if (model == SimplePinhole){ | ||
| cam->fx = readBinary<double>(camf); | ||
| cam->fy = cam->fx; | ||
| cam->cx = readBinary<double>(camf); | ||
| cam->cy = readBinary<double>(camf); | ||
| }else if (model == Pinhole){ | ||
| cam->fx = readBinary<double>(camf); | ||
| cam->fy = readBinary<double>(camf); | ||
| cam->cx = readBinary<double>(camf); | ||
| cam->cy = readBinary<double>(camf); | ||
| }else if (model == OpenCV){ | ||
| cam->fx = readBinary<double>(camf); | ||
| cam->fy = readBinary<double>(camf); | ||
| cam->cx = readBinary<double>(camf); | ||
| cam->cy = readBinary<double>(camf); | ||
| cam->k1 = readBinary<double>(camf); | ||
| cam->k2 = readBinary<double>(camf); | ||
| cam->p1 = readBinary<double>(camf); | ||
| cam->p2 = readBinary<double>(camf); | ||
| }else{ | ||
| throw std::runtime_error("Unsupported camera model: " + std::to_string(model)); | ||
| } | ||
|
|
||
| camMap[cam->id] = cam; | ||
| } | ||
|
|
||
| camf.close(); | ||
|
|
||
|
|
||
| size_t numImages = readBinary<uint64_t>(imgf); | ||
| torch::Tensor unorientedPoses = torch::zeros({static_cast<long int>(numImages), 4, 4}, torch::kFloat32); | ||
|
|
||
| for (size_t i = 0; i < numImages; i++){ | ||
| readBinary<uint32_t>(imgf); // imageId | ||
|
|
||
| torch::Tensor qVec = torch::tensor({ | ||
| readBinary<double>(imgf), | ||
| readBinary<double>(imgf), | ||
| readBinary<double>(imgf), | ||
| readBinary<double>(imgf) | ||
| }, torch::kFloat32); | ||
| torch::Tensor R = quatToRotMat(qVec); | ||
| torch::Tensor T = torch::tensor({ | ||
| { readBinary<double>(imgf) }, | ||
| { readBinary<double>(imgf) }, | ||
| { readBinary<double>(imgf) } | ||
| }, torch::kFloat32); | ||
|
|
||
| torch::Tensor Rinv = R.transpose(0, 1); | ||
| torch::Tensor Tinv = torch::matmul(-Rinv, T); | ||
|
|
||
| uint32_t camId = readBinary<uint32_t>(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<uint64_t>(imgf); | ||
| for (size_t j = 0; j < numPoints2D; j++){ | ||
| readBinary<double>(imgf); // x | ||
| readBinary<double>(imgf); // y | ||
| readBinary<uint64_t>(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<float>(); | ||
| 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; | ||
| } | ||
|
|
||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,17 @@ | ||
| #ifndef COLMAP_H | ||
| #define COLMAP_H | ||
|
|
||
| #include <fstream> | ||
| #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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,151 @@ | ||
| #include <filesystem> | ||
| #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<float>(cImg.rows) / static_cast<float>(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<float> 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<float>(); | ||
| fy = K[1][1].item<float>(); | ||
| cx = K[0][2].item<float>(); | ||
| cy = K[1][2].item<float>(); | ||
| } | ||
|
|
||
| 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<float> Camera::undistortionParameters(){ | ||
| std::vector<float> 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<int>(static_cast<float>(height) * scaleFactor); | ||
| width = static_cast<int>(static_cast<float>(width) * scaleFactor); | ||
| } | ||
|
|
||
| std::tuple<std::vector<Camera>, 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<Camera> 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); | ||
| } | ||
| } |
Oops, something went wrong.