From 69c3d37d3a45785cdff81e5f6ce3942d4685e13f Mon Sep 17 00:00:00 2001
From: steffen <urbste@googlemail.com>
Date: Tue, 28 Nov 2023 20:50:00 +0100
Subject: [PATCH] added some more wrappers

Change-Id: Iee59a9838af3c4c3d86670f8c7aeeecea6eb901a
---
 README.md                    |  2 +-
 setup.py                     |  2 +-
 src/pytheia/sfm/sfm.cc       | 22 ++++++++++++++++++----
 src/theia/sfm/sfm_wrapper.cc | 16 ++++++++++++++++
 src/theia/sfm/sfm_wrapper.h  |  9 +++++++++
 5 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 0ef3e1e..53a7a5e 100644
--- a/README.md
+++ b/README.md
@@ -166,7 +166,7 @@ opts.loss_function_type = pt.sfm.LossFunctionType.HUBER
 
 res = BundleAdjustReconstruction(opts, recon)
 res = BundleAdjustPartialReconstruction(opts, {view_ids}, {track_ids}, recon)
-res = BundleAdjustPartialViewConstant(opts, {var_view_ids}, {const_view_ids}, recon)
+res = BundleAdjustPartialViewsConstant(opts, {var_view_ids}, {const_view_ids}, recon)
 
 # optimize absolute pose on normalized 2D 3D correspondences
 res = pt.sfm.OptimizeAbsolutePoseOnNormFeatures(
diff --git a/setup.py b/setup.py
index 23a4b99..00fd2f5 100644
--- a/setup.py
+++ b/setup.py
@@ -65,7 +65,7 @@ def create_package():
 
 setuptools.setup(
     name='pytheia',
-    version='0.2.0',
+    version='0.2.1',
     description='A performant Structure from Motion library for Python',
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',
diff --git a/src/pytheia/sfm/sfm.cc b/src/pytheia/sfm/sfm.cc
index d165143..765c35c 100644
--- a/src/pytheia/sfm/sfm.cc
+++ b/src/pytheia/sfm/sfm.cc
@@ -1401,13 +1401,27 @@ void pytheia_sfm_classes(py::module& m) {
       //&theia::Reconstruction::GetSubReconstructionWrapper)
       ;
 
+  // Reconstruction Estimator Helpers
   m.def("SetUnderconstrainedTracksToUnestimated",
-        theia::SetUnderconstrainedTracksToUnestimated);
+        &theia::SetUnderconstrainedTracksToUnestimated);
   m.def("SetUnderconstrainedViewsToUnestimated",
-        theia::SetUnderconstrainedViewsToUnestimated);
-
+        &theia::SetUnderconstrainedViewsToUnestimated);
+  m.def("NumEstimatedViews",
+        &theia::NumEstimatedViews);
+  m.def("NumEstimatedTracks",
+        &theia::NumEstimatedTracks);
+  m.def("SetReconstructionFromEstimatedPoses",
+        &theia::SetReconstructionFromEstimatedPoses);
+  m.def("CreateEstimatedSubreconstruction",
+        &theia::CreateEstimatedSubreconstruction);
+  m.def("RelativeRotationsFromViewGraph",
+        &theia::RelativeRotationsFromViewGraph);
+  m.def("GetEstimatedViewsFromReconstruction",
+        &theia::GetEstimatedViewsFromReconstructionWrapper);    
+  m.def("GetEstimatedTracksFromReconstruction",
+        &theia::GetEstimatedTracksFromReconstructionWrapper);
+        
   // Reconstruction Estimator
-
   // TwoViewInfo
   py::class_<theia::TwoViewInfo>(m, "TwoViewInfo")
       .def(py::init<>())
diff --git a/src/theia/sfm/sfm_wrapper.cc b/src/theia/sfm/sfm_wrapper.cc
index c6fd290..a868e42 100644
--- a/src/theia/sfm/sfm_wrapper.cc
+++ b/src/theia/sfm/sfm_wrapper.cc
@@ -55,4 +55,20 @@ int SetOutlierTracksToUnestimatedWrapper(
   return num_features_rm;
 }
 
+// Outputs the ViewId of all estimated views in the reconstruction.
+std::unordered_set<ViewId> GetEstimatedViewsFromReconstructionWrapper(
+    const Reconstruction& reconstruction) {
+  std::unordered_set<ViewId> estimated_views;
+  theia::GetEstimatedViewsFromReconstruction(reconstruction, &estimated_views);
+  return estimated_views;
+}
+
+// Outputs the TrackId of all estimated tracks in the reconstruction.
+std::unordered_set<TrackId> GetEstimatedTracksFromReconstructionWrapper(
+    const Reconstruction& reconstruction) {
+  std::unordered_set<TrackId> estimated_tracks;
+  theia::GetEstimatedTracksFromReconstruction(reconstruction, &estimated_tracks);
+  return estimated_tracks;
+}
+
 }  // namespace theia
diff --git a/src/theia/sfm/sfm_wrapper.h b/src/theia/sfm/sfm_wrapper.h
index f011d29..f1fdda0 100644
--- a/src/theia/sfm/sfm_wrapper.h
+++ b/src/theia/sfm/sfm_wrapper.h
@@ -14,6 +14,7 @@
 #include "theia/sfm/set_camera_intrinsics_from_priors.h"
 #include "theia/sfm/set_outlier_tracks_to_unestimated.h"
 #include "theia/sfm/undistort_image.h"
+#include "theia/sfm/reconstruction_estimator_utils.h"
 
 namespace theia {
 
@@ -43,4 +44,12 @@ int SetOutlierTracksToUnestimatedWrapper(
     const double min_triangulation_angle_degrees,
     Reconstruction& reconstruction);
 
+// Outputs the ViewId of all estimated views in the reconstruction.
+std::unordered_set<ViewId> GetEstimatedViewsFromReconstructionWrapper(
+    const Reconstruction& reconstruction);
+
+// Outputs the TrackId of all estimated tracks in the reconstruction.
+std::unordered_set<TrackId> GetEstimatedTracksFromReconstructionWrapper(
+    const Reconstruction& reconstruction);
+
 }  // namespace theia