Upload folder using huggingface_hub

OpenPath/dinov2/hub/__init__.py

@@ -1,4 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the Apache License, Version 2.0
-# found in the LICENSE file in the root directory of this source tree.

OpenPath/dinov2/hub/backbones.py

@@ -1,174 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the Apache License, Version 2.0
-# found in the LICENSE file in the root directory of this source tree.
-
-from enum import Enum
-from pathlib import Path
-from typing import Optional, Union
-from urllib.parse import urlparse
-
-import torch
-
-from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name
-
-
-class Weights(Enum):
-    LVD142M = "LVD142M"
-    XRAY_DINO = "XRay-DINO"
-
-
-def is_url(path: str) -> bool:
-    parsed = urlparse(path)
-    return parsed.scheme in ("https", "file")
-
-
-def convert_path_or_url_to_url(path: str) -> str:
-    if is_url(path):
-        return path
-    return Path(path).expanduser().resolve().as_uri()
-
-
-def _make_dinov2_model(
-    *,
-    arch_name: str = "vit_large",
-    img_size: int = 518,
-    patch_size: int = 14,
-    init_values: float = 1.0,
-    ffn_layer: str = "mlp",
-    block_chunks: int = 0,
-    num_register_tokens: int = 0,
-    interpolate_antialias: bool = False,
-    interpolate_offset: float = 0.1,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.LVD142M,
-    hash: Optional[str] = None,
-    check_hash: bool = False,
-    **kwargs,
-):
-    from ..models import vision_transformer as vits
-
-    model_base_name = _make_dinov2_model_name(arch_name, patch_size)
-    vit_kwargs = dict(
-        img_size=img_size,
-        patch_size=patch_size,
-        init_values=init_values,
-        ffn_layer=ffn_layer,
-        block_chunks=block_chunks,
-        num_register_tokens=num_register_tokens,
-        interpolate_antialias=interpolate_antialias,
-        interpolate_offset=interpolate_offset,
-    )
-    vit_kwargs.update(**kwargs)
-    model = vits.__dict__[arch_name](**vit_kwargs)
-
-    if pretrained:
-        if type(weights) is Weights and weights not in {
-            Weights.LVD142M,
-            Weights.XRAY_DINO,
-        }:
-            raise ValueError(f"Unsupported weights for the backbone: {weights}")
-        elif type(weights) is Weights:
-            model_full_name = _make_dinov2_model_name(arch_name, patch_size, num_register_tokens)
-            url = _DINOV2_BASE_URL + f"/{model_base_name}/{model_full_name}_pretrain.pth"
-        else:
-            url = convert_path_or_url_to_url(weights)
-        state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu", check_hash=check_hash)
-        model.load_state_dict(state_dict, strict=True)
-
-    return model
-
-
-def dinov2_vits14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
-    """
-    DINOv2 ViT-S/14 model (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(arch_name="vit_small", pretrained=pretrained, weights=weights, **kwargs)
-
-
-def dinov2_vitb14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
-    """
-    DINOv2 ViT-B/14 model (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(arch_name="vit_base", pretrained=pretrained, weights=weights, **kwargs)
-
-
-def dinov2_vitl14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
-    """
-    DINOv2 ViT-L/14 model (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(arch_name="vit_large", pretrained=pretrained, weights=weights, **kwargs)
-
-
-def dinov2_vitg14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
-    """
-    DINOv2 ViT-g/14 model (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_giant2",
-        ffn_layer="swiglufused",
-        weights=weights,
-        pretrained=pretrained,
-        **kwargs,
-    )
-
-
-def dinov2_vits14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
-    """
-    DINOv2 ViT-S/14 model with registers (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_small",
-        pretrained=pretrained,
-        weights=weights,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )
-
-
-def dinov2_vitb14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
-    """
-    DINOv2 ViT-B/14 model with registers (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_base",
-        pretrained=pretrained,
-        weights=weights,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )
-
-
-def dinov2_vitl14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
-    """
-    DINOv2 ViT-L/14 model with registers (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_large",
-        pretrained=pretrained,
-        weights=weights,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )
-
-
-def dinov2_vitg14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
-    """
-    DINOv2 ViT-g/14 model with registers (optionally) pretrained on the LVD-142M dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_giant2",
-        ffn_layer="swiglufused",
-        weights=weights,
-        pretrained=pretrained,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )

OpenPath/dinov2/hub/cell_dino/backbones.py

@@ -1,182 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the CC-by-NC licence,
-# found in the LICENSE_CELL_DINO_CODE file in the root directory of this source tree.
-
-from enum import Enum
-from typing import Optional, Union
-
-import torch
-
-
-class Weights(Enum):
-    CELL_DINO = "CELL-DINO"
-
-
-def _make_cell_dino_model(
-    *,
-    arch_name: str = "vit_large",
-    img_size: int = 518,
-    patch_size: int = 14,
-    init_values: float = 1.0,
-    ffn_layer: str = "mlp",
-    block_chunks: int = 0,
-    num_register_tokens: int = 0,
-    interpolate_antialias: bool = False,
-    interpolate_offset: float = 0.1,
-    pretrained: bool = True,
-    channel_adaptive: bool = False,
-    weights: Union[Weights, str] = Weights.CELL_DINO,
-    pretrained_url: Optional[str] = None,
-    pretrained_path: Optional[str] = None,
-    **kwargs,
-):
-    from ...models import vision_transformer as vits
-
-    if isinstance(weights, str):
-        try:
-            weights = Weights[weights]
-        except KeyError:
-            raise AssertionError(f"Unsupported weights: {weights}")
-
-    vit_kwargs = dict(
-        img_size=img_size,
-        patch_size=patch_size,
-        init_values=init_values,
-        ffn_layer=ffn_layer,
-        block_chunks=block_chunks,
-        num_register_tokens=num_register_tokens,
-        interpolate_antialias=interpolate_antialias,
-        interpolate_offset=interpolate_offset,
-        channel_adaptive=channel_adaptive,
-    )
-    vit_kwargs.update(**kwargs)
-    model = vits.__dict__[arch_name](**vit_kwargs)
-
-    if pretrained:
-        if pretrained_path is not None:
-            state_dict = torch.load(pretrained_path, map_location="cpu")
-        else:
-            pretrained_url is not None
-            state_dict = torch.hub.load_state_dict_from_url(pretrained_url, map_location="cpu")
-        model.load_state_dict(state_dict, strict=True)
-
-    return model
-
-
-def cell_dino_hpa_vitl16(
-    *,
-    pretrained_url: Optional[str] = None,
-    pretrained_path: Optional[str] = None,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.CELL_DINO,
-    in_channels: int = 4,
-    **kwargs,
-):
-    """
-    Cell-DINO ViT-L/16 model dataset pretrained on HPA dataset.
-    """
-    return _make_cell_dino_model(
-        arch_name="vit_large",
-        patch_size=16,
-        img_size=224,
-        num_register_tokens=0,
-        interpolate_antialias=False,
-        interpolate_offset=0.1,
-        block_chunks=4,
-        pretrained_url=pretrained_url,
-        pretrained_path=pretrained_path,
-        pretrained=pretrained,
-        weights=weights,
-        in_chans=in_channels,
-        **kwargs,
-    )
-
-
-def cell_dino_hpa_vitl14(
-    *,
-    pretrained_url: Optional[str] = None,
-    pretrained_path: Optional[str] = None,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.CELL_DINO,
-    in_channels: int = 4,
-    **kwargs,
-):
-    """
-    Cell-DINO ViT-L/14 model dataset pretrained on LVD, then on HPA dataset.
-    """
-    return _make_cell_dino_model(
-        arch_name="vit_large",
-        patch_size=14,
-        img_size=518,
-        num_register_tokens=0,
-        interpolate_antialias=False,
-        interpolate_offset=0.1,
-        block_chunks=4,
-        pretrained_url=pretrained_url,
-        pretrained_path=pretrained_path,
-        pretrained=pretrained,
-        weights=weights,
-        in_chans=in_channels,
-        **kwargs,
-    )
-
-
-def cell_dino_cp_vits8(
-    *,
-    pretrained_url: Optional[str] = None,
-    pretrained_path: Optional[str] = None,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.CELL_DINO,
-    in_channels: int = 5,
-    **kwargs,
-):
-    """
-    Cell-DINO ViT-S/8 model dataset pretrained on the combined cell painting dataset.
-    """
-    return _make_cell_dino_model(
-        arch_name="vit_small",
-        patch_size=8,
-        img_size=128,
-        num_register_tokens=0,
-        interpolate_antialias=False,
-        interpolate_offset=0.1,
-        block_chunks=4,
-        pretrained_url=pretrained_url,
-        pretrained_path=pretrained_path,
-        pretrained=pretrained,
-        weights=weights,
-        in_chans=in_channels,
-        **kwargs,
-    )
-
-
-def channel_adaptive_dino_vitl16(
-    *,
-    pretrained_url: Optional[str] = None,
-    pretrained_path: Optional[str] = None,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.CELL_DINO,
-    in_channels: int = 1,
-    channel_adaptive: bool = True,
-    **kwargs,
-):
-    """
-    Cell-DINO ViT-L/16 model dataset pretrained on HPA dataset.
-    """
-    return _make_cell_dino_model(
-        arch_name="vit_large",
-        patch_size=16,
-        img_size=224,
-        num_register_tokens=0,
-        interpolate_antialias=False,
-        interpolate_offset=0.1,
-        block_chunks=4,
-        pretrained_url=pretrained_url,
-        pretrained_path=pretrained_path,
-        pretrained=pretrained,
-        weights=weights,
-        in_chans=in_channels,
-        channel_adaptive=channel_adaptive,
-        **kwargs,
-    )

OpenPath/dinov2/hub/classifiers.py

