File size: 4,730 Bytes
80df7e4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """Named SPECTRE architectures and their published weights.
This module is vendored into the Hugging Face Hub repository by `scripts/export_hf.py`, so it
must import nothing beyond the standard library.
"""
from __future__ import annotations
import copy
import difflib
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
__all__ = [
'SpectrePreset',
'PRESETS',
'BACKBONE_KWARGS',
'COMBINER_KWARGS',
'get_preset',
'list_presets',
'list_pretrained',
]
_HF_WEIGHTS = "https://huggingface.co/cclaess/SPECTRE/resolve/main/{}?download=true"
# The only place the SPECTRE architecture kwargs are written down. `spectre.model`, the Hugging
# Face SpectreConfig and scripts/export_hf.py all resolve back to here rather than restating them.
# The backbone sees raw CT crops; the feature combiner sees the backbone's per-crop embeddings,
# hence the different RoPE bases.
BACKBONE_KWARGS: Dict[str, Any] = {
"num_classes": 0,
"global_pool": '',
"pos_embed": "rope",
"rope_kwargs": {"base": 1000.0},
"init_values": 1.0,
}
COMBINER_KWARGS: Dict[str, Any] = {
"num_classes": 0,
"global_pool": '',
"pos_embed": "rope",
"rope_kwargs": {"base": 100.0},
"init_values": 1.0,
}
@dataclass(frozen=True)
class SpectrePreset:
"""One named SPECTRE architecture, optionally with published weights.
There is deliberately no crop size here: it is `backbone.patch_embed.img_size`, so the
architecture stays the single source of truth.
"""
name: str
backbone: str
feature_combiner: str
backbone_kwargs: Dict[str, Any] = field(default_factory=dict)
feature_combiner_kwargs: Dict[str, Any] = field(default_factory=dict)
backbone_weights: Optional[str] = None
feature_combiner_weights: Optional[str] = None
description: str = ""
@property
def has_pretrained_weights(self) -> bool:
return self.backbone_weights is not None
PRESETS: Dict[str, SpectrePreset] = {
"spectre-small": SpectrePreset(
name="spectre-small",
backbone="vit_small_patch16_128",
feature_combiner="feat_vit_small",
backbone_kwargs=copy.deepcopy(BACKBONE_KWARGS),
feature_combiner_kwargs=copy.deepcopy(COMBINER_KWARGS),
description="SPECTRE with a ViT-Small backbone. No published weights yet.",
),
"spectre-base": SpectrePreset(
name="spectre-base",
backbone="vit_base_patch16_128",
feature_combiner="feat_vit_base",
backbone_kwargs=copy.deepcopy(BACKBONE_KWARGS),
feature_combiner_kwargs=copy.deepcopy(COMBINER_KWARGS),
description="SPECTRE with a ViT-Base backbone. No published weights yet.",
),
"spectre-large": SpectrePreset(
name="spectre-large",
backbone="vit_large_patch16_128",
feature_combiner="feat_vit_large",
backbone_kwargs=copy.deepcopy(BACKBONE_KWARGS),
feature_combiner_kwargs=copy.deepcopy(COMBINER_KWARGS),
backbone_weights=_HF_WEIGHTS.format("spectre_backbone_vit_large_patch16_128.pt"),
feature_combiner_weights=_HF_WEIGHTS.format("spectre_combiner_feature_vit_large.pt"),
description="SPECTRE with a ViT-Large backbone, pretrained with SSL + vision-language alignment.",
),
}
def list_presets() -> List[str]:
"""Every known architecture name."""
return list(PRESETS)
def list_pretrained() -> List[str]:
"""The architecture names that have published weights."""
return [name for name, preset in PRESETS.items() if preset.has_pretrained_weights]
def get_preset(name: str) -> SpectrePreset:
"""Look up a preset by name.
The returned preset owns deep copies of its kwargs, so callers can mutate them without
corrupting the registry for the rest of the process.
"""
try:
preset = PRESETS[name]
except KeyError:
suggestion = difflib.get_close_matches(str(name), list(PRESETS), n=1)
hint = f" Did you mean {suggestion[0]!r}?" if suggestion else ""
raise ValueError(
f"Unknown SPECTRE preset {name!r}. Available: {list_presets()}.{hint}"
) from None
return SpectrePreset(
name=preset.name,
backbone=preset.backbone,
feature_combiner=preset.feature_combiner,
backbone_kwargs=copy.deepcopy(preset.backbone_kwargs),
feature_combiner_kwargs=copy.deepcopy(preset.feature_combiner_kwargs),
backbone_weights=preset.backbone_weights,
feature_combiner_weights=preset.feature_combiner_weights,
description=preset.description,
)
|