Update h5_to_db.py

h5_to_db.py

@@ -69,18 +69,39 @@ def create_camera(db, image_path, camera_model):
     return db.add_camera(model, width, height, param_arr)
 
 
+
 def add_keypoints(db, h5_path, image_path, img_ext, camera_model, single_camera=True):
+    import h5py
+    import numpy as np
+    import os
+    import glob
+    from tqdm import tqdm
+    from PIL import Image
+    
     keypoint_f = h5py.File(os.path.join(h5_path, 'keypoints.h5'), 'r')
     camera_id = None
     fname_to_id = {}
     
-    import glob
     all_images = {}
     for ext in ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG']:
         for img_file in glob.glob(os.path.join(image_path, f'*{ext}')):
             base_name = os.path.splitext(os.path.basename(img_file))[0]
             all_images[base_name] = img_file
     
+    # ★★★ カメラモデルのマッピング ★★★
+    camera_model_ids = {
+        'SIMPLE_PINHOLE': 0,
+        'PINHOLE': 1,
+        'SIMPLE_RADIAL': 2,
+        'RADIAL': 3,
+        'OPENCV': 4,
+    }
+    
+    if camera_model not in camera_model_ids:
+        raise ValueError(f"Unknown camera model: {camera_model}")
+    
+    model_id = camera_model_ids[camera_model]
+    
     for filename in tqdm(list(keypoint_f.keys())):
         keypoints = keypoint_f[filename][()]
         
@@ -92,22 +113,35 @@ def add_keypoints(db, h5_path, image_path, img_ext, camera_model, single_camera=
         
         # カメラは1回だけ作成（single_camera=Trueの場合）
         if camera_id is None:
-            from PIL import Image
             img = Image.open(path)
             width, height = img.size
             
             focal_length = float(max(width, height))
             cx = float(width / 2.0)
             cy = float(height / 2.0)
-            k = 0.0
             
-            print(f"カメラ作成: f={focal_length}, cx={cx}, cy={cy}, k={k}")
+            # ★★★ カメラモデルに応じてパラメータを設定 ★★★
+            if camera_model == 'PINHOLE':
+                # PINHOLE: [fx, fy, cx, cy]
+                params = np.array([focal_length, focal_length, cx, cy], dtype=np.float64)
+                print(f"カメラ作成: PINHOLE, f={focal_length}, cx={cx}, cy={cy}")
+            elif camera_model == 'SIMPLE_PINHOLE':
+                # SIMPLE_PINHOLE: [f, cx, cy]
+                params = np.array([focal_length, cx, cy], dtype=np.float64)
+                print(f"カメラ作成: SIMPLE_PINHOLE, f={focal_length}, cx={cx}, cy={cy}")
+            elif camera_model == 'SIMPLE_RADIAL':
+                # SIMPLE_RADIAL: [f, cx, cy, k]
+                k = 0.0
+                params = np.array([focal_length, cx, cy, k], dtype=np.float64)
+                print(f"カメラ作成: SIMPLE_RADIAL, f={focal_length}, cx={cx}, cy={cy}, k={k}")
+            else:
+                raise ValueError(f"Unsupported camera model: {camera_model}")
             
             camera_id = db.add_camera(
-                model=2,  # SIMPLE_RADIAL
+                model=model_id,  # ★★★ 引数に基づいて設定 ★★★
                 width=width,
                 height=height,
-                params=np.array([focal_length, cx, cy, k], dtype=np.float64)
+                params=params
             )
         
         image_id = db.add_image(fname_with_ext, camera_id)
@@ -115,7 +149,8 @@ def add_keypoints(db, h5_path, image_path, img_ext, camera_model, single_camera=
         db.add_keypoints(image_id, keypoints)
     
     return fname_to_id
-    
+
+
 
 def add_matches(db, h5_path, fname_to_id):
     """