wip

app.py

@@ -20,7 +20,7 @@ from surge.inference import (
     point_cloud_from_result,
     resize_to_max_size,
 )
-from surge.model.surge import SurGe
+from surge import SurGe
 from surge.utils.background import apply_background_filters
 from surge.utils.io import save_glb, save_ply
 

surge/__init__.py

@@ -0,0 +1,3 @@
+from surge.models import SurGe
+
+__all__ = ["SurGe"]

surge/cli.py

@@ -25,7 +25,7 @@ from surge.inference import (
     resize_to_max_size,
 )
 from surge.utils.background import apply_background_filters
-from surge.model.surge import SurGe
+from surge import SurGe
 from surge.utils.io import read_image, save_glb, save_ply
 
 try:

surge/inference.py

@@ -8,7 +8,7 @@ import numpy as np
 import torch
 import utils3d
 
-from surge.model.surge import SurGe
+from surge import SurGe
 from surge.utils.geometry_numpy import depth_occlusion_edge_numpy
 
 DEFAULT_CHECKPOINT = "karimknaebel/surge-large"

surge/models/__init__.py

@@ -0,0 +1,3 @@
+from surge.models.surge import SurGe
+
+__all__ = ["SurGe"]

surge/{model → models}/surge.py

@@ -13,8 +13,8 @@ from einops import rearrange
 from huggingface_hub import PyTorchModelHubMixin
 from specbuild import build
 
-from surge.modules.base import BaseEncoder
-from surge.modules.nad import NAD
+from surge.modules.encoders.base import BaseEncoder
+from surge.modules.heads.nad import NAD
 from surge.utils.geometry_torch import normalized_view_plane_uv, recover_focal_shift
 
 

surge/modules/__init__.py

@@ -1,5 +1,5 @@
 # import for registry
-import surge.modules.dinov2_encoder
-import surge.modules.dinov3_encoder
-import surge.modules.nad
-import surge.modules.timm_encoder
+import surge.modules.encoders.dinov2
+import surge.modules.encoders.dinov3
+import surge.modules.encoders.timm
+import surge.modules.heads.nad

surge/modules/encoders/__init__.py

@@ -0,0 +1,6 @@
+from surge.modules.encoders.base import BaseEncoder
+from surge.modules.encoders.dinov2 import DINOv2Encoder
+from surge.modules.encoders.dinov3 import DINOv3Encoder
+from surge.modules.encoders.timm import TimmEncoder
+
+__all__ = ["BaseEncoder", "DINOv2Encoder", "DINOv3Encoder", "TimmEncoder"]

surge/modules/{dinov2_encoder.py → encoders/dinov2.py}

@@ -6,8 +6,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 from specbuild import REGISTRY
 
-from surge.modules.base import BaseEncoder
-from surge.modules.dinov2.models.vision_transformer import DinoVisionTransformer
+from surge._vendor.dinov2.models.vision_transformer import DinoVisionTransformer
+from surge.modules.encoders.base import BaseEncoder
 
 
 @REGISTRY.register()
@@ -24,7 +24,8 @@ class DINOv2Encoder(BaseEncoder):
         **kwargs,
     ):
         hub_loader = getattr(
-            importlib.import_module(".dinov2.hub.backbones", __package__), backbone
+            importlib.import_module("surge._vendor.dinov2.hub.backbones"),
+            backbone,
         )
         _backbone = hub_loader(pretrained=pretrained)
         super().__init__(

surge/modules/{dinov3_encoder.py → encoders/dinov3.py}

@@ -6,7 +6,7 @@ from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
     checkpoint_wrapper,
 )
 
-from surge.modules.base import BaseEncoder
+from surge.modules.encoders.base import BaseEncoder
 
 
 @REGISTRY.register()

surge/modules/{timm_encoder.py → encoders/timm.py}

@@ -3,7 +3,7 @@ import torch
 from specbuild import REGISTRY
 from timm.models.vision_transformer import VisionTransformer
 
-from surge.modules.base import BaseEncoder
+from surge.modules.encoders.base import BaseEncoder
 
 
 @REGISTRY.register()

surge/modules/heads/__init__.py

@@ -0,0 +1,3 @@
+from surge.modules.heads.nad import NAD
+
+__all__ = ["NAD"]