@@ -1,268 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the Apache License, Version 2.0
-# found in the LICENSE file in the root directory of this source tree.
-
-from enum import Enum
-from typing import Union
-
-import torch
-import torch.nn as nn
-
-from .backbones import _make_dinov2_model
-from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name
-
-
-class Weights(Enum):
-    IMAGENET1K = "IMAGENET1K"
-
-
-def _make_dinov2_linear_classification_head(
-    *,
-    arch_name: str = "vit_large",
-    patch_size: int = 14,
-    embed_dim: int = 1024,
-    layers: int = 4,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.IMAGENET1K,
-    num_register_tokens: int = 0,
-    **kwargs,
-):
-    if layers not in (1, 4):
-        raise AssertionError(f"Unsupported number of layers: {layers}")
-    if isinstance(weights, str):
-        try:
-            weights = Weights[weights]
-        except KeyError:
-            raise AssertionError(f"Unsupported weights: {weights}")
-
-    linear_head = nn.Linear((1 + layers) * embed_dim, 1_000)
-
-    if pretrained:
-        model_base_name = _make_dinov2_model_name(arch_name, patch_size)
-        model_full_name = _make_dinov2_model_name(arch_name, patch_size, num_register_tokens)
-        layers_str = str(layers) if layers == 4 else ""
-        url = _DINOV2_BASE_URL + f"/{model_base_name}/{model_full_name}_linear{layers_str}_head.pth"
-        state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu")
-        linear_head.load_state_dict(state_dict, strict=True)
-
-    return linear_head
-
-
-class _LinearClassifierWrapper(nn.Module):
-    def __init__(self, *, backbone: nn.Module, linear_head: nn.Module, layers: int = 4):
-        super().__init__()
-        self.backbone = backbone
-        self.linear_head = linear_head
-        self.layers = layers
-
-    def forward(self, x):
-        if self.layers == 1:
-            x = self.backbone.forward_features(x)
-            cls_token = x["x_norm_clstoken"]
-            patch_tokens = x["x_norm_patchtokens"]
-            # fmt: off
-            linear_input = torch.cat([
-                cls_token,
-                patch_tokens.mean(dim=1),
-            ], dim=1)
-            # fmt: on
-        elif self.layers == 4:
-            x = self.backbone.get_intermediate_layers(x, n=4, return_class_token=True)
-            # fmt: off
-            linear_input = torch.cat([
-                x[0][1],
-                x[1][1],
-                x[2][1],
-                x[3][1],
-                x[3][0].mean(dim=1),
-            ], dim=1)
-            # fmt: on
-        else:
-            assert False, f"Unsupported number of layers: {self.layers}"
-        return self.linear_head(linear_input)
-
-
-def _make_dinov2_linear_classifier(
-    *,
-    arch_name: str = "vit_large",
-    layers: int = 4,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.IMAGENET1K,
-    num_register_tokens: int = 0,
-    interpolate_antialias: bool = False,
-    interpolate_offset: float = 0.1,
-    **kwargs,
-):
-    backbone = _make_dinov2_model(
-        arch_name=arch_name,
-        pretrained=pretrained,
-        num_register_tokens=num_register_tokens,
-        interpolate_antialias=interpolate_antialias,
-        interpolate_offset=interpolate_offset,
-        **kwargs,
-    )
-
-    embed_dim = backbone.embed_dim
-    patch_size = backbone.patch_size
-    linear_head = _make_dinov2_linear_classification_head(
-        arch_name=arch_name,
-        patch_size=patch_size,
-        embed_dim=embed_dim,
-        layers=layers,
-        pretrained=pretrained,
-        weights=weights,
-        num_register_tokens=num_register_tokens,
-    )
-
-    return _LinearClassifierWrapper(backbone=backbone, linear_head=linear_head, layers=layers)
-
-
-def dinov2_vits14_lc(
-    *,
-    layers: int = 4,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.IMAGENET1K,
-    **kwargs,
-):
-    """
-    Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-S/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
-    """
-    return _make_dinov2_linear_classifier(
-        arch_name="vit_small",
-        layers=layers,
-        pretrained=pretrained,
-        weights=weights,
-        **kwargs,
-    )
-
-
-def dinov2_vitb14_lc(
-    *,
-    layers: int = 4,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.IMAGENET1K,
-    **kwargs,
-):
-    """
-    Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-B/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
-    """
-    return _make_dinov2_linear_classifier(
-        arch_name="vit_base",
-        layers=layers,
-        pretrained=pretrained,
-        weights=weights,
-        **kwargs,
-    )
-
-
-def dinov2_vitl14_lc(
-    *,
-    layers: int = 4,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.IMAGENET1K,
-    **kwargs,
-):
-    """
-    Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-L/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
-    """
-    return _make_dinov2_linear_classifier(
-        arch_name="vit_large",
-        layers=layers,
-        pretrained=pretrained,
-        weights=weights,
-        **kwargs,
-    )
-
-
-def dinov2_vitg14_lc(
-    *,
-    layers: int = 4,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.IMAGENET1K,
-    **kwargs,
-):
-    """
-    Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-g/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
-    """
-    return _make_dinov2_linear_classifier(
-        arch_name="vit_giant2",
-        layers=layers,
-        ffn_layer="swiglufused",
-        pretrained=pretrained,
-        weights=weights,
-        **kwargs,
-    )
-
-
-def dinov2_vits14_reg_lc(
-    *, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.IMAGENET1K, **kwargs
-):
-    """
-    Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-S/14 backbone with registers (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
-    """
-    return _make_dinov2_linear_classifier(
-        arch_name="vit_small",
-        layers=layers,
-        pretrained=pretrained,
-        weights=weights,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )
-
-
-def dinov2_vitb14_reg_lc(
-    *, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.IMAGENET1K, **kwargs
-):
-    """
-    Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-B/14 backbone with registers (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
-    """
-    return _make_dinov2_linear_classifier(
-        arch_name="vit_base",
-        layers=layers,
-        pretrained=pretrained,
-        weights=weights,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )
-
-
-def dinov2_vitl14_reg_lc(
-    *, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.IMAGENET1K, **kwargs
-):
-    """
-    Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-L/14 backbone with registers (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
-    """
-    return _make_dinov2_linear_classifier(
-        arch_name="vit_large",
-        layers=layers,
-        pretrained=pretrained,
-        weights=weights,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )
-
-
-def dinov2_vitg14_reg_lc(
-    *, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.IMAGENET1K, **kwargs
-):
-    """
-    Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-g/14 backbone with registers (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k.
-    """
-    return _make_dinov2_linear_classifier(
-        arch_name="vit_giant2",
-        layers=layers,
-        ffn_layer="swiglufused",
-        pretrained=pretrained,
-        weights=weights,
-        num_register_tokens=4,
-        interpolate_antialias=True,
-        interpolate_offset=0.0,
-        **kwargs,
-    )

OpenPath/dinov2/hub/depth/__init__.py

@@ -1,7 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the Apache License, Version 2.0
-# found in the LICENSE file in the root directory of this source tree.
-
-from .decode_heads import BNHead, DPTHead
-from .encoder_decoder import DepthEncoderDecoder

OpenPath/dinov2/hub/depth/decode_heads.py

