| 12 | |
| 13 | |
| 14 | class DepthAnythingWrapper(nn.Module): |
| 15 | def __init__(self, model_name: str): |
| 16 | super().__init__() |
| 17 | assert model_name in model_configs.keys(), f"model_name should be in {model_configs.keys()}" |
| 18 | self.model = self._build_depth_anything(model_name) |
| 19 | self._freeze() |
| 20 | |
| 21 | def _build_depth_anything(self, model_name): |
| 22 | from importlib import import_module |
| 23 | da2_hub = import_module(".depth_anything.depth_anything_v2.dpt", package=__package__) |
| 24 | model_fn = getattr(da2_hub, "DepthAnythingV2") |
| 25 | model = model_fn(**model_configs[model_name]) |
| 26 | checkpoint_path = os.path.join(os.path.dirname(__file__), f"../checkpoints/depth_anything_v2_{model_name}.pth") |
| 27 | model.load_state_dict(torch.load(checkpoint_path, map_location='cpu', weights_only=True)) |
| 28 | return model |
| 29 | |
| 30 | def _freeze(self): |
| 31 | # logger.warning(f"======== Freezing Dinov2Wrapper ========") |
| 32 | self.model.eval() |
| 33 | for name, param in self.model.named_parameters(): |
| 34 | param.requires_grad = False |
| 35 | |
| 36 | def forward(self, x): |
| 37 | return self.model(x) |