torch_em.model.util

Shared configuration for vision transformer models and raw transforms.

View Source

 1"""Shared configuration for vision transformer models and raw transforms."""
 2
 3from dataclasses import dataclass
 4from typing import Optional, Tuple
 5
 6
 7@dataclass(frozen=True)
 8class ViTConfig:
 9    """Defaults shared by ViT construction, normalization, and UNETR."""
10
11    models: Tuple[str, ...]
12    img_size: int
13    stats: Optional[str]
14    dependency: str
15    clip_max: int = 1
16    skip_models: Tuple[str, ...] = ()
17    weight_suffixes: bool = False
18
19
20VIT_CONFIGS = {
21    "sam": ViTConfig(
22        models=("vit_b", "vit_l", "vit_h"), img_size=1024, stats="sam", dependency="sam", clip_max=255,
23        skip_models=("vit_b", "vit_l", "vit_h"), weight_suffixes=True
24    ),
25    "sam2": ViTConfig(
26        models=("hvit_t", "hvit_s", "hvit_b", "hvit_l"), img_size=1024, stats="sam", dependency="sam2",
27        weight_suffixes=True
28    ),
29    "sam3": ViTConfig(models=("vit_pe",), img_size=1008, stats="sam", dependency="sam3", weight_suffixes=True),
30    "cellpose_sam": ViTConfig(models=("vit_l",), img_size=256, stats=None, dependency="sam"),
31    "mae": ViTConfig(
32        models=("vit_b", "vit_l", "vit_h"), img_size=224, stats="mae", dependency="timm", skip_models=("vit_b", "vit_l")
33    ),
34    "scalemae": ViTConfig(models=("vit_b", "vit_l", "vit_h"), img_size=224, stats="mae", dependency="timm"),
35    "dinov2": ViTConfig(
36        models=("vit_s", "vit_b", "vit_l", "vit_g", "vit_s_reg4", "vit_b_reg4", "vit_l_reg4", "vit_g_reg4"),
37        img_size=518, stats="dino", dependency="dinov2"
38    ),
39    "dinov3": ViTConfig(
40        models=("vit_s", "vit_s+", "vit_b", "vit_l", "vit_l+", "vit_h+", "vit_7b"),
41        img_size=224, stats="dino", dependency="dinov3",
42        skip_models=("vit_s", "vit_s+", "vit_b", "vit_l", "vit_l+", "vit_h+", "vit_7b")
43    ),
44    "torchvision": ViTConfig(
45        models=("vit_b_16", "vit_b_32", "vit_l_16", "vit_l_32", "vit_h_14"),
46        img_size=224, stats="imagenet", dependency="torchvision", skip_models=("vit_b_16", "vit_l_16")
47    ),
48}

@dataclass(frozen=True)

class ViTConfig: View Source

 8@dataclass(frozen=True)
 9class ViTConfig:
10    """Defaults shared by ViT construction, normalization, and UNETR."""
11
12    models: Tuple[str, ...]
13    img_size: int
14    stats: Optional[str]
15    dependency: str
16    clip_max: int = 1
17    skip_models: Tuple[str, ...] = ()
18    weight_suffixes: bool = False

Defaults shared by ViT construction, normalization, and UNETR.

ViTConfig( models: Tuple[str, ...], img_size: int, stats: Optional[str], dependency: str, clip_max: int = 1, skip_models: Tuple[str, ...] = (), weight_suffixes: bool = False)

models: Tuple[str, ...]

img_size: int

stats: Optional[str]

dependency: str

clip_max: int = 1

skip_models: Tuple[str, ...] = ()

weight_suffixes: bool = False

VIT_CONFIGS = {'sam': ViTConfig(models=('vit_b', 'vit_l', 'vit_h'), img_size=1024, stats='sam', dependency='sam', clip_max=255, skip_models=('vit_b', 'vit_l', 'vit_h'), weight_suffixes=True), 'sam2': ViTConfig(models=('hvit_t', 'hvit_s', 'hvit_b', 'hvit_l'), img_size=1024, stats='sam', dependency='sam2', clip_max=1, skip_models=(), weight_suffixes=True), 'sam3': ViTConfig(models=('vit_pe',), img_size=1008, stats='sam', dependency='sam3', clip_max=1, skip_models=(), weight_suffixes=True), 'cellpose_sam': ViTConfig(models=('vit_l',), img_size=256, stats=None, dependency='sam', clip_max=1, skip_models=(), weight_suffixes=False), 'mae': ViTConfig(models=('vit_b', 'vit_l', 'vit_h'), img_size=224, stats='mae', dependency='timm', clip_max=1, skip_models=('vit_b', 'vit_l'), weight_suffixes=False), 'scalemae': ViTConfig(models=('vit_b', 'vit_l', 'vit_h'), img_size=224, stats='mae', dependency='timm', clip_max=1, skip_models=(), weight_suffixes=False), 'dinov2': ViTConfig(models=('vit_s', 'vit_b', 'vit_l', 'vit_g', 'vit_s_reg4', 'vit_b_reg4', 'vit_l_reg4', 'vit_g_reg4'), img_size=518, stats='dino', dependency='dinov2', clip_max=1, skip_models=(), weight_suffixes=False), 'dinov3': ViTConfig(models=('vit_s', 'vit_s+', 'vit_b', 'vit_l', 'vit_l+', 'vit_h+', 'vit_7b'), img_size=224, stats='dino', dependency='dinov3', clip_max=1, skip_models=('vit_s', 'vit_s+', 'vit_b', 'vit_l', 'vit_l+', 'vit_h+', 'vit_7b'), weight_suffixes=False), 'torchvision': ViTConfig(models=('vit_b_16', 'vit_b_32', 'vit_l_16', 'vit_l_32', 'vit_h_14'), img_size=224, stats='imagenet', dependency='torchvision', clip_max=1, skip_models=('vit_b_16', 'vit_l_16'), weight_suffixes=False)}