@@ -1,747 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the Apache License, Version 2.0
-# found in the LICENSE file in the root directory of this source tree.
-
-import copy
-from functools import partial
-import math
-import warnings
-
-import torch
-import torch.nn as nn
-
-from .ops import resize
-
-
-# XXX: (Untested) replacement for mmcv.imdenormalize()
-def _imdenormalize(img, mean, std, to_bgr=True):
-    import numpy as np
-
-    mean = mean.reshape(1, -1).astype(np.float64)
-    std = std.reshape(1, -1).astype(np.float64)
-    img = (img * std) + mean
-    if to_bgr:
-        img = img[::-1]
-    return img
-
-
-class DepthBaseDecodeHead(nn.Module):
-    """Base class for BaseDecodeHead.
-
-    Args:
-        in_channels (List): Input channels.
-        channels (int): Channels after modules, before conv_depth.
-        conv_layer (nn.Module): Conv layers. Default: None.
-        act_layer (nn.Module): Activation layers. Default: nn.ReLU.
-        loss_decode (dict): Config of decode loss.
-            Default: ().
-        sampler (dict|None): The config of depth map sampler.
-            Default: None.
-        align_corners (bool): align_corners argument of F.interpolate.
-            Default: False.
-        min_depth (int): Min depth in dataset setting.
-            Default: 1e-3.
-        max_depth (int): Max depth in dataset setting.
-            Default: None.
-        norm_layer (dict|None): Norm layers.
-            Default: None.
-        classify (bool): Whether predict depth in a cls.-reg. manner.
-            Default: False.
-        n_bins (int): The number of bins used in cls. step.
-            Default: 256.
-        bins_strategy (str): The discrete strategy used in cls. step.
-            Default: 'UD'.
-        norm_strategy (str): The norm strategy on cls. probability
-            distribution. Default: 'linear'
-        scale_up (str): Whether predict depth in a scale-up manner.
-            Default: False.
-    """
-
-    def __init__(
-        self,
-        in_channels,
-        conv_layer=None,
-        act_layer=nn.ReLU,
-        channels=96,
-        loss_decode=(),
-        sampler=None,
-        align_corners=False,
-        min_depth=1e-3,
-        max_depth=None,
-        norm_layer=None,
-        classify=False,
-        n_bins=256,
-        bins_strategy="UD",
-        norm_strategy="linear",
-        scale_up=False,
-    ):
-        super(DepthBaseDecodeHead, self).__init__()
-
-        self.in_channels = in_channels
-        self.channels = channels
-        self.conf_layer = conv_layer
-        self.act_layer = act_layer
-        self.loss_decode = loss_decode
-        self.align_corners = align_corners
-        self.min_depth = min_depth
-        self.max_depth = max_depth
-        self.norm_layer = norm_layer
-        self.classify = classify
-        self.n_bins = n_bins
-        self.scale_up = scale_up
-
-        if self.classify:
-            assert bins_strategy in ["UD", "SID"], "Support bins_strategy: UD, SID"
-            assert norm_strategy in ["linear", "softmax", "sigmoid"], "Support norm_strategy: linear, softmax, sigmoid"
-
-            self.bins_strategy = bins_strategy
-            self.norm_strategy = norm_strategy
-            self.softmax = nn.Softmax(dim=1)
-            self.conv_depth = nn.Conv2d(channels, n_bins, kernel_size=3, padding=1, stride=1)
-        else:
-            self.conv_depth = nn.Conv2d(channels, 1, kernel_size=3, padding=1, stride=1)
-
-        self.relu = nn.ReLU()
-        self.sigmoid = nn.Sigmoid()
-
-    def forward(self, inputs, img_metas):
-        """Placeholder of forward function."""
-        pass
-
-    def forward_train(self, img, inputs, img_metas, depth_gt):
-        """Forward function for training.
-        Args:
-            inputs (list[Tensor]): List of multi-level img features.
-            img_metas (list[dict]): List of image info dict where each dict
-                has: 'img_shape', 'scale_factor', 'flip', and may also contain
-                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
-                For details on the values of these keys see
-                `depth/datasets/pipelines/formatting.py:Collect`.
-            depth_gt (Tensor): GT depth
-
-        Returns:
-            dict[str, Tensor]: a dictionary of loss components
-        """
-        depth_pred = self.forward(inputs, img_metas)
-        losses = self.losses(depth_pred, depth_gt)
-
-        log_imgs = self.log_images(img[0], depth_pred[0], depth_gt[0], img_metas[0])
-        losses.update(**log_imgs)
-
-        return losses
-
-    def forward_test(self, inputs, img_metas):
-        """Forward function for testing.
-        Args:
-            inputs (list[Tensor]): List of multi-level img features.
-            img_metas (list[dict]): List of image info dict where each dict
-                has: 'img_shape', 'scale_factor', 'flip', and may also contain
-                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
-                For details on the values of these keys see
-                `depth/datasets/pipelines/formatting.py:Collect`.
-
-        Returns:
-            Tensor: Output depth map.
-        """
-        return self.forward(inputs, img_metas)
-
-    def depth_pred(self, feat):
-        """Prediction each pixel."""
-        if self.classify:
-            logit = self.conv_depth(feat)
-
-            if self.bins_strategy == "UD":
-                bins = torch.linspace(self.min_depth, self.max_depth, self.n_bins, device=feat.device)
-            elif self.bins_strategy == "SID":
-                bins = torch.logspace(self.min_depth, self.max_depth, self.n_bins, device=feat.device)
-
-            # following Adabins, default linear
-            if self.norm_strategy == "linear":
-                logit = torch.relu(logit)
-                eps = 0.1
-                logit = logit + eps
-                logit = logit / logit.sum(dim=1, keepdim=True)
-            elif self.norm_strategy == "softmax":
-                logit = torch.softmax(logit, dim=1)
-            elif self.norm_strategy == "sigmoid":
-                logit = torch.sigmoid(logit)
-                logit = logit / logit.sum(dim=1, keepdim=True)
-
-            output = torch.einsum("ikmn,k->imn", [logit, bins]).unsqueeze(dim=1)
-
-        else:
-            if self.scale_up:
-                output = self.sigmoid(self.conv_depth(feat)) * self.max_depth
-            else:
-                output = self.relu(self.conv_depth(feat)) + self.min_depth
-        return output
-
-    def losses(self, depth_pred, depth_gt):
-        """Compute depth loss."""
-        loss = dict()
-        depth_pred = resize(
-            input=depth_pred, size=depth_gt.shape[2:], mode="bilinear", align_corners=self.align_corners, warning=False
-        )
-        if not isinstance(self.loss_decode, nn.ModuleList):
-            losses_decode = [self.loss_decode]
-        else:
-            losses_decode = self.loss_decode
-        for loss_decode in losses_decode:
-            if loss_decode.loss_name not in loss:
-                loss[loss_decode.loss_name] = loss_decode(depth_pred, depth_gt)
-            else:
-                loss[loss_decode.loss_name] += loss_decode(depth_pred, depth_gt)
-        return loss
-
-    def log_images(self, img_path, depth_pred, depth_gt, img_meta):
-        import numpy as np
-
-        show_img = copy.deepcopy(img_path.detach().cpu().permute(1, 2, 0))
-        show_img = show_img.numpy().astype(np.float32)
-        show_img = _imdenormalize(
-            show_img,
-            img_meta["img_norm_cfg"]["mean"],
-            img_meta["img_norm_cfg"]["std"],
-            img_meta["img_norm_cfg"]["to_rgb"],
-        )
-        show_img = np.clip(show_img, 0, 255)
-        show_img = show_img.astype(np.uint8)
-        show_img = show_img[:, :, ::-1]
-        show_img = show_img.transpose(0, 2, 1)
-        show_img = show_img.transpose(1, 0, 2)
-
-        depth_pred = depth_pred / torch.max(depth_pred)
-        depth_gt = depth_gt / torch.max(depth_gt)
-
-        depth_pred_color = copy.deepcopy(depth_pred.detach().cpu())
-        depth_gt_color = copy.deepcopy(depth_gt.detach().cpu())
-
-        return {"img_rgb": show_img, "img_depth_pred": depth_pred_color, "img_depth_gt": depth_gt_color}
-
-
-class BNHead(DepthBaseDecodeHead):
-    """Just a batchnorm."""
-
-    def __init__(self, input_transform="resize_concat", in_index=(0, 1, 2, 3), upsample=1, **kwargs):
-        super().__init__(**kwargs)
-        self.input_transform = input_transform
-        self.in_index = in_index
-        self.upsample = upsample
-        # self.bn = nn.SyncBatchNorm(self.in_channels)
-        if self.classify:
-            self.conv_depth = nn.Conv2d(self.channels, self.n_bins, kernel_size=1, padding=0, stride=1)
-        else:
-            self.conv_depth = nn.Conv2d(self.channels, 1, kernel_size=1, padding=0, stride=1)
-
-    def _transform_inputs(self, inputs):
-        """Transform inputs for decoder.
-        Args:
-            inputs (list[Tensor]): List of multi-level img features.
-        Returns:
-            Tensor: The transformed inputs
-        """
-
-        if "concat" in self.input_transform:
-            inputs = [inputs[i] for i in self.in_index]
-            if "resize" in self.input_transform:
-                inputs = [
-                    resize(
-                        input=x,
-                        size=[s * self.upsample for s in inputs[0].shape[2:]],
-                        mode="bilinear",
-                        align_corners=self.align_corners,
-                    )
-                    for x in inputs
-                ]
-            inputs = torch.cat(inputs, dim=1)
-        elif self.input_transform == "multiple_select":
-            inputs = [inputs[i] for i in self.in_index]
-        else:
-            inputs = inputs[self.in_index]
-
-        return inputs
-
-    def _forward_feature(self, inputs, img_metas=None, **kwargs):
-        """Forward function for feature maps before classifying each pixel with
-        ``self.cls_seg`` fc.
-        Args:
-            inputs (list[Tensor]): List of multi-level img features.
-        Returns:
-            feats (Tensor): A tensor of shape (batch_size, self.channels,
-                H, W) which is feature map for last layer of decoder head.
-        """
-        # accept lists (for cls token)
-        inputs = list(inputs)
-        for i, x in enumerate(inputs):
-            if len(x) == 2:
-                x, cls_token = x[0], x[1]
-                if len(x.shape) == 2:
-                    x = x[:, :, None, None]
-                cls_token = cls_token[:, :, None, None].expand_as(x)
-                inputs[i] = torch.cat((x, cls_token), 1)
-            else:
-                x = x[0]
-                if len(x.shape) == 2:
-                    x = x[:, :, None, None]
-                inputs[i] = x
-        x = self._transform_inputs(inputs)
-        # feats = self.bn(x)
-        return x
-
-    def forward(self, inputs, img_metas=None, **kwargs):
-        """Forward function."""
-        output = self._forward_feature(inputs, img_metas=img_metas, **kwargs)
-        output = self.depth_pred(output)
-        return output
-
-
-class ConvModule(nn.Module):
-    """A conv block that bundles conv/norm/activation layers.
-
-    This block simplifies the usage of convolution layers, which are commonly
-    used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU).
-    It is based upon three build methods: `build_conv_layer()`,
-    `build_norm_layer()` and `build_activation_layer()`.
-
-    Besides, we add some additional features in this module.
-    1. Automatically set `bias` of the conv layer.
-    2. Spectral norm is supported.
-    3. More padding modes are supported. Before PyTorch 1.5, nn.Conv2d only
-    supports zero and circular padding, and we add "reflect" padding mode.
-
-    Args:
-        in_channels (int): Number of channels in the input feature map.
-            Same as that in ``nn._ConvNd``.
-        out_channels (int): Number of channels produced by the convolution.
-            Same as that in ``nn._ConvNd``.
-        kernel_size (int | tuple[int]): Size of the convolving kernel.
-            Same as that in ``nn._ConvNd``.
-        stride (int | tuple[int]): Stride of the convolution.
-            Same as that in ``nn._ConvNd``.
-        padding (int | tuple[int]): Zero-padding added to both sides of
-            the input. Same as that in ``nn._ConvNd``.
-        dilation (int | tuple[int]): Spacing between kernel elements.
-            Same as that in ``nn._ConvNd``.
-        groups (int): Number of blocked connections from input channels to
-            output channels. Same as that in ``nn._ConvNd``.
-        bias (bool | str): If specified as `auto`, it will be decided by the
-            norm_layer. Bias will be set as True if `norm_layer` is None, otherwise
-            False. Default: "auto".
-        conv_layer (nn.Module): Convolution layer. Default: None,
-            which means using conv2d.
-        norm_layer (nn.Module): Normalization layer. Default: None.
-        act_layer (nn.Module): Activation layer. Default: nn.ReLU.
-        inplace (bool): Whether to use inplace mode for activation.
-            Default: True.
-        with_spectral_norm (bool): Whether use spectral norm in conv module.
-            Default: False.
-        padding_mode (str): If the `padding_mode` has not been supported by
-            current `Conv2d` in PyTorch, we will use our own padding layer
-            instead. Currently, we support ['zeros', 'circular'] with official
-            implementation and ['reflect'] with our own implementation.
-            Default: 'zeros'.
-        order (tuple[str]): The order of conv/norm/activation layers. It is a
-            sequence of "conv", "norm" and "act". Common examples are
-            ("conv", "norm", "act") and ("act", "conv", "norm").
-            Default: ('conv', 'norm', 'act').
-    """
-
-    _abbr_ = "conv_block"
-
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        groups=1,
-        bias="auto",
-        conv_layer=nn.Conv2d,
-        norm_layer=None,
-        act_layer=nn.ReLU,
-        inplace=True,
-        with_spectral_norm=False,
-        padding_mode="zeros",
-        order=("conv", "norm", "act"),
-    ):
-        super(ConvModule, self).__init__()
-        official_padding_mode = ["zeros", "circular"]
-        self.conv_layer = conv_layer
-        self.norm_layer = norm_layer
-        self.act_layer = act_layer
-        self.inplace = inplace
-        self.with_spectral_norm = with_spectral_norm
-        self.with_explicit_padding = padding_mode not in official_padding_mode
-        self.order = order
-        assert isinstance(self.order, tuple) and len(self.order) == 3
-        assert set(order) == set(["conv", "norm", "act"])
-
-        self.with_norm = norm_layer is not None
-        self.with_activation = act_layer is not None
-        # if the conv layer is before a norm layer, bias is unnecessary.
-        if bias == "auto":
-            bias = not self.with_norm
-        self.with_bias = bias
-
-        if self.with_explicit_padding:
-            if padding_mode == "zeros":
-                padding_layer = nn.ZeroPad2d
-            else:
-                raise AssertionError(f"Unsupported padding mode: {padding_mode}")
-            self.pad = padding_layer(padding)
-
-        # reset padding to 0 for conv module
-        conv_padding = 0 if self.with_explicit_padding else padding
-        # build convolution layer
-        self.conv = self.conv_layer(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride=stride,
-            padding=conv_padding,
-            dilation=dilation,
-            groups=groups,
-            bias=bias,
-        )
-        # export the attributes of self.conv to a higher level for convenience
-        self.in_channels = self.conv.in_channels
-        self.out_channels = self.conv.out_channels
-        self.kernel_size = self.conv.kernel_size
-        self.stride = self.conv.stride
-        self.padding = padding
-        self.dilation = self.conv.dilation
-        self.transposed = self.conv.transposed
-        self.output_padding = self.conv.output_padding
-        self.groups = self.conv.groups
-
-        if self.with_spectral_norm:
-            self.conv = nn.utils.spectral_norm(self.conv)
-
-        # build normalization layers
-        if self.with_norm:
-            # norm layer is after conv layer
-            if order.index("norm") > order.index("conv"):
-                norm_channels = out_channels
-            else:
-                norm_channels = in_channels
-            norm = partial(norm_layer, num_features=norm_channels)
-            self.add_module("norm", norm)
-            if self.with_bias:
-                from torch.nnModules.batchnorm import _BatchNorm
-                from torch.nnModules.instancenorm import _InstanceNorm
-
-                if isinstance(norm, (_BatchNorm, _InstanceNorm)):
-                    warnings.warn("Unnecessary conv bias before batch/instance norm")
-        else:
-            self.norm_name = None
-
-        # build activation layer
-        if self.with_activation:
-            # nn.Tanh has no 'inplace' argument
-            # (nn.Tanh, nn.PReLU, nn.Sigmoid, nn.HSigmoid, nn.Swish, nn.GELU)
-            if not isinstance(act_layer, (nn.Tanh, nn.PReLU, nn.Sigmoid, nn.GELU)):
-                act_layer = partial(act_layer, inplace=inplace)
-            self.activate = act_layer()
-
-        # Use msra init by default
-        self.init_weights()
-
-    @property
-    def norm(self):
-        if self.norm_name:
-            return getattr(self, self.norm_name)
-        else:
-            return None
-
-    def init_weights(self):
-        # 1. It is mainly for customized conv layers with their own
-        #    initialization manners by calling their own ``init_weights()``,
-        #    and we do not want ConvModule to override the initialization.
-        # 2. For customized conv layers without their own initialization
-        #    manners (that is, they don't have their own ``init_weights()``)
-        #    and PyTorch's conv layers, they will be initialized by
-        #    this method with default ``kaiming_init``.
-        # Note: For PyTorch's conv layers, they will be overwritten by our
-        #    initialization implementation using default ``kaiming_init``.
-        if not hasattr(self.conv, "init_weights"):
-            if self.with_activation and isinstance(self.act_layer, nn.LeakyReLU):
-                nonlinearity = "leaky_relu"
-                a = 0.01  # XXX: default negative_slope
-            else:
-                nonlinearity = "relu"
-                a = 0
-            if hasattr(self.conv, "weight") and self.conv.weight is not None:
-                nn.init.kaiming_normal_(self.conv.weight, a=a, mode="fan_out", nonlinearity=nonlinearity)
-            if hasattr(self.conv, "bias") and self.conv.bias is not None:
-                nn.init.constant_(self.conv.bias, 0)
-        if self.with_norm:
-            if hasattr(self.norm, "weight") and self.norm.weight is not None:
-                nn.init.constant_(self.norm.weight, 1)
-            if hasattr(self.norm, "bias") and self.norm.bias is not None:
-                nn.init.constant_(self.norm.bias, 0)
-
-    def forward(self, x, activate=True, norm=True):
-        for layer in self.order:
-            if layer == "conv":
-                if self.with_explicit_padding:
-                    x = self.pad(x)
-                x = self.conv(x)
-            elif layer == "norm" and norm and self.with_norm:
-                x = self.norm(x)
-            elif layer == "act" and activate and self.with_activation:
-                x = self.activate(x)
-        return x
-
-
-class Interpolate(nn.Module):
-    def __init__(self, scale_factor, mode, align_corners=False):
-        super(Interpolate, self).__init__()
-        self.interp = nn.functional.interpolate
-        self.scale_factor = scale_factor
-        self.mode = mode
-        self.align_corners = align_corners
-
-    def forward(self, x):
-        x = self.interp(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners)
-        return x
-
-
-class HeadDepth(nn.Module):
-    def __init__(self, features):
-        super(HeadDepth, self).__init__()
-        self.head = nn.Sequential(
-            nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),
-            Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
-            nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
-            nn.ReLU(),
-            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
-        )
-
-    def forward(self, x):
-        x = self.head(x)
-        return x
-
-
-class ReassembleBlocks(nn.Module):
-    """ViTPostProcessBlock, process cls_token in ViT backbone output and
-    rearrange the feature vector to feature map.
-    Args:
-        in_channels (int): ViT feature channels. Default: 768.
-        out_channels (List): output channels of each stage.
-            Default: [96, 192, 384, 768].
-        readout_type (str): Type of readout operation. Default: 'ignore'.
-        patch_size (int): The patch size. Default: 16.
-    """
-
-    def __init__(self, in_channels=768, out_channels=[96, 192, 384, 768], readout_type="ignore", patch_size=16):
-        super(ReassembleBlocks, self).__init__()
-
-        assert readout_type in ["ignore", "add", "project"]
-        self.readout_type = readout_type
-        self.patch_size = patch_size
-
-        self.projects = nn.ModuleList(
-            [
-                ConvModule(
-                    in_channels=in_channels,
-                    out_channels=out_channel,
-                    kernel_size=1,
-                    act_layer=None,
-                )
-                for out_channel in out_channels
-            ]
-        )
-
-        self.resize_layers = nn.ModuleList(
-            [
-                nn.ConvTranspose2d(
-                    in_channels=out_channels[0], out_channels=out_channels[0], kernel_size=4, stride=4, padding=0
-                ),
-                nn.ConvTranspose2d(
-                    in_channels=out_channels[1], out_channels=out_channels[1], kernel_size=2, stride=2, padding=0
-                ),
-                nn.Identity(),
-                nn.Conv2d(
-                    in_channels=out_channels[3], out_channels=out_channels[3], kernel_size=3, stride=2, padding=1
-                ),
-            ]
-        )
-        if self.readout_type == "project":
-            self.readout_projects = nn.ModuleList()
-            for _ in range(len(self.projects)):
-                self.readout_projects.append(nn.Sequential(nn.Linear(2 * in_channels, in_channels), nn.GELU()))
-
-    def forward(self, inputs):
-        assert isinstance(inputs, list)
-        out = []
-        for i, x in enumerate(inputs):
-            assert len(x) == 2
-            x, cls_token = x[0], x[1]
-            feature_shape = x.shape
-            if self.readout_type == "project":
-                x = x.flatten(2).permute((0, 2, 1))
-                readout = cls_token.unsqueeze(1).expand_as(x)
-                x = self.readout_projects[i](torch.cat((x, readout), -1))
-                x = x.permute(0, 2, 1).reshape(feature_shape)
-            elif self.readout_type == "add":
-                x = x.flatten(2) + cls_token.unsqueeze(-1)
-                x = x.reshape(feature_shape)
-            else:
-                pass
-            x = self.projects[i](x)
-            x = self.resize_layers[i](x)
-            out.append(x)
-        return out
-
-
-class PreActResidualConvUnit(nn.Module):
-    """ResidualConvUnit, pre-activate residual unit.
-    Args:
-        in_channels (int): number of channels in the input feature map.
-        act_layer (nn.Module): activation layer.
-        norm_layer (nn.Module): norm layer.
-        stride (int): stride of the first block. Default: 1
-        dilation (int): dilation rate for convs layers. Default: 1.
-    """
-
-    def __init__(self, in_channels, act_layer, norm_layer, stride=1, dilation=1):
-        super(PreActResidualConvUnit, self).__init__()
-
-        self.conv1 = ConvModule(
-            in_channels,
-            in_channels,
-            3,
-            stride=stride,
-            padding=dilation,
-            dilation=dilation,
-            norm_layer=norm_layer,
-            act_layer=act_layer,
-            bias=False,
-            order=("act", "conv", "norm"),
-        )
-
-        self.conv2 = ConvModule(
-            in_channels,
-            in_channels,
-            3,
-            padding=1,
-            norm_layer=norm_layer,
-            act_layer=act_layer,
-            bias=False,
-            order=("act", "conv", "norm"),
-        )
-
-    def forward(self, inputs):
-        inputs_ = inputs.clone()
-        x = self.conv1(inputs)
-        x = self.conv2(x)
-        return x + inputs_
-
-
-class FeatureFusionBlock(nn.Module):
-    """FeatureFusionBlock, merge feature map from different stages.
-    Args:
-        in_channels (int): Input channels.
-        act_layer (nn.Module): activation layer for ResidualConvUnit.
-        norm_layer (nn.Module): normalization layer.
-        expand (bool): Whether expand the channels in post process block.
-            Default: False.
-        align_corners (bool): align_corner setting for bilinear upsample.
-            Default: True.
-    """
-
-    def __init__(self, in_channels, act_layer, norm_layer, expand=False, align_corners=True):
-        super(FeatureFusionBlock, self).__init__()
-
-        self.in_channels = in_channels
-        self.expand = expand
-        self.align_corners = align_corners
-
-        self.out_channels = in_channels
-        if self.expand:
-            self.out_channels = in_channels // 2
-
-        self.project = ConvModule(self.in_channels, self.out_channels, kernel_size=1, act_layer=None, bias=True)
-
-        self.res_conv_unit1 = PreActResidualConvUnit(
-            in_channels=self.in_channels, act_layer=act_layer, norm_layer=norm_layer
-        )
-        self.res_conv_unit2 = PreActResidualConvUnit(
-            in_channels=self.in_channels, act_layer=act_layer, norm_layer=norm_layer
-        )
-
-    def forward(self, *inputs):
-        x = inputs[0]
-        if len(inputs) == 2:
-            if x.shape != inputs[1].shape:
-                res = resize(inputs[1], size=(x.shape[2], x.shape[3]), mode="bilinear", align_corners=False)
-            else:
-                res = inputs[1]
-            x = x + self.res_conv_unit1(res)
-        x = self.res_conv_unit2(x)
-        x = resize(x, scale_factor=2, mode="bilinear", align_corners=self.align_corners)
-        x = self.project(x)
-        return x
-
-
-class DPTHead(DepthBaseDecodeHead):
-    """Vision Transformers for Dense Prediction.
-    This head is implemented of `DPT <https://arxiv.org/abs/2103.13413>`_.
-    Args:
-        embed_dims (int): The embed dimension of the ViT backbone.
-            Default: 768.
-        post_process_channels (List): Out channels of post process conv
-            layers. Default: [96, 192, 384, 768].
-        readout_type (str): Type of readout operation. Default: 'ignore'.
-        patch_size (int): The patch size. Default: 16.
-        expand_channels (bool): Whether expand the channels in post process
-            block. Default: False.
-    """
-
-    def __init__(
-        self,
-        embed_dims=768,
-        post_process_channels=[96, 192, 384, 768],
-        readout_type="ignore",
-        patch_size=16,
-        expand_channels=False,
-        **kwargs,
-    ):
-        super(DPTHead, self).__init__(**kwargs)
-
-        self.in_channels = self.in_channels
-        self.expand_channels = expand_channels
-        self.reassemble_blocks = ReassembleBlocks(embed_dims, post_process_channels, readout_type, patch_size)
-
-        self.post_process_channels = [
-            channel * math.pow(2, i) if expand_channels else channel for i, channel in enumerate(post_process_channels)
-        ]
-        self.convs = nn.ModuleList()
-        for channel in self.post_process_channels:
-            self.convs.append(ConvModule(channel, self.channels, kernel_size=3, padding=1, act_layer=None, bias=False))
-        self.fusion_blocks = nn.ModuleList()
-        for _ in range(len(self.convs)):
-            self.fusion_blocks.append(FeatureFusionBlock(self.channels, self.act_layer, self.norm_layer))
-        self.fusion_blocks[0].res_conv_unit1 = None
-        self.project = ConvModule(self.channels, self.channels, kernel_size=3, padding=1, norm_layer=self.norm_layer)
-        self.num_fusion_blocks = len(self.fusion_blocks)
-        self.num_reassemble_blocks = len(self.reassemble_blocks.resize_layers)
-        self.num_post_process_channels = len(self.post_process_channels)
-        assert self.num_fusion_blocks == self.num_reassemble_blocks
-        assert self.num_reassemble_blocks == self.num_post_process_channels
-        self.conv_depth = HeadDepth(self.channels)
-
-    def forward(self, inputs, img_metas):
-        assert len(inputs) == self.num_reassemble_blocks
-        x = [inp for inp in inputs]
-        x = self.reassemble_blocks(x)
-        x = [self.convs[i](feature) for i, feature in enumerate(x)]
-        out = self.fusion_blocks[0](x[-1])
-        for i in range(1, len(self.fusion_blocks)):
-            out = self.fusion_blocks[i](out, x[-(i + 1)])
-        out = self.project(out)
-        out = self.depth_pred(out)
-        return out

