elev estimate synced from GitHub

elevation_estimate/estimate_wild_imgs.py

@@ -0,0 +1,10 @@
+import os.path as osp
+from .utils.elev_est_api import elev_est_api
+
+def estimate_elev(root_dir):
+    img_dir = osp.join(root_dir, "stage2_8")
+    img_paths = []
+    for i in range(4):
+        img_paths.append(f"{img_dir}/0_{i}.png")
+    elev = elev_est_api(img_paths)
+    return elev

{one2345_elev_est → elevation_estimate}/pyproject.toml

@@ -1,5 +1,5 @@
 [project]
-name = "one2345_elev_est"
+name = "elevation_estimate"
 version = "0.1"
 
 [tool.setuptools.packages.find]

{one2345_elev_est/oee → elevation_estimate}/utils/elev_est_api.py

@@ -1,21 +1,19 @@
-import matplotlib.pyplot as plt
-import warnings
-
-import numpy as np
-import cv2
 import os
+import cv2
+import numpy as np
 import os.path as osp
 import imageio
 from copy import deepcopy
 
 import loguru
 import torch
-from oee.models.loftr import LoFTR, default_cfg
 import matplotlib.cm as cm
+import matplotlib.pyplot as plt
 
-from oee.utils import plt_utils
-from oee.utils.plotting import make_matching_figure
-from oee.utils.utils3d import rect_to_img, canonical_to_camera, calc_pose
+from ..loftr import LoFTR, default_cfg
+from . import plt_utils
+from .plotting import make_matching_figure
+from .utils3d import rect_to_img, canonical_to_camera, calc_pose
 
 
 class ElevEstHelper:
@@ -28,13 +26,10 @@ class ElevEstHelper:
             _default_cfg = deepcopy(default_cfg)
             _default_cfg['coarse']['temp_bug_fix'] = True  # set to False when using the old ckpt
             matcher = LoFTR(config=_default_cfg)
-            ckpt_path = "weights/indoor_ds_new.ckpt"
+            current_dir = os.path.dirname(os.path.abspath(__file__))
+            ckpt_path = os.path.join(current_dir, "weights/indoor_ds_new.ckpt")
             if not osp.exists(ckpt_path):
-                loguru.logger.info("Downloading feature matcher...")
-                os.makedirs("weights", exist_ok=True)
-                import gdown
-                gdown.cached_download(url="https://drive.google.com/uc?id=19s3QvcCWQ6g-N1PrYlDCg-2mOJZ3kkgS",
-                                      path=ckpt_path)
+                raise FileNotFoundError(f"Checkpoint not found at {ckpt_path}")
             matcher.load_state_dict(torch.load(ckpt_path)['state_dict'])
             matcher = matcher.eval().cuda()
             cls._feature_matcher = matcher

one2345_elev_est/tools/estimate_wild_imgs.py

@@ -1,35 +0,0 @@
-import tqdm
-import imageio
-import json
-import os.path as osp
-import os
-
-from oee.utils import plt_utils
-from oee.utils.elev_est_api import elev_est_api
-
-
-def visualize(img_paths, elev):
-    imgs = [imageio.imread_v2(img_path) for img_path in img_paths]
-    plt_utils.image_grid(imgs, 2, 2, label=f"elev={elev}")
-
-
-def estimate_elev(root_dir):
-    # root_dir = "/home/linghao/Datasets/objaverse-processed/zero12345_img/wild"
-    # dataset = "supp_fail"
-    # root_dir = "/home/chao/chao/OpenComplete/zero123/zero123/gradio_tmp/"
-    # obj_names = sorted(os.listdir(root_dir))
-    # results = {}
-    # for obj_name in tqdm.tqdm(obj_names):
-    img_dir = osp.join(root_dir, "stage2_8")
-    img_paths = []
-    for i in range(4):
-        img_paths.append(f"{img_dir}/0_{i}.png")
-    elev = elev_est_api(img_paths)
-    # visualize(img_paths, elev)
-    # results[obj_name] = elev
-    # json.dump(results, open(osp.join(root_dir, f"../{dataset}_elev.json"), "w"), indent=4)
-    return elev
-
-
-# if __name__ == '__main__':
-#     main()

one2345_elev_est/tools/example.py

@@ -1,38 +0,0 @@
-import imageio
-import numpy as np
-
-from oee.utils import plt_utils
-from oee.utils.elev_est_api import elev_est_api
-import argparse
-
-
-def visualize(img_paths, elev):
-    imgs = [imageio.imread_v2(img_path) for img_path in img_paths]
-    plt_utils.image_grid(imgs, 2, 2, label=f"elev={elev}")
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--img_paths", type=str, nargs=4, help="image paths",
-                        default=["assets/example_data/0_0.png",
-                                 "assets/example_data/0_1.png",
-                                 "assets/example_data/0_2.png",
-                                 "assets/example_data/0_3.png"])
-    parser.add_argument("--min_elev", type=float, default=30, help="min elevation")
-    parser.add_argument("--max_elev", type=float, default=150, help="max elevation")
-    parser.add_argument("--dbg", default=False, action="store_true", help="debug mode")
-    parser.add_argument("--K_path", type=str, default=None, help="path to K")
-    args = parser.parse_args()
-
-    if args.K_path is not None:
-        K = np.loadtxt(args.K_path)
-    else:
-        K = None
-
-    elev = elev_est_api(args.img_paths, args.min_elev, args.max_elev, K, args.dbg)
-
-    visualize(args.img_paths, elev)
-
-
-if __name__ == '__main__':
-    main()