Update h5_to_db.py

h5_to_db.py

@@ -71,25 +71,45 @@ def create_camera(db, image_path, camera_model):
 
 def add_keypoints(db, h5_path, image_path, img_ext, camera_model, single_camera = True):
     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
+    
     for filename in tqdm(list(keypoint_f.keys())):
         keypoints = keypoint_f[filename][()]
-
-        fname_with_ext = filename + img_ext
-        path = os.path.join(image_path, fname_with_ext)
-        if not os.path.isfile(path):
-            raise IOError(f'Invalid image path {path}')
-
+        
+        # 拡張子なしのファイル名で検索
+        if filename not in all_images:
+            raise IOError(f'Image not found for key: {filename}')
+        
+        path = all_images[filename]  # 実際のフルパス
+        fname_with_ext = os.path.basename(path)  # 拡張子付きファイル名
+        
         if camera_id is None or not single_camera:
-            camera_id = create_camera(db, path, camera_model)
+            from PIL import Image
+            img = Image.open(path)
+            width, height = img.size
+            
+            camera_id = db.add_camera(
+                model=2,  # SIMPLE_RADIAL
+                width=width,
+                height=height,
+                params=np.array([0.0, 0.0, 0.0, 0.0])  # COLMAPが最適化
+            )
+        
         image_id = db.add_image(fname_with_ext, camera_id)
         fname_to_id[filename] = image_id
-
         db.add_keypoints(image_id, keypoints)
-
+    
     return fname_to_id
+    
 
 def add_matches(db, h5_path, fname_to_id):
     match_file = h5py.File(os.path.join(h5_path, 'matches.h5'), 'r')