Update database with q/t

kaiyi98 · Jan 4, 2022 · f23666b · f23666b
1 parent 55e6cde
commit f23666b
Showing 1 changed file with 13 additions and 5 deletions.
diff --git a/hloc/utils/database.py b/hloc/utils/database.py
@@ -79,7 +79,9 @@
     config INTEGER NOT NULL,
     F BLOB,
     E BLOB,
-    H BLOB)
+    H BLOB,
+    qvec BLOB,
+    tvec BLOB)
 """
 
 CREATE_KEYPOINTS_TABLE = """CREATE TABLE IF NOT EXISTS keypoints (
@@ -171,7 +173,8 @@ def add_camera(self, model, width, height, params,
         return cursor.lastrowid
 
     def add_image(self, name, camera_id,
-                  prior_q=np.zeros(4), prior_t=np.zeros(3), image_id=None):
+                  prior_q=np.full(4, np.NaN), prior_t=np.full(3, np.NaN),
+                  image_id=None):
         cursor = self.execute(
             "INSERT INTO images VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
             (image_id, name, camera_id, prior_q[0], prior_q[1], prior_q[2],
@@ -207,7 +210,9 @@ def add_matches(self, image_id1, image_id2, matches):
             (pair_id,) + matches.shape + (array_to_blob(matches),))
 
     def add_two_view_geometry(self, image_id1, image_id2, matches,
-                              F=np.eye(3), E=np.eye(3), H=np.eye(3), config=2):
+                              F=np.eye(3), E=np.eye(3), H=np.eye(3),
+                              qvec=np.array([1.0, 0.0, 0.0, 0.0]),
+                              tvec=np.zeros(3), config=2):
         assert(len(matches.shape) == 2)
         assert(matches.shape[1] == 2)
 
@@ -219,10 +224,13 @@ def add_two_view_geometry(self, image_id1, image_id2, matches,
         F = np.asarray(F, dtype=np.float64)
         E = np.asarray(E, dtype=np.float64)
         H = np.asarray(H, dtype=np.float64)
+        qvec = np.asarray(qvec, dtype=np.float64)
+        tvec = np.asarray(tvec, dtype=np.float64)
         self.execute(
-            "INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+            "INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
             (pair_id,) + matches.shape + (array_to_blob(matches), config,
-             array_to_blob(F), array_to_blob(E), array_to_blob(H)))
+             array_to_blob(F), array_to_blob(E), array_to_blob(H),
+             array_to_blob(qvec), array_to_blob(tvec)))
 
 
 def example_usage():