Upload folder using huggingface_hub

geocalib/extractor.py

@@ -22,14 +22,9 @@ class GeoCalib(nn.Module):
             weights (str): trained variant, "pinhole" (default) or "distorted".
         """
         super().__init__()
-        if weights == "pinhole":
-            url = "https://github.com/cvg/GeoCalib/releases/download/v1.0/geocalib-pinhole.tar"
-        elif weights == "distorted":
-            url = (
-                "https://github.com/cvg/GeoCalib/releases/download/v1.0/geocalib-simple_radial.tar"
-            )
-        else:
+        if weights not in {"pinhole", "distorted"}:
             raise ValueError(f"Unknown weights: {weights}")
+        url = f"https://github.com/cvg/GeoCalib/releases/download/v1.0/geocalib-{weights}.tar"
 
         # load checkpoint
         model_dir = f"{torch.hub.get_dir()}/geocalib"

gradio_app.py

@@ -8,7 +8,7 @@ import numpy as np
 import spaces
 import torch
 
-from geocalib import viz2d
+from geocalib import logger, viz2d
 from geocalib.camera import camera_models
 from geocalib.extractor import GeoCalib
 from geocalib.perspective_fields import get_perspective_field
@@ -77,7 +77,9 @@ def format_output(results):
 @spaces.GPU(duration=10)
 def inference(img, camera_model):
     out = model.calibrate(img.to(device), camera_model=camera_model)
-    save_keys = ["camera", "gravity"] + [f"{k}_uncertainty" for k in ["roll", "pitch", "vfov", "focal"]]
+    save_keys = ["camera", "gravity"] + [
+        f"{k}_uncertainty" for k in ["roll", "pitch", "vfov", "focal"]
+    ]
     res = {k: v.cpu() for k, v in out.items() if k in save_keys}
     # not converting to numpy results in gpu abort
     res["up_confidence"] = out["up_confidence"].cpu().numpy()
@@ -100,10 +102,9 @@ def process_results(
         raise gr.Error("Please upload an image first.")
 
     img = model.load_image(image_path)
-    print("Running inference...")
     start = time()
     inference_result = inference(img, camera_model)
-    print(f"Done ({time() - start:.2f}s)")
+    logger.info(f"Calibration took {time() - start:.2f} sec. ({camera_model})")
     inference_result["image"] = img.cpu()
 
     if inference_result is None:
@@ -158,7 +159,9 @@ def update_plot(
         viz2d.plot_confidences([torch.tensor(inference_result["up_confidence"][0])], axes=[ax[0]])
 
     if plot_latitude_confidence:
-        viz2d.plot_confidences([torch.tensor(inference_result["latitude_confidence"][0])], axes=[ax[0]])
+        viz2d.plot_confidences(
+            [torch.tensor(inference_result["latitude_confidence"][0])], axes=[ax[0]]
+        )
 
     fig.canvas.draw()
     img = np.array(fig.canvas.renderer.buffer_rgba())

siclib/models/extractor.py

@@ -22,14 +22,9 @@ class GeoCalib(nn.Module):
             weights (str, optional): Weights to load. Defaults to "pinhole".
         """
         super().__init__()
-        if weights == "pinhole":
-            url = "https://github.com/cvg/GeoCalib/releases/download/v1.0/geocalib-pinhole.tar"
-        elif weights == "distorted":
-            url = (
-                "https://github.com/cvg/GeoCalib/releases/download/v1.0/geocalib-simple_radial.tar"
-            )
-        else:
+        if weights not in {"pinhole", "distorted"}:
             raise ValueError(f"Unknown weights: {weights}")
+        url = f"https://github.com/cvg/GeoCalib/releases/download/v1.0/geocalib-{weights}.tar"
 
         # load checkpoint
         model_dir = f"{torch.hub.get_dir()}/geocalib"