OpenPath/dinov2/hub/depth/encoder_decoder.py

@@ -1,351 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the Apache License, Version 2.0
-# found in the LICENSE file in the root directory of this source tree.
-
-from collections import OrderedDict
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .ops import resize
-
-
-def add_prefix(inputs, prefix):
-    """Add prefix for dict.
-
-    Args:
-        inputs (dict): The input dict with str keys.
-        prefix (str): The prefix to add.
-
-    Returns:
-
-        dict: The dict with keys updated with ``prefix``.
-    """
-
-    outputs = dict()
-    for name, value in inputs.items():
-        outputs[f"{prefix}.{name}"] = value
-
-    return outputs
-
-
-class DepthEncoderDecoder(nn.Module):
-    """Encoder Decoder depther.
-
-    EncoderDecoder typically consists of backbone and decode_head.
-    """
-
-    def __init__(self, backbone, decode_head):
-        super(DepthEncoderDecoder, self).__init__()
-
-        self.backbone = backbone
-        self.decode_head = decode_head
-        self.align_corners = self.decode_head.align_corners
-
-    def extract_feat(self, img):
-        """Extract features from images."""
-        return self.backbone(img)
-
-    def encode_decode(self, img, img_metas, rescale=True, size=None):
-        """Encode images with backbone and decode into a depth estimation
-        map of the same size as input."""
-        x = self.extract_feat(img)
-        out = self._decode_head_forward_test(x, img_metas)
-        # crop the pred depth to the certain range.
-        out = torch.clamp(out, min=self.decode_head.min_depth, max=self.decode_head.max_depth)
-        if rescale:
-            if size is None:
-                if img_metas is not None:
-                    size = img_metas[0]["ori_shape"][:2]
-                else:
-                    size = img.shape[2:]
-            out = resize(input=out, size=size, mode="bilinear", align_corners=self.align_corners)
-        return out
-
-    def _decode_head_forward_train(self, img, x, img_metas, depth_gt, **kwargs):
-        """Run forward function and calculate loss for decode head in
-        training."""
-        losses = dict()
-        loss_decode = self.decode_head.forward_train(img, x, img_metas, depth_gt, **kwargs)
-        losses.update(add_prefix(loss_decode, "decode"))
-        return losses
-
-    def _decode_head_forward_test(self, x, img_metas):
-        """Run forward function and calculate loss for decode head in
-        inference."""
-        depth_pred = self.decode_head.forward_test(x, img_metas)
-        return depth_pred
-
-    def forward_dummy(self, img):
-        """Dummy forward function."""
-        depth = self.encode_decode(img, None)
-
-        return depth
-
-    def forward_train(self, img, img_metas, depth_gt, **kwargs):
-        """Forward function for training.
-
-        Args:
-            img (Tensor): Input images.
-            img_metas (list[dict]): List of image info dict where each dict
-                has: 'img_shape', 'scale_factor', 'flip', and may also contain
-                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
-                For details on the values of these keys see
-                `depth/datasets/pipelines/formatting.py:Collect`.
-            depth_gt (Tensor): Depth gt
-                used if the architecture supports depth estimation task.
-
-        Returns:
-            dict[str, Tensor]: a dictionary of loss components
-        """
-
-        x = self.extract_feat(img)
-
-        losses = dict()
-
-        # the last of x saves the info from neck
-        loss_decode = self._decode_head_forward_train(img, x, img_metas, depth_gt, **kwargs)
-
-        losses.update(loss_decode)
-
-        return losses
-
-    def whole_inference(self, img, img_meta, rescale, size=None):
-        """Inference with full image."""
-        return self.encode_decode(img, img_meta, rescale, size=size)
-
-    def slide_inference(self, img, img_meta, rescale, stride, crop_size):
-        """Inference by sliding-window with overlap.
-
-        If h_crop > h_img or w_crop > w_img, the small patch will be used to
-        decode without padding.
-        """
-
-        h_stride, w_stride = stride
-        h_crop, w_crop = crop_size
-        batch_size, _, h_img, w_img = img.size()
-        h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
-        w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
-        preds = img.new_zeros((batch_size, 1, h_img, w_img))
-        count_mat = img.new_zeros((batch_size, 1, h_img, w_img))
-        for h_idx in range(h_grids):
-            for w_idx in range(w_grids):
-                y1 = h_idx * h_stride
-                x1 = w_idx * w_stride
-                y2 = min(y1 + h_crop, h_img)
-                x2 = min(x1 + w_crop, w_img)
-                y1 = max(y2 - h_crop, 0)
-                x1 = max(x2 - w_crop, 0)
-                crop_img = img[:, :, y1:y2, x1:x2]
-                depth_pred = self.encode_decode(crop_img, img_meta, rescale)
-                preds += F.pad(depth_pred, (int(x1), int(preds.shape[3] - x2), int(y1), int(preds.shape[2] - y2)))
-
-                count_mat[:, :, y1:y2, x1:x2] += 1
-        assert (count_mat == 0).sum() == 0
-        if torch.onnx.is_in_onnx_export():
-            # cast count_mat to constant while exporting to ONNX
-            count_mat = torch.from_numpy(count_mat.cpu().detach().numpy()).to(device=img.device)
-        preds = preds / count_mat
-        return preds
-
-    def inference(self, img, img_meta, rescale, size=None, mode="whole"):
-        """Inference with slide/whole style.
-
-        Args:
-            img (Tensor): The input image of shape (N, 3, H, W).
-            img_meta (dict): Image info dict where each dict has: 'img_shape',
-                'scale_factor', 'flip', and may also contain
-                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
-                For details on the values of these keys see
-                `depth/datasets/pipelines/formatting.py:Collect`.
-            rescale (bool): Whether rescale back to original shape.
-
-        Returns:
-            Tensor: The output depth map.
-        """
-
-        assert mode in ["slide", "whole"]
-        ori_shape = img_meta[0]["ori_shape"]
-        assert all(_["ori_shape"] == ori_shape for _ in img_meta)
-        if mode == "slide":
-            depth_pred = self.slide_inference(img, img_meta, rescale)
-        else:
-            depth_pred = self.whole_inference(img, img_meta, rescale, size=size)
-        output = depth_pred
-        flip = img_meta[0]["flip"]
-        if flip:
-            flip_direction = img_meta[0]["flip_direction"]
-            assert flip_direction in ["horizontal", "vertical"]
-            if flip_direction == "horizontal":
-                output = output.flip(dims=(3,))
-            elif flip_direction == "vertical":
-                output = output.flip(dims=(2,))
-
-        return output
-
-    def simple_test(self, img, img_meta, rescale=True):
-        """Simple test with single image."""
-        depth_pred = self.inference(img, img_meta, rescale)
-        if torch.onnx.is_in_onnx_export():
-            # our inference backend only support 4D output
-            depth_pred = depth_pred.unsqueeze(0)
-            return depth_pred
-        depth_pred = depth_pred.cpu().numpy()
-        # unravel batch dim
-        depth_pred = list(depth_pred)
-        return depth_pred
-
-    def aug_test(self, imgs, img_metas, rescale=True):
-        """Test with augmentations.
-
-        Only rescale=True is supported.
-        """
-        # aug_test rescale all imgs back to ori_shape for now
-        assert rescale
-        # to save memory, we get augmented depth logit inplace
-        depth_pred = self.inference(imgs[0], img_metas[0], rescale)
-        for i in range(1, len(imgs)):
-            cur_depth_pred = self.inference(imgs[i], img_metas[i], rescale, size=depth_pred.shape[-2:])
-            depth_pred += cur_depth_pred
-        depth_pred /= len(imgs)
-        depth_pred = depth_pred.cpu().numpy()
-        # unravel batch dim
-        depth_pred = list(depth_pred)
-        return depth_pred
-
-    def forward_test(self, imgs, img_metas, **kwargs):
-        """
-        Args:
-            imgs (List[Tensor]): the outer list indicates test-time
-                augmentations and inner Tensor should have a shape NxCxHxW,
-                which contains all images in the batch.
-            img_metas (List[List[dict]]): the outer list indicates test-time
-                augs (multiscale, flip, etc.) and the inner list indicates
-                images in a batch.
-        """
-        for var, name in [(imgs, "imgs"), (img_metas, "img_metas")]:
-            if not isinstance(var, list):
-                raise TypeError(f"{name} must be a list, but got " f"{type(var)}")
-        num_augs = len(imgs)
-        if num_augs != len(img_metas):
-            raise ValueError(f"num of augmentations ({len(imgs)}) != " f"num of image meta ({len(img_metas)})")
-        # all images in the same aug batch all of the same ori_shape and pad
-        # shape
-        for img_meta in img_metas:
-            ori_shapes = [_["ori_shape"] for _ in img_meta]
-            assert all(shape == ori_shapes[0] for shape in ori_shapes)
-            img_shapes = [_["img_shape"] for _ in img_meta]
-            assert all(shape == img_shapes[0] for shape in img_shapes)
-            pad_shapes = [_["pad_shape"] for _ in img_meta]
-            assert all(shape == pad_shapes[0] for shape in pad_shapes)
-
-        if num_augs == 1:
-            return self.simple_test(imgs[0], img_metas[0], **kwargs)
-        else:
-            return self.aug_test(imgs, img_metas, **kwargs)
-
-    def forward(self, img, img_metas, return_loss=True, **kwargs):
-        """Calls either :func:`forward_train` or :func:`forward_test` depending
-        on whether ``return_loss`` is ``True``.
-
-        Note this setting will change the expected inputs. When
-        ``return_loss=True``, img and img_meta are single-nested (i.e. Tensor
-        and List[dict]), and when ``resturn_loss=False``, img and img_meta
-        should be double nested (i.e.  List[Tensor], List[List[dict]]), with
-        the outer list indicating test time augmentations.
-        """
-        if return_loss:
-            return self.forward_train(img, img_metas, **kwargs)
-        else:
-            return self.forward_test(img, img_metas, **kwargs)
-
-    def train_step(self, data_batch, optimizer, **kwargs):
-        """The iteration step during training.
-
-        This method defines an iteration step during training, except for the
-        back propagation and optimizer updating, which are done in an optimizer
-        hook. Note that in some complicated cases or models, the whole process
-        including back propagation and optimizer updating is also defined in
-        this method, such as GAN.
-
-        Args:
-            data (dict): The output of dataloader.
-            optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of
-                runner is passed to ``train_step()``. This argument is unused
-                and reserved.
-
-        Returns:
-            dict: It should contain at least 3 keys: ``loss``, ``log_vars``,
-                ``num_samples``.
-                ``loss`` is a tensor for back propagation, which can be a
-                weighted sum of multiple losses.
-                ``log_vars`` contains all the variables to be sent to the
-                logger.
-                ``num_samples`` indicates the batch size (when the model is
-                DDP, it means the batch size on each GPU), which is used for
-                averaging the logs.
-        """
-        losses = self(**data_batch)
-
-        # split losses and images
-        real_losses = {}
-        log_imgs = {}
-        for k, v in losses.items():
-            if "img" in k:
-                log_imgs[k] = v
-            else:
-                real_losses[k] = v
-
-        loss, log_vars = self._parse_losses(real_losses)
-
-        outputs = dict(loss=loss, log_vars=log_vars, num_samples=len(data_batch["img_metas"]), log_imgs=log_imgs)
-
-        return outputs
-
-    def val_step(self, data_batch, **kwargs):
-        """The iteration step during validation.
-
-        This method shares the same signature as :func:`train_step`, but used
-        during val epochs. Note that the evaluation after training epochs is
-        not implemented with this method, but an evaluation hook.
-        """
-        output = self(**data_batch, **kwargs)
-        return output
-
-    @staticmethod
-    def _parse_losses(losses):
-        import torch.distributed as dist
-
-        """Parse the raw outputs (losses) of the network.
-
-        Args:
-            losses (dict): Raw output of the network, which usually contain
-                losses and other necessary information.
-
-        Returns:
-            tuple[Tensor, dict]: (loss, log_vars), loss is the loss tensor
-                which may be a weighted sum of all losses, log_vars contains
-                all the variables to be sent to the logger.
-        """
-        log_vars = OrderedDict()
-        for loss_name, loss_value in losses.items():
-            if isinstance(loss_value, torch.Tensor):
-                log_vars[loss_name] = loss_value.mean()
-            elif isinstance(loss_value, list):
-                log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value)
-            else:
-                raise TypeError(f"{loss_name} is not a tensor or list of tensors")
-
-        loss = sum(_value for _key, _value in log_vars.items() if "loss" in _key)
-
-        log_vars["loss"] = loss
-        for loss_name, loss_value in log_vars.items():
-            # reduce loss when distributed training
-            if dist.is_available() and dist.is_initialized():
-                loss_value = loss_value.data.clone()
-                dist.all_reduce(loss_value.div_(dist.get_world_size()))
-            log_vars[loss_name] = loss_value.item()
-
-        return loss, log_vars

