Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| from typing import List, Tuple | |
| import timm | |
| import torch | |
| from PIL import Image | |
| from timm.data import create_transform, resolve_data_config | |
| class ImageBackbone(torch.nn.Module): | |
| def __init__(self, model_name: str, pretrained: bool = True) -> None: | |
| super().__init__() | |
| self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=0) | |
| cfg = resolve_data_config(self.model.pretrained_cfg, model=self.model) | |
| self.transform = create_transform(**cfg, is_training=False) | |
| self.output_dim = getattr(self.model, "num_features", None) | |
| if self.output_dim is None: | |
| raise RuntimeError(f"Could not infer output dim for {model_name}") | |
| def forward_pil(self, images: List[Image.Image], device: torch.device) -> torch.Tensor: | |
| batch = torch.stack([self.transform(img) for img in images], dim=0).to(device) | |
| return self.model(batch) | |
| def build_backbones( | |
| use_clip: bool, | |
| use_dino: bool, | |
| clip_model_name: str, | |
| dino_model_name: str, | |
| clip_pretrained: bool, | |
| dino_pretrained: bool, | |
| ) -> Tuple[ImageBackbone | None, ImageBackbone | None, int]: | |
| clip = None | |
| dino = None | |
| feature_dim = 0 | |
| if use_clip: | |
| clip = ImageBackbone(model_name=clip_model_name, pretrained=clip_pretrained) | |
| feature_dim += int(clip.output_dim) | |
| if use_dino: | |
| dino = ImageBackbone(model_name=dino_model_name, pretrained=dino_pretrained) | |
| feature_dim += int(dino.output_dim) | |
| if feature_dim <= 0: | |
| raise RuntimeError("At least one encoder must be enabled for baseline training.") | |
| return clip, dino, feature_dim | |