Initial media depth project backup
This commit is contained in:
@@ -0,0 +1,39 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user