From 0c9cdf111ff995263e5b3fbf6d07f78ce4c3dcd5 Mon Sep 17 00:00:00 2001
From: Paidi Akileswar <155959310+paidiakileswar@users.noreply.github.com>
Date: Tue, 30 Jul 2024 23:00:54 +0530
Subject: [PATCH] Adding Code to Tweak b/w cameras model & single camera

---
 pixsfm/refine_hloc.py | 29 +++++++++++++++++++++++++----
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/pixsfm/refine_hloc.py b/pixsfm/refine_hloc.py
index 9918f00..530d9cd 100644
--- a/pixsfm/refine_hloc.py
+++ b/pixsfm/refine_hloc.py
@@ -145,6 +145,13 @@ def reconstruction(
             reference_model_path=None, cache_path=cache_path,
             feature_manager=feature_manager, **hloc_args)
 
+def add_camera_args(parser):
+    parser.add_argument(
+        '--camera_model', type=str, default='OPENCV',
+        help='Camera model for the image reader options.')
+    parser.add_argument(
+        '--single_camera', action='store_true',
+        help='Enable single camera mode for reconstruction.')
 
 def add_hloc_args(parser):
     parser.add_argument(
@@ -170,7 +177,7 @@ def add_hloc_args(parser):
         help='Output HDF5 file where the refine keypoints will be written.')
     add_hloc_args(parser_ka)
     add_common_args(parser_ka)
-
+    
     parser_rec = subparsers.add_parser(
         'reconstructor', aliases=['rec', 'sfm'],
         help='Full 3D reconstruction with keypoint and bundle adjustments.')
@@ -178,6 +185,7 @@ def add_hloc_args(parser):
         '--sfm_dir', type=Path, required=True, help='Output SfM model.')
     add_hloc_args(parser_rec)
     add_common_args(parser_rec)
+    add_camera_args(parser_rec)
 
     parser_tri = subparsers.add_parser(
         'triangulator', aliases=['tri'],
@@ -197,19 +205,32 @@ def add_hloc_args(parser):
     if args.config is not None:
         conf = OmegaConf.merge(OmegaConf.load(args.config), conf)
     sfm = PixSfM(conf)
+    
+    camera_model = args.camera_model if 'camera_model' in args else 'OPENCV'
+    image_options = pycolmap.ImageReaderOptions(camera_model=camera_model)
 
     if args.command == 'keypoint_adjuster':
         logger.info("Will perform keypoint adjustment.")
+        logger.info("Received arguments for 3D triangulation: %s", args)
         sfm.refine_keypoints(
             args.output_path, args.features_path, args.image_dir,
             args.pairs_path, args.matches_path, cache_path=args.cache_path)
     elif args.command == 'reconstructor':
         logger.info("Will perform full 3D reconstruction.")
-        sfm.reconstruction(
-            args.sfm_dir, args.image_dir, args.pairs_path, args.features_path,
-            args.matches_path, cache_path=args.cache_path)
+        logger.info("Received arguments for 3D triangulation: %s", args)
+        if args.single_camera:
+            sfm.reconstruction(
+                args.sfm_dir, args.image_dir, args.pairs_path, args.features_path,
+                args.matches_path, cache_path=args.cache_path, camera_mode=pycolmap.CameraMode.SINGLE,
+                image_options=image_options)
+        else:
+            sfm.reconstruction(
+                args.sfm_dir, args.image_dir, args.pairs_path, args.features_path,
+                args.matches_path, cache_path=args.cache_path,
+                image_options=image_options)
     elif args.command == 'triangulator':
         logger.info("Will perform 3D triangulation.")
+        logger.info("Received arguments for 3D triangulation: %s", args)
         sfm.triangulation(
             args.sfm_dir, args.reference_sfm_model,
             args.image_dir, args.pairs_path, args.features_path,