Spaces:
Sleeping
Sleeping
File size: 1,672 Bytes
4687353 | 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 | 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
|