OpenPath/dinov2/hub/depth/ops.py

@@ -1,28 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the Apache License, Version 2.0
-# found in the LICENSE file in the root directory of this source tree.
-
-import warnings
-
-import torch.nn.functional as F
-
-
-def resize(input, size=None, scale_factor=None, mode="nearest", align_corners=None, warning=False):
-    if warning:
-        if size is not None and align_corners:
-            input_h, input_w = tuple(int(x) for x in input.shape[2:])
-            output_h, output_w = tuple(int(x) for x in size)
-            if output_h > input_h or output_w > output_h:
-                if (
-                    (output_h > 1 and output_w > 1 and input_h > 1 and input_w > 1)
-                    and (output_h - 1) % (input_h - 1)
-                    and (output_w - 1) % (input_w - 1)
-                ):
-                    warnings.warn(
-                        f"When align_corners={align_corners}, "
-                        "the output would more aligned if "
-                        f"input size {(input_h, input_w)} is `x+1` and "
-                        f"out size {(output_h, output_w)} is `nx+1`"
-                    )
-    return F.interpolate(input, size, scale_factor, mode, align_corners)

OpenPath/dinov2/hub/depthers.py

@@ -1,246 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the Apache License, Version 2.0
-# found in the LICENSE file in the root directory of this source tree.
-
-from enum import Enum
-from functools import partial
-from typing import Optional, Tuple, Union
-
-import torch
-
-from .backbones import _make_dinov2_model
-from .depth import BNHead, DepthEncoderDecoder, DPTHead
-from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name, CenterPadding
-
-
-class Weights(Enum):
-    NYU = "NYU"
-    KITTI = "KITTI"
-
-
-def _get_depth_range(pretrained: bool, weights: Weights = Weights.NYU) -> Tuple[float, float]:
-    if not pretrained:  # Default
-        return (0.001, 10.0)
-
-    # Pretrained, set according to the training dataset for the provided weights
-    if weights == Weights.KITTI:
-        return (0.001, 80.0)
-
-    if weights == Weights.NYU:
-        return (0.001, 10.0)
-
-    return (0.001, 10.0)
-
-
-def _make_dinov2_linear_depth_head(
-    *,
-    embed_dim: int,
-    layers: int,
-    min_depth: float,
-    max_depth: float,
-    **kwargs,
-):
-    if layers not in (1, 4):
-        raise AssertionError(f"Unsupported number of layers: {layers}")
-
-    if layers == 1:
-        in_index = [0]
-    else:
-        assert layers == 4
-        in_index = [0, 1, 2, 3]
-
-    return BNHead(
-        classify=True,
-        n_bins=256,
-        bins_strategy="UD",
-        norm_strategy="linear",
-        upsample=4,
-        in_channels=[embed_dim] * len(in_index),
-        in_index=in_index,
-        input_transform="resize_concat",
-        channels=embed_dim * len(in_index) * 2,
-        align_corners=False,
-        min_depth=0.001,
-        max_depth=80,
-        loss_decode=(),
-    )
-
-
-def _make_dinov2_linear_depther(
-    *,
-    arch_name: str = "vit_large",
-    layers: int = 4,
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.NYU,
-    depth_range: Optional[Tuple[float, float]] = None,
-    **kwargs,
-):
-    if layers not in (1, 4):
-        raise AssertionError(f"Unsupported number of layers: {layers}")
-    if isinstance(weights, str):
-        try:
-            weights = Weights[weights]
-        except KeyError:
-            raise AssertionError(f"Unsupported weights: {weights}")
-
-    if depth_range is None:
-        depth_range = _get_depth_range(pretrained, weights)
-    min_depth, max_depth = depth_range
-
-    backbone = _make_dinov2_model(arch_name=arch_name, pretrained=pretrained, **kwargs)
-
-    embed_dim = backbone.embed_dim
-    patch_size = backbone.patch_size
-    model_name = _make_dinov2_model_name(arch_name, patch_size)
-    linear_depth_head = _make_dinov2_linear_depth_head(
-        embed_dim=embed_dim,
-        layers=layers,
-        min_depth=min_depth,
-        max_depth=max_depth,
-    )
-
-    layer_count = {
-        "vit_small": 12,
-        "vit_base": 12,
-        "vit_large": 24,
-        "vit_giant2": 40,
-    }[arch_name]
-
-    if layers == 4:
-        out_index = {
-            "vit_small": [2, 5, 8, 11],
-            "vit_base": [2, 5, 8, 11],
-            "vit_large": [4, 11, 17, 23],
-            "vit_giant2": [9, 19, 29, 39],
-        }[arch_name]
-    else:
-        assert layers == 1
-        out_index = [layer_count - 1]
-
-    model = DepthEncoderDecoder(backbone=backbone, decode_head=linear_depth_head)
-    model.backbone.forward = partial(
-        backbone.get_intermediate_layers,
-        n=out_index,
-        reshape=True,
-        return_class_token=True,
-        norm=False,
-    )
-    model.backbone.register_forward_pre_hook(lambda _, x: CenterPadding(patch_size)(x[0]))
-
-    if pretrained:
-        layers_str = str(layers) if layers == 4 else ""
-        weights_str = weights.value.lower()
-        url = _DINOV2_BASE_URL + f"/{model_name}/{model_name}_{weights_str}_linear{layers_str}_head.pth"
-        checkpoint = torch.hub.load_state_dict_from_url(url, map_location="cpu")
-        if "state_dict" in checkpoint:
-            state_dict = checkpoint["state_dict"]
-        model.load_state_dict(state_dict, strict=False)
-
-    return model
-
-
-def dinov2_vits14_ld(*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
-    return _make_dinov2_linear_depther(
-        arch_name="vit_small", layers=layers, pretrained=pretrained, weights=weights, **kwargs
-    )
-
-
-def dinov2_vitb14_ld(*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
-    return _make_dinov2_linear_depther(
-        arch_name="vit_base", layers=layers, pretrained=pretrained, weights=weights, **kwargs
-    )
-
-
-def dinov2_vitl14_ld(*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
-    return _make_dinov2_linear_depther(
-        arch_name="vit_large", layers=layers, pretrained=pretrained, weights=weights, **kwargs
-    )
-
-
-def dinov2_vitg14_ld(*, layers: int = 4, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
-    return _make_dinov2_linear_depther(
-        arch_name="vit_giant2", layers=layers, ffn_layer="swiglufused", pretrained=pretrained, weights=weights, **kwargs
-    )
-
-
-def _make_dinov2_dpt_depth_head(*, embed_dim: int, min_depth: float, max_depth: float):
-    return DPTHead(
-        in_channels=[embed_dim] * 4,
-        channels=256,
-        embed_dims=embed_dim,
-        post_process_channels=[embed_dim // 2 ** (3 - i) for i in range(4)],
-        readout_type="project",
-        min_depth=min_depth,
-        max_depth=max_depth,
-        loss_decode=(),
-    )
-
-
-def _make_dinov2_dpt_depther(
-    *,
-    arch_name: str = "vit_large",
-    pretrained: bool = True,
-    weights: Union[Weights, str] = Weights.NYU,
-    depth_range: Optional[Tuple[float, float]] = None,
-    **kwargs,
-):
-    if isinstance(weights, str):
-        try:
-            weights = Weights[weights]
-        except KeyError:
-            raise AssertionError(f"Unsupported weights: {weights}")
-
-    if depth_range is None:
-        depth_range = _get_depth_range(pretrained, weights)
-    min_depth, max_depth = depth_range
-
-    backbone = _make_dinov2_model(arch_name=arch_name, pretrained=pretrained, **kwargs)
-
-    model_name = _make_dinov2_model_name(arch_name, backbone.patch_size)
-    dpt_depth_head = _make_dinov2_dpt_depth_head(embed_dim=backbone.embed_dim, min_depth=min_depth, max_depth=max_depth)
-
-    out_index = {
-        "vit_small": [2, 5, 8, 11],
-        "vit_base": [2, 5, 8, 11],
-        "vit_large": [4, 11, 17, 23],
-        "vit_giant2": [9, 19, 29, 39],
-    }[arch_name]
-
-    model = DepthEncoderDecoder(backbone=backbone, decode_head=dpt_depth_head)
-    model.backbone.forward = partial(
-        backbone.get_intermediate_layers,
-        n=out_index,
-        reshape=True,
-        return_class_token=True,
-        norm=False,
-    )
-    model.backbone.register_forward_pre_hook(lambda _, x: CenterPadding(backbone.patch_size)(x[0]))
-
-    if pretrained:
-        weights_str = weights.value.lower()
-        url = _DINOV2_BASE_URL + f"/{model_name}/{model_name}_{weights_str}_dpt_head.pth"
-        checkpoint = torch.hub.load_state_dict_from_url(url, map_location="cpu")
-        if "state_dict" in checkpoint:
-            state_dict = checkpoint["state_dict"]
-        model.load_state_dict(state_dict, strict=False)
-
-    return model
-
-
-def dinov2_vits14_dd(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
-    return _make_dinov2_dpt_depther(arch_name="vit_small", pretrained=pretrained, weights=weights, **kwargs)
-
-
-def dinov2_vitb14_dd(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
-    return _make_dinov2_dpt_depther(arch_name="vit_base", pretrained=pretrained, weights=weights, **kwargs)
-
-
-def dinov2_vitl14_dd(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
-    return _make_dinov2_dpt_depther(arch_name="vit_large", pretrained=pretrained, weights=weights, **kwargs)
-
-
-def dinov2_vitg14_dd(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.NYU, **kwargs):
-    return _make_dinov2_dpt_depther(
-        arch_name="vit_giant2", ffn_layer="swiglufused", pretrained=pretrained, weights=weights, **kwargs
-    )

OpenPath/dinov2/hub/dinotxt.py

@@ -1,77 +0,0 @@
-import torch
-import math
-
-from .backbones import dinov2_vitl14_reg
-from .utils import _DINOV2_BASE_URL
-
-
-def dinov2_vitl14_reg4_dinotxt_tet1280d20h24l():
-    from .text.dinotxt_model import DinoTxtConfig, DinoTxt
-    from .text.dinov2_wrapper import DINOv2Wrapper
-    from .text.text_transformer import TextTransformer
-
-    dinotxt_config = DinoTxtConfig(
-        embed_dim=2048,
-        vision_model_freeze_backbone=True,
-        vision_model_train_img_size=224,
-        vision_model_use_class_token=True,
-        vision_model_use_patch_tokens=True,
-        vision_model_num_head_blocks=2,
-        vision_model_head_blocks_drop_path=0.3,
-        vision_model_use_linear_projection=False,
-        vision_model_patch_tokens_pooler_type="mean",
-        vision_model_patch_token_layer=1,  # which layer to take patch tokens from
-        # 1 - last layer, 2 - second last layer, etc.
-        text_model_freeze_backbone=False,
-        text_model_num_head_blocks=0,
-        text_model_head_blocks_is_causal=False,
-        text_model_head_blocks_drop_prob=0.0,
-        text_model_tokens_pooler_type="argmax",
-        text_model_use_linear_projection=True,
-        init_logit_scale=math.log(1 / 0.07),
-        init_logit_bias=None,
-        freeze_logit_scale=False,
-    )
-    vision_backbone = DINOv2Wrapper(dinov2_vitl14_reg())
-    text_backbone = TextTransformer(
-        context_length=77,
-        vocab_size=49408,
-        dim=1280,
-        num_heads=20,
-        num_layers=24,
-        ffn_ratio=4,
-        is_causal=True,
-        ls_init_value=None,
-        dropout_prob=0.0,
-    )
-    model = DinoTxt(dinotxt_config, vision_backbone, text_backbone)
-    model.init_weights()
-    model.visual_model.backbone = vision_backbone
-    model.eval()
-
-    visual_model_head_state_dict = torch.hub.load_state_dict_from_url(
-        _DINOV2_BASE_URL + "/dinov2_vitl14/dinov2_vitl14_reg4_dinotxt_tet1280d20h24l_vision_head.pth",
-        map_location="cpu",
-    )
-    text_model_state_dict = torch.hub.load_state_dict_from_url(
-        _DINOV2_BASE_URL + "/dinov2_vitl14/dinov2_vitl14_reg4_dinotxt_tet1280d20h24l_text_encoder.pth",
-        map_location="cpu",
-    )
-    model.visual_model.head.load_state_dict(visual_model_head_state_dict, strict=True)
-    model.text_model.load_state_dict(text_model_state_dict, strict=True)
-    return model
-
-
-def get_tokenizer():
-    from .text.tokenizer import Tokenizer
-    import requests
-    from io import BytesIO
-
-    url = _DINOV2_BASE_URL + "/thirdparty/bpe_simple_vocab_16e6.txt.gz"
-    try:
-        response = requests.get(url)
-        response.raise_for_status()
-        file_buf = BytesIO(response.content)
-        return Tokenizer(vocab_path=file_buf)
-    except Exception as e:
-        raise FileNotFoundError(f"Failed to download file from url {url} with error last: {e}")

OpenPath/dinov2/hub/text/dinotxt_model.py

@@ -1,130 +0,0 @@
-import math
-from dataclasses import dataclass
-from typing import Optional, Tuple
-
-import torch
-import torch.nn.functional as F
-from torch import nn, Tensor
-
-from .vision_tower import VisionTower
-from .text_tower import TextTower
-
-
-@dataclass
-class DinoTxtConfig:
-    embed_dim: int
-    vision_model_freeze_backbone: bool = True
-    vision_model_train_img_size: int = 224
-    vision_model_use_class_token: bool = True
-    vision_model_use_patch_tokens: bool = False
-    vision_model_num_head_blocks: int = 0
-    vision_model_head_blocks_drop_path: float = 0.3
-    vision_model_use_linear_projection: bool = False
-    vision_model_patch_tokens_pooler_type: str = "mean"
-    vision_model_patch_token_layer: int = 1  # which layer to take patch tokens from
-    # 1 - last layer, 2 - second last layer, etc.
-    text_model_freeze_backbone: bool = False
-    text_model_num_head_blocks: int = 0
-    text_model_head_blocks_is_causal: bool = False
-    text_model_head_blocks_drop_prob: float = 0.0
-    text_model_tokens_pooler_type: str = "first"
-    text_model_use_linear_projection: bool = False
-    init_logit_scale: float = math.log(1 / 0.07)
-    init_logit_bias: Optional[float] = None
-    freeze_logit_scale: bool = False
-
-
-class DinoTxt(nn.Module):
-    def __init__(
-        self,
-        model_config: DinoTxtConfig,
-        vision_backbone: nn.Module,
-        text_backbone: nn.Module,
-    ):
-        super().__init__()
-        self.model_config = model_config
-        self.visual_model = VisionTower(
-            vision_backbone,
-            model_config.vision_model_freeze_backbone,
-            model_config.embed_dim,
-            model_config.vision_model_num_head_blocks,
-            model_config.vision_model_head_blocks_drop_path,
-            model_config.vision_model_use_class_token,
-            model_config.vision_model_use_patch_tokens,
-            model_config.vision_model_patch_token_layer,
-            model_config.vision_model_patch_tokens_pooler_type,
-            model_config.vision_model_use_linear_projection,
-        )
-        self.text_model = TextTower(
-            text_backbone,
-            model_config.text_model_freeze_backbone,
-            model_config.embed_dim,
-            model_config.text_model_num_head_blocks,
-            model_config.text_model_head_blocks_is_causal,
-            model_config.text_model_head_blocks_drop_prob,
-            model_config.text_model_tokens_pooler_type,
-            model_config.text_model_use_linear_projection,
-        )
-        self.logit_scale = nn.Parameter(torch.ones(1) * model_config.init_logit_scale)
-        if model_config.freeze_logit_scale:
-            self.logit_scale.requires_grad = False
-
-    def init_weights(self):
-        self.visual_model.init_weights()
-        self.text_model.init_weights()
-
-    def get_visual_class_and_patch_tokens(self, image: Tensor) -> Tuple[Tensor, Tensor]:
-        return self.visual_model.get_class_and_patch_tokens(image)
-
-    def encode_image(
-        self,
-        image: Tensor,
-        normalize: bool = False,
-    ) -> Tensor:
-        """
-        Encode an image into a vector descriptor containing both global and local features.
-
-        Args:
-            image (Tensor): Tensor of shape `(batch_size, rgb, height, width)`, normalized using ImageNet mean and std.
-            normalize (bool, optional): Whether to normalize the output vectors. Default is False.
-                Image features should always be normalized before comparing them with text features:
-        Returns:
-            Tensor: Tensor of shape `(batch_size, embed_dim)` containing the image features.
-                The first half of the vector corresponds to the global features (class token),
-                and the second half corresponds to the pooled patch features.
-        """
-        features = self.visual_model(image)
-        return F.normalize(features, dim=-1) if normalize else features
-
-    def encode_text(self, text: Tensor, normalize: bool = False) -> Tensor:
-        """
-        Encode a text input into a vector descriptor.
-
-        Args:
-            text (Tensor): Tensor of shape `(batch_size, seq_len)` containing token indices.
-            normalize (bool, optional): Whether to normalize the output vectors. Default is False.
-                Text features should be normalized before comparing them with image features:
-        Returns:
-            Tensor: Tensor of shape `(batch_size, embed_dim)` containing the text features.
-                As a consequence of the training procedure, assume that the first half of the tensor corresponds
-                to global image features and the second half to pooled patch features.
-        """
-        features = self.text_model(text)
-        return F.normalize(features, dim=-1) if normalize else features
-
-    def get_logits(self, image: Tensor, text: Tensor) -> Tuple[Tensor, Tensor]:
-        text_features = self.encode_text(text, normalize=True)
-        image_features = self.encode_image(image, normalize=True)
-        image_logits = self.logit_scale.exp() * image_features @ text_features.T
-        text_logits = image_logits.T
-        return image_logits, text_logits
-
-    def forward(
-        self,
-        image: Tensor,
-        text: Tensor,
-    ) -> Tuple[Tensor, Tensor, Tensor]:
-
-        text_features = self.encode_text(text, normalize=True)
-        image_features = self.encode_image(image, normalize=True)
-        return image_features, text_features, self.logit_scale.exp()

OpenPath/dinov2/hub/text/dinov2_wrapper.py

@@ -1,59 +0,0 @@
-from typing import Sequence
-
-import torch
-
-
-class DINOv2Wrapper(torch.nn.Module):
-    def __init__(self, model):
-        super().__init__()
-        self.model = model
-        self.embed_dim = model.embed_dim
-        self.num_heads = model.num_heads
-        self.num_register_tokens = model.num_register_tokens
-
-    # Same as the original forward, but assert is_training and rename x_norm_regtokens -> x_storage_tokens
-    def forward(self, img, is_training: bool):
-        assert is_training
-        H, W = img.shape[-2:]
-        P = self.model.patch_size
-        x_dict = self.model(img, is_training=True)
-        x_dict["h"] = h = H // P
-        x_dict["w"] = w = W // P
-        assert x_dict["x_norm_patchtokens"].shape[-2] == h * w
-        return x_dict
-
-    # Same as the original get_intermediate_layers, but allow returining extra tokens (registers)
-    def get_intermediate_layers(
-        self,
-        x: torch.Tensor,
-        n: int | Sequence[int] = 1,  # Layers or n last layers to take
-        reshape: bool = False,
-        return_class_token: bool = False,
-        return_register_tokens: bool = False,
-        norm=True,
-    ) -> tuple[torch.Tensor] | tuple[tuple[torch.Tensor, ...], ...]:
-        if self.model.chunked_blocks:
-            outputs = self.model._get_intermediate_layers_chunked(x, n)
-        else:
-            outputs = self.model._get_intermediate_layers_not_chunked(x, n)
-        if norm:
-            outputs = [self.model.norm(out) for out in outputs]
-        class_tokens = [out[:, 0] for out in outputs]
-        register_tokens = [out[:, 1 : 1 + self.model.num_register_tokens] for out in outputs]
-        outputs = [out[:, 1 + self.model.num_register_tokens :] for out in outputs]
-        if reshape:
-            B, _, h, w = x.shape
-            outputs = [
-                out.reshape(B, h // self.model.patch_size, w // self.model.patch_size, -1)
-                .permute(0, 3, 1, 2)
-                .contiguous()
-                for out in outputs
-            ]
-
-        if not return_class_token and not return_register_tokens:
-            return tuple(outputs)
-        if return_class_token and not return_register_tokens:
-            return tuple(zip(outputs, class_tokens))
-        if not return_class_token and return_register_tokens:
-            return tuple(zip(outputs, register_tokens))
-        return tuple(zip(outputs, class_tokens, register_tokens))

OpenPath/dinov2/hub/text/text_tower.py

@@ -1,99 +0,0 @@
-import torch
-from torch import nn, Tensor
-
-from dinov2.layers import (
-    CausalAttentionBlock,
-)
-
-
-class TextHead(nn.Module):
-    def __init__(
-        self,
-        input_dim: int,
-        embed_dim: int,
-        num_heads: int,
-        num_blocks: int,
-        block_drop_prob: float,
-        is_causal: bool,
-        use_linear_projection: bool,
-    ):
-        super().__init__()
-        block_list = [nn.Identity()]
-        self.ln_final = nn.Identity()
-        if num_blocks > 0:
-            block_list = [
-                CausalAttentionBlock(
-                    dim=input_dim,
-                    num_heads=num_heads,
-                    is_causal=is_causal,
-                    dropout_prob=block_drop_prob,
-                )
-                for _ in range(num_blocks)
-            ]
-            self.ln_final = nn.LayerNorm(input_dim)
-        self.block_list = nn.ModuleList(block_list)
-        self.num_blocks = num_blocks
-        self.linear_projection = nn.Identity()
-        if input_dim != embed_dim or use_linear_projection:
-            self.linear_projection = nn.Linear(input_dim, embed_dim, bias=False)
-
-    def init_weights(self):
-        if self.num_blocks > 0:
-            for i in range(self.num_blocks):
-                self.block_list[i].init_weights()
-            self.ln_final.reset_parameters()
-        if isinstance(self.linear_projection, nn.Linear):
-            nn.init.normal_(self.linear_projection.weight, std=self.linear_projection.in_features**-0.5)
-
-    def forward(self, text_tokens: Tensor) -> Tensor:
-        for block in self.block_list:
-            text_tokens = block(text_tokens)
-        text_tokens = self.ln_final(text_tokens)
-        return self.linear_projection(text_tokens)
-
-
-class TextTower(nn.Module):
-    def __init__(
-        self,
-        backbone: nn.Module,
-        freeze_backbone: bool,
-        embed_dim: int,
-        num_head_blocks: int,
-        head_blocks_is_causal: bool,
-        head_blocks_block_drop_prob: float,
-        tokens_pooler_type: str,
-        use_linear_projection: bool,
-    ):
-        super().__init__()
-        self.backbone = backbone
-        self.freeze_backbone = freeze_backbone
-        backbone_out_dim = backbone.embed_dim
-        self.backbone = backbone
-        self.head = TextHead(
-            backbone_out_dim,
-            embed_dim,
-            self.backbone.num_heads,
-            num_head_blocks,
-            head_blocks_block_drop_prob,
-            head_blocks_is_causal,
-            use_linear_projection,
-        )
-        self.tokens_pooler_type = tokens_pooler_type
-
-    def init_weights(self):
-        self.backbone.init_weights()
-        self.head.init_weights()
-
-    def forward(self, token_indices: Tensor) -> Tensor:
-        text_tokens = self.backbone(token_indices)
-        text_tokens = self.head(text_tokens)
-        if self.tokens_pooler_type == "first":
-            features = text_tokens[:, 0]
-        elif self.tokens_pooler_type == "last":
-            features = text_tokens[:, -1]
-        elif self.tokens_pooler_type == "argmax":
-            assert token_indices is not None
-            features = text_tokens[torch.arange(text_tokens.shape[0]), token_indices.argmax(dim=-1)]
-        else:
-            raise ValueError(f"Unknown text tokens pooler type: {self.pooler_type}")
-        return features

OpenPath/dinov2/hub/text/text_transformer.py

@@ -1,67 +0,0 @@
-from typing import Callable, Optional, Tuple
-
-import torch
-from torch import nn, Tensor
-
-
-from dinov2.layers import CausalAttentionBlock
-
-
-class TextTransformer(nn.Module):
-    def __init__(
-        self,
-        context_length: int,
-        vocab_size: int,
-        dim: int,
-        num_heads: int,
-        num_layers: int,
-        ffn_ratio: float,
-        is_causal: bool,
-        ls_init_value: Optional[float] = None,
-        act_layer: Callable = nn.GELU,
-        norm_layer: Callable = nn.LayerNorm,
-        dropout_prob: float = 0.0,
-    ):
-        super().__init__()
-        self.vocab_size = vocab_size
-        self.embed_dim = dim
-        self.num_heads = num_heads
-
-        self.token_embedding = nn.Embedding(vocab_size, dim)
-        self.positional_embedding = nn.Parameter(torch.empty(context_length, dim))
-        self.dropout = nn.Dropout(dropout_prob)
-        self.num_layers = num_layers
-        block_list = [
-            CausalAttentionBlock(
-                dim=dim,
-                num_heads=num_heads,
-                ffn_ratio=ffn_ratio,
-                ls_init_value=ls_init_value,
-                is_causal=is_causal,
-                act_layer=act_layer,
-                norm_layer=norm_layer,
-                dropout_prob=dropout_prob,
-            )
-            for _ in range(num_layers)
-        ]
-        self.blocks = nn.ModuleList(block_list)
-        self.ln_final = norm_layer(dim)
-
-    def init_weights(self):
-        nn.init.normal_(self.token_embedding.weight, std=0.02)
-        nn.init.normal_(self.positional_embedding, std=0.01)
-        init_attn_std = self.embed_dim**-0.5
-        init_proj_std = (self.embed_dim**-0.5) * ((2 * self.num_layers) ** -0.5)
-        init_fc_std = (2 * self.embed_dim) ** -0.5
-        for block in self.blocks:
-            block.init_weights(init_attn_std, init_proj_std, init_fc_std)
-        self.ln_final.reset_parameters()
-
-    def forward(self, token_indices: Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        _, N = token_indices.size()
-        x = self.token_embedding(token_indices) + self.positional_embedding[:N]
-        x = self.dropout(x)
-        for block in self.blocks:
-            x = block(x)
-        x = self.ln_final(x)
-        return x

OpenPath/dinov2/hub/text/tokenizer.py

@@ -1,40 +0,0 @@
-import torch
-from typing import List, Union
-
-
-from dinov2.thirdparty.CLIP.clip.simple_tokenizer import SimpleTokenizer
-
-
-class Tokenizer(SimpleTokenizer):
-    def __init__(self, vocab_path: str):
-        SimpleTokenizer.__init__(self, bpe_path=vocab_path)
-
-    def tokenize(self, texts: Union[str, List[str]], context_length: int = 77) -> torch.LongTensor:
-        """
-        Returns the tokenized representation of given input string(s)
-
-        Parameters
-        ----------
-        texts : Union[str, List[str]]
-            An input string or a list of input strings to tokenize
-        context_length : int
-            The context length to use; all CLIP models use 77 as the context length
-
-        Returns
-        -------
-        A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
-        """
-        if isinstance(texts, str):
-            texts = [texts]
-        sot_token = self.encoder["<|startoftext|>"]
-        eot_token = self.encoder["<|endoftext|>"]
-        all_tokens = [[sot_token] + self.encode(text) + [eot_token] for text in texts]
-        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
-
-        for i, tokens in enumerate(all_tokens):
-            if len(tokens) > context_length:
-                tokens = tokens[:context_length]  # Truncate
-                tokens[-1] = eot_token
-            result[i, : len(tokens)] = torch.tensor(tokens)
-
-        return result

OpenPath/dinov2/hub/text/vision_tower.py

@@ -1,174 +0,0 @@
-from functools import partial
-from typing import Callable, Tuple
-
-import torch
-from torch import nn, Tensor
-
-from dinov2.layers import (
-    LayerScale,
-    NestedTensorBlock as AttentionBlock,
-    SwiGLUFFNAligned as SwiGLUFFN,
-)
-
-
-def init_weights_vit_timm(module: nn.Module, name: str = ""):
-    """ViT weight initialization, original timm impl (for reproducibility)"""
-    if isinstance(module, nn.Linear):
-        nn.init.trunc_normal_(module.weight, std=0.02)
-        if module.bias is not None:
-            nn.init.zeros_(module.bias)
-    if isinstance(module, nn.LayerNorm):
-        module.reset_parameters()
-    if isinstance(module, LayerScale):
-        module.reset_parameters()
-    if isinstance(module, nn.Conv2d):
-        module.reset_parameters()
-
-
-def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
-    if not depth_first and include_root:
-        fn(module=module, name=name)
-    for child_name, child_module in module.named_children():
-        child_name = ".".join((name, child_name)) if name else child_name
-        named_apply(
-            fn=fn,
-            module=child_module,
-            name=child_name,
-            depth_first=depth_first,
-            include_root=True,
-        )
-    if depth_first and include_root:
-        fn(module=module, name=name)
-    return module
-
-
-class VisionHead(nn.Module):
-    def __init__(
-        self,
-        input_dim: int,
-        embed_dim: int,
-        num_heads: int,
-        num_blocks: int,
-        blocks_drop_path: float,
-        use_class_token: bool,
-        use_patch_tokens: bool,
-        use_linear_projection: bool,
-    ):
-        super().__init__()
-        block_list = [nn.Identity()]
-        self.ln_final = nn.Identity()
-        if num_blocks > 0:
-            block_list = [
-                AttentionBlock(
-                    input_dim,
-                    num_heads,
-                    ffn_layer=partial(SwiGLUFFN, align_to=64),
-                    init_values=1e-5,
-                    drop_path=blocks_drop_path,
-                )
-                for _ in range(num_blocks)
-            ]
-            self.ln_final = nn.LayerNorm(input_dim)
-        self.block_list = nn.ModuleList(block_list)
-        self.num_blocks = num_blocks
-        multiplier = 2 if use_class_token and use_patch_tokens else 1
-        self.linear_projection = nn.Identity()
-        if multiplier * input_dim != embed_dim or use_linear_projection:
-            assert embed_dim % multiplier == 0, f"Expects {embed_dim} to be divisible by {multiplier}"
-            self.linear_projection = nn.Linear(input_dim, embed_dim // multiplier, bias=False)
-
-    def init_weights(self):
-        if self.num_blocks > 0:
-            for i in range(self.num_blocks):
-                block = self.block_list[i]
-                named_apply(init_weights_vit_timm, block)
-            self.ln_final.reset_parameters()
-        if isinstance(self.linear_projection, nn.Linear):
-            nn.init.normal_(self.linear_projection.weight, std=self.linear_projection.in_features**-0.5)
-
-    def forward(self, image_tokens: Tensor) -> Tensor:
-        for block in self.block_list:
-            image_tokens = block(image_tokens)
-        image_tokens = self.ln_final(image_tokens)
-        return self.linear_projection(image_tokens)
-
-
-class VisionTower(nn.Module):
-    def __init__(
-        self,
-        backbone: nn.Module,
-        freeze_backbone: bool,
-        embed_dim: int,
-        num_head_blocks: int,
-        head_blocks_block_drop_path: float,
-        use_class_token: bool,
-        use_patch_tokens: bool,
-        patch_token_layer: int,
-        patch_tokens_pooler_type: str,
-        use_linear_projection: bool,
-    ):
-        super().__init__()
-        self.backbone = backbone
-        self.freeze_backbone = freeze_backbone
-        self.use_class_token = use_class_token
-        self.use_patch_tokens = use_patch_tokens
-        self.patch_token_layer = patch_token_layer
-        self.patch_tokens_pooler_type = patch_tokens_pooler_type
-        self.num_register_tokens = 0
-        if hasattr(self.backbone, "num_register_tokens"):
-            self.num_register_tokens = self.backbone.num_register_tokens
-        elif hasattr(self.backbone, "n_storage_tokens"):
-            self.num_register_tokens = self.backbone.n_storage_tokens
-        backbone_out_dim = self.backbone.embed_dim
-        self.head = VisionHead(
-            backbone_out_dim,
-            embed_dim,
-            self.backbone.num_heads,
-            num_head_blocks,
-            head_blocks_block_drop_path,
-            use_class_token,
-            use_patch_tokens,
-            use_linear_projection,
-        )
-
-    def init_weights(self):
-        if not self.freeze_backbone:
-            self.backbone.init_weights()
-        self.head.init_weights()
-
-    def get_backbone_features(self, images: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
-        tokens = self.backbone.get_intermediate_layers(
-            images,
-            n=self.patch_token_layer,
-            return_class_token=True,
-            return_register_tokens=True,
-        )
-        class_token = tokens[-1][1]
-        patch_tokens = tokens[0][0]
-        register_tokens = tokens[0][2]
-        return class_token, patch_tokens, register_tokens
-
-    def get_class_and_patch_tokens(self, images: Tensor) -> Tuple[Tensor, Tensor]:
-        class_token, patch_tokens, register_tokens = self.get_backbone_features(images)
-        image_tokens = self.head(torch.cat([class_token.unsqueeze(1), register_tokens, patch_tokens], dim=1))
-        class_token, patch_tokens = image_tokens[:, 0], image_tokens[:, self.num_register_tokens + 1 :]
-        return class_token, patch_tokens
-
-    def forward(self, images: Tensor) -> Tensor:
-        class_token, patch_tokens = self.get_class_and_patch_tokens(images)
-        features = []
-        if self.use_class_token:
-            features.append(class_token)
-        if self.use_patch_tokens:
-            if self.patch_tokens_pooler_type == "mean":
-                features.append(torch.mean(patch_tokens, dim=1))
-            elif self.patch_tokens_pooler_type == "max":
-                features.append(torch.max(patch_tokens, dim=1).values)
-            elif self.patch_tokens_pooler_type == "gem":
-                power = 3
-                eps = 1e-6
-                patch_tokens_power = patch_tokens.clamp(min=eps).pow(power)
-                features.append(torch.mean(patch_tokens_power, dim=1).pow(1 / power))
-            else:
-                raise ValueError(f"Unknown patch tokens pooler type: {self.patch_tokens_pooler_type}")
-        return torch.cat(features, dim=-1)

OpenPath/dinov2/hub/utils.py

@@ -1,39 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the Apache License, Version 2.0
-# found in the LICENSE file in the root directory of this source tree.
-
-import itertools
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-_DINOV2_BASE_URL = "https://dl.fbaipublicfiles.com/dinov2"
-
-
-def _make_dinov2_model_name(arch_name: str, patch_size: int, num_register_tokens: int = 0) -> str:
-    compact_arch_name = arch_name.replace("_", "")[:4]
-    registers_suffix = f"_reg{num_register_tokens}" if num_register_tokens else ""
-    return f"dinov2_{compact_arch_name}{patch_size}{registers_suffix}"
-
-
-class CenterPadding(nn.Module):
-    def __init__(self, multiple):
-        super().__init__()
-        self.multiple = multiple
-
-    def _get_pad(self, size):
-        new_size = math.ceil(size / self.multiple) * self.multiple
-        pad_size = new_size - size
-        pad_size_left = pad_size // 2
-        pad_size_right = pad_size - pad_size_left
-        return pad_size_left, pad_size_right
-
-    @torch.inference_mode()
-    def forward(self, x):
-        pads = list(itertools.chain.from_iterable(self._get_pad(m) for m in x.shape[:1:-1]))
-        output = F.pad(x, pads)
-        return output

OpenPath/dinov2/hub/xray_dino/backbones.py

@@ -1,28 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the licence
-# found in the LICENSE_XRAY_DINO_MODEL file in the root directory of this source tree.
-
-from typing import Union
-
-from ..backbones import Weights, _make_dinov2_model
-
-
-def xray_dino_vitl16(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.XRAY_DINO, **kwargs):
-    """
-    XRay-DINO ViT-L/16 model (optionally) pretrained on the XRay-DINO dataset.
-    """
-    return _make_dinov2_model(
-        arch_name="vit_large",
-        patch_size=16,
-        img_size=512,
-        num_register_tokens=0,
-        interpolate_antialias=False,
-        interpolate_offset=0.1,
-        block_chunks=4,
-        pretrained=pretrained,
-        weights=weights,
-        hash="ad31c2b0",
-        check_hash=True,
-        **kwargs,
-    )

OpenPath/dinov2/loss/gram_loss.py

@@ -1,7 +1,13 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright (c) 2026 OpenPath authors.
 #
-# This software may be used and distributed in accordance with
-# the terms of the DINOv3 License Agreement.
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+#
+# Clean-room re-implementation of the gram-anchoring loss (technique from DINOv3,
+# Siméoni et al., 2025). Written from the mathematical description only — the MSE
+# between the (optionally L2-normalized) patch-token Gram / self-similarity matrices
+# of the student and a frozen anchor — with no reference to the original DINOv3
+# source, so this file carries the same Apache-2.0 license as the rest of the fork.
 
 import torch
 import torch.nn as nn
@@ -9,76 +15,54 @@ import torch.nn.functional as F
 
 
 class GramLoss(nn.Module):
-    """Implementation of the gram loss"""
-
-    def __init__(
-        self,
-        apply_norm=True,
-        img_level=True,
-        remove_neg=True,
-        remove_only_teacher_neg=False,
-    ):
+    """Gram-anchoring loss.
+
+    For patch-token features from the student and a frozen anchor, build each model's
+    per-image Gram matrix (token-by-token self-similarity) and penalize their squared
+    difference. This anchors the student's *relational* (dense/patch) structure to the
+    anchor's while DINO/iBOT keep optimizing the global CLS representation.
+
+    Args:
+        apply_norm: L2-normalize tokens along the feature dim before the Gram product,
+            so each Gram entry is a cosine similarity in [-1, 1].
+        img_level: build one Gram per image, shape (B, N, N) (default path).
+        remove_neg: clamp negative similarities to 0 before the comparison.
+        remove_only_teacher_neg: when clamping, clamp the anchor's Gram only.
+    """
+
+    def __init__(self, apply_norm=True, img_level=True, remove_neg=True,
+                 remove_only_teacher_neg=False):
         super().__init__()
-
-        # Loss
-        self.mse_loss = torch.nn.MSELoss()
-
-        # Parameters
         self.apply_norm = apply_norm
+        self.img_level = img_level
         self.remove_neg = remove_neg
         self.remove_only_teacher_neg = remove_only_teacher_neg
 
-        if self.remove_neg or self.remove_only_teacher_neg:
-            assert self.remove_neg != self.remove_only_teacher_neg
-
-    def forward(self, output_feats, target_feats, img_level=True):
-        """Compute the MSE loss between the gram matrix of the input and target features.
-
-        Args:
-            output_feats: Pytorch tensor (B, N, dim) or (B*N, dim) if img_level == False
-            target_feats: Pytorch tensor (B, N, dim) or (B*N, dim) if img_level == False
-            img_level: bool, if true gram computed at the image level only else over the entire batch
-        Returns:
-            loss: scalar
-        """
+    @staticmethod
+    def _gram(x):
+        # x: (..., N, D) -> (..., N, N) self-similarity
+        return x @ x.transpose(-1, -2)
 
-        # Dimensions of the tensor should be (B, N, dim)
-        if img_level:
-            assert len(target_feats.shape) == 3 and len(output_feats.shape) == 3
+    def forward(self, student_tokens, teacher_tokens, img_level=None):
+        img_level = self.img_level if img_level is None else img_level
 
-        # Float casting
-        output_feats = output_feats.float()
-        target_feats = target_feats.float()
-
-        # SSL correlation
-        if self.apply_norm:
-            target_feats = F.normalize(target_feats, dim=-1)
-
-        if not img_level and len(target_feats.shape) == 3:
-            # Flatten (B, N, D) into  (B*N, D)
-            target_feats = target_feats.flatten(0, 1)
-
-        # Compute similarities
-        target_sim = torch.matmul(target_feats, target_feats.transpose(-1, -2))
-
-        # Patch correlation
+        xs = student_tokens.float()
+        xt = teacher_tokens.float()
         if self.apply_norm:
-            output_feats = F.normalize(output_feats, dim=-1)
+            xs = F.normalize(xs, dim=-1)
+            xt = F.normalize(xt, dim=-1)
 
-        if not img_level and len(output_feats.shape) == 3:
-            # Flatten (B, N, D) into  (B*N, D)
-            output_feats = output_feats.flatten(0, 1)
+        if not img_level:
+            # collapse the batch into one token set (rarely used)
+            xs = xs.reshape(-1, xs.shape[-1])
+            xt = xt.reshape(-1, xt.shape[-1])
 
-        # Compute similarities
-        student_sim = torch.matmul(output_feats, output_feats.transpose(-1, -2))
+        gs = self._gram(xs)
+        gt = self._gram(xt)
 
         if self.remove_neg:
-            target_sim[target_sim < 0] = 0.0
-            student_sim[student_sim < 0] = 0.0
-
-        elif self.remove_only_teacher_neg:
-            # Remove only the negative sim values of the teacher
-            target_sim[target_sim < 0] = 0.0
-            student_sim[(student_sim < 0) & (target_sim < 0)] = 0.0
+            gt = gt.clamp(min=0.0)
+            if not self.remove_only_teacher_neg:
+                gs = gs.clamp(min=0.0)
 
-        return self.mse_loss(student_sim, target_sim)
+        return F.mse_loss(gs, gt)

OpenPath/dinov2/train/ssl_meta_arch.py

@@ -123,7 +123,7 @@ class SSLMetaArch(nn.Module):
         for p in self.teacher.parameters():
             p.requires_grad = False
 
-        # ---- Gram anchoring (DINOv3 기법 이식) ----
+        # ---- Gram anchoring (technique from DINOv3; loss re-implemented clean-room, Apache-2.0) ----
         # 얼린 앵커(=좋은 peak 체크포인트)의 patch 구조에 student를 붙잡아 학습 후반 degradation 방지.
         self.do_gram = bool(getattr(cfg, "gram", None)) and bool(cfg.gram.use_loss)
         if self.do_gram:

README.md

@@ -29,7 +29,7 @@ The corpus and checkpoints are hosted separately (see below).
 > See [Evaluation](#evaluation).
 
 - **Encoder:** ViT-g/14 (reg4), 1536-dim CLS embedding
-- **Objective:** DINO + iBOT + KDE (DINOv2) with **gram anchoring** ported from DINOv3
+- **Objective:** DINO + iBOT + KDE (DINOv2) with **gram anchoring** (technique from DINOv3, re-implemented)
 - **Data:** public pathology WSIs only (TCGA, TCIA, GTEx, CAMELYON, ACROBAT, SurGen, …), re-tiled at native 40×
 - **Warm start:** Meta DINOv2 ViT-g/14-reg
 - **Training:** FSDP (SHARD_GRAD_OP), bf16, flat learning-rate schedule, 40× B200 (multi-node)
@@ -228,5 +228,12 @@ of the Republic of Korea, Ministry of Science and ICT; and the National Research
 
 ## License
 
-Code: derived from the OpenMidnight / DINOv2 fork — see `OpenPath/LICENSE`.
-Training data: public pathology datasets under CC-BY / CC0 / NIH-open terms (redistributable).
+**Code — Apache-2.0.** This repository is a fork of **DINOv2 / OpenMidnight** (both Apache-2.0); see
+`OpenPath/LICENSE`. The gram-anchoring loss (`OpenPath/dinov2/loss/gram_loss.py`) is a **clean-room
+re-implementation** of the DINOv3 technique — written from its mathematical description and verified
+to be numerically equivalent — so it is Apache-2.0 as well, and the codebase contains **no
+non-commercial (DINOv3-licensed) code**.
+
+**Weights — Apache-2.0** (warm-started from Meta DINOv2 ViT-g/14-reg, itself Apache-2.0).
+
+**Training data:** public pathology datasets under CC-BY / CC0 / NIH-open terms (redistributable).