birder-clip 0.0.2.dev0__py3-none-any.whl

birder_clip/__init__.py

@@ -0,0 +1,7 @@
+from birder_clip.model_registry.model_registry import list_pretrained_models
+from birder_clip.version import __version__
+
+__all__ = [
+    "list_pretrained_models",
+    "__version__",
+]

birder_clip/common/fs_ops.py

@@ -0,0 +1,111 @@
+import logging
+from pathlib import Path
+from typing import Any
+from typing import NamedTuple
+from typing import Optional
+
+import torch
+from birder.conf import settings
+from birder.data.transforms.classification import RGBType
+
+from birder_clip.common import lib
+from birder_clip.model_registry import registry
+from birder_clip.net.base import BaseNet
+
+logger = logging.getLogger(__name__)
+
+
+class ModelInfo(NamedTuple):
+    signature: dict[str, Any]
+    rgb_stats: RGBType
+    custom_config: Optional[dict[str, Any]] = None
+
+
+def model_path(network_name: str, *, epoch: Optional[int | str] = None) -> Path:
+    if epoch is not None:
+        file_name = f"{network_name}_{epoch}.pt"
+    else:
+        file_name = f"{network_name}.pt"
+
+    return settings.MODELS_DIR.joinpath(file_name)
+
+
+def load_model(
+    device: torch.device,
+    network: str,
+    *,
+    path: Optional[str | Path] = None,
+    config: Optional[dict[str, Any]] = None,
+    tag: Optional[str] = None,
+    image_encoder: Optional[str] = None,
+    text_encoder: Optional[str] = None,
+    embed_dim: Optional[int] = None,
+    tokenizer: Optional[str] = None,
+    image_encoder_config: Optional[dict[str, Any]] = None,
+    text_encoder_config: Optional[dict[str, Any]] = None,
+    epoch: Optional[int] = None,
+    new_size: Optional[tuple[int, int]] = None,
+    inference: bool,
+    dtype: Optional[torch.dtype] = None,
+) -> tuple[BaseNet, ModelInfo]:
+    if path is None:
+        _network_name = lib.get_image_text_network_name(
+            network,
+            tag=tag,
+            image_encoder=image_encoder,
+            text_encoder=text_encoder,
+            embed_dim=embed_dim,
+            tokenizer=tokenizer,
+        )
+        path = model_path(_network_name, epoch=epoch)
+
+    logger.info(f"Loading model from {path} on device {device}...")
+    model_dict: dict[str, Any] = torch.load(path, map_location=device, weights_only=True)
+
+    loaded_config: dict[str, Any] = model_dict.get("config", {})
+    merged_config = {**loaded_config}
+    if tokenizer is not None:
+        tokenizer_name: Optional[str] = tokenizer
+        merged_config["tokenizer"] = tokenizer
+    else:
+        tokenizer_name = loaded_config.get("tokenizer")
+
+    if image_encoder is not None and text_encoder is not None and embed_dim is not None and tokenizer_name is not None:
+        image_size = loaded_config.get("image", {}).get("size")
+        merged_config.update(
+            lib.get_image_text_network_config(
+                image_encoder,
+                text_encoder,
+                embed_dim,
+                tokenizer_name,
+                image_size=image_size,
+                image_config=image_encoder_config,
+                text_config=text_encoder_config,
+            )
+        )
+    if config is not None:
+        merged_config.update(config)
+    if len(merged_config) == 0:
+        merged_config = None  # type: ignore[assignment]
+
+    net = registry.net_factory(network, config=merged_config)
+    net.load_state_dict(model_dict["state"])
+    if new_size is not None:
+        net.adjust_image_size(new_size)
+
+    net.to(device)
+    if dtype is not None:
+        net.to(dtype)
+    if inference is True:
+        for param in net.parameters():
+            param.requires_grad_(False)
+
+        net.eval()
+
+    if len(loaded_config) == 0:
+        custom_config = None
+    else:
+        custom_config = loaded_config
+        logger.debug(f"Model loaded with custom config: {custom_config}")
+
+    return (net, ModelInfo(model_dict["signature"], model_dict["rgb_stats"], custom_config))

birder_clip/common/lib.py

@@ -0,0 +1,65 @@
+from typing import Any
+from typing import Optional
+
+
+def get_size_from_signature(signature: dict[str, Any]) -> tuple[int, int]:
+    return tuple(signature["inputs"][0]["data_shape"][2:4])
+
+
+def get_channels_from_signature(signature: dict[str, Any]) -> int:
+    return signature["inputs"][0]["data_shape"][1]  # type: ignore[no-any-return]
+
+
+def get_image_text_network_name(
+    network: str,
+    tag: Optional[str] = None,
+    image_encoder: Optional[str] = None,
+    text_encoder: Optional[str] = None,
+    embed_dim: Optional[int] = None,
+    tokenizer: Optional[str] = None,
+) -> str:
+    parts = [network]
+    if image_encoder is not None:
+        parts.append(image_encoder)
+    if text_encoder is not None:
+        parts.append(text_encoder)
+    if tokenizer is not None:
+        parts.append(tokenizer)
+    if embed_dim is not None:
+        parts.append(f"d{embed_dim}")
+
+    network_name = "_".join(parts)
+    if tag is not None:
+        network_name = f"{network_name}_{tag}"
+
+    return network_name
+
+
+def get_image_text_network_config(
+    image_encoder: str,
+    text_encoder: str,
+    embed_dim: int,
+    tokenizer: Optional[str],
+    *,
+    image_size: Optional[tuple[int, int]] = None,
+    image_config: Optional[dict[str, Any]] = None,
+    text_config: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    image: dict[str, Any] = {"network": image_encoder}
+    if image_size is not None:
+        image["size"] = image_size
+    if image_config is not None:
+        image["config"] = image_config
+
+    text: dict[str, Any] = {"network": text_encoder}
+    if text_config is not None:
+        text["config"] = text_config
+
+    network_config = {
+        "image": image,
+        "text": text,
+        "embed_dim": embed_dim,
+        "tokenizer": tokenizer,
+    }
+
+    return network_config

birder_clip/conf/settings.py

@@ -0,0 +1,27 @@
+import logging.config
+import os
+from typing import Any
+
+# Logging
+# https://docs.python.org/3/library/logging.config.html
+LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO").upper()
+LOGGING: dict[str, Any] = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "formatters": {
+        "verbose": {
+            "format": "[{asctime}.{msecs:04.0f} {levelname} {filename}:{lineno:<4d}] {message}",
+            "style": "{",
+            "datefmt": "%d/%b/%Y %H:%M:%S",
+        },
+        "simple": {"format": "[{asctime} {levelname}] {message}", "style": "{"},
+    },
+    "handlers": {
+        "console": {"class": "logging.StreamHandler", "level": "DEBUG", "formatter": "verbose"},
+    },
+    "loggers": {
+        "birder_clip": {"handlers": ["console"], "level": LOG_LEVEL, "propagate": False},
+    },
+}
+
+logging.config.dictConfig(LOGGING)

birder_clip/inference/zero_shot.py

@@ -0,0 +1,50 @@
+"""
+Zero-shot text embedding helpers
+
+Zero-shot classification compares image features against one text feature per
+candidate class. When multiple prompt templates are used, this module follows
+the OpenCLIP/OpenAI CLIP convention: encode every class/template prompt,
+normalize prompt embeddings, average them per class and normalize the averaged
+class embedding again.
+"""
+
+from collections.abc import Sequence
+
+import torch
+import torch.nn.functional as F
+
+from birder_clip.net.base import BaseNet
+from birder_clip.tokenizers.base import Tokenizer
+
+
+def render_prompts(class_names: Sequence[str], templates: Sequence[str]) -> list[str]:
+    return [template.format(class_name) for class_name in class_names for template in templates]
+
+
+def build_class_text_embeddings(
+    model: BaseNet,
+    tokenizer: Tokenizer,
+    class_names: Sequence[str],
+    templates: Sequence[str],
+    *,
+    device: torch.device,
+    context_length: int | None = None,
+    batch_size: int | None = None,
+) -> torch.Tensor:
+    num_templates = len(templates)
+    if batch_size is None:
+        batch_size = len(class_names)
+
+    class_text_embeddings = []
+    with torch.inference_mode():
+        for start in range(0, len(class_names), batch_size):
+            batch_class_names = class_names[start : start + batch_size]
+            prompts = render_prompts(batch_class_names, templates)
+            tokens = tokenizer(prompts, context_length=context_length).to(device)
+            class_embeddings = model.encode_text(tokens, normalize=True)
+
+            class_embeddings = class_embeddings.reshape(len(batch_class_names), num_templates, -1).mean(dim=1)
+            class_embeddings = F.normalize(class_embeddings, dim=-1)
+            class_text_embeddings.append(class_embeddings)
+
+    return torch.concat(class_text_embeddings, dim=0)

birder_clip/inference/zero_shot_templates.py

@@ -0,0 +1,229 @@
+"""
+Zero-shot classification prompt templates
+
+OpenAI simple ImageNet templates adapted from OpenCLIP:
+https://github.com/mlfoundations/open_clip/blob/main/src/open_clip/zero_shot_metadata.py
+
+OpenAI templates adapted from:
+https://github.com/openai/CLIP/blob/main/data/prompts.md
+"""
+
+# Reference license: MIT
+
+DEFAULT_TEMPLATES = ("a photo of a {}.",)
+
+IDENTITY_TEMPLATES = ("{}",)
+
+SIMPLE_IMAGENET_TEMPLATES = (
+    "itap of a {}.",
+    "a bad photo of the {}.",
+    "a origami {}.",
+    "a photo of the large {}.",
+    "a {} in a video game.",
+    "art of the {}.",
+    "a photo of the small {}.",
+)
+
+OPENAI_IMAGENET_TEMPLATES = (
+    "a bad photo of a {}.",
+    "a photo of many {}.",
+    "a sculpture of a {}.",
+    "a photo of the hard to see {}.",
+    "a low resolution photo of the {}.",
+    "a rendering of a {}.",
+    "graffiti of a {}.",
+    "a bad photo of the {}.",
+    "a cropped photo of the {}.",
+    "a tattoo of a {}.",
+    "the embroidered {}.",
+    "a photo of a hard to see {}.",
+    "a bright photo of a {}.",
+    "a photo of a clean {}.",
+    "a photo of a dirty {}.",
+    "a dark photo of the {}.",
+    "a drawing of a {}.",
+    "a photo of my {}.",
+    "the plastic {}.",
+    "a photo of the cool {}.",
+    "a close-up photo of a {}.",
+    "a black and white photo of the {}.",
+    "a painting of the {}.",
+    "a painting of a {}.",
+    "a pixelated photo of the {}.",
+    "a sculpture of the {}.",
+    "a bright photo of the {}.",
+    "a cropped photo of a {}.",
+    "a plastic {}.",
+    "a photo of the dirty {}.",
+    "a jpeg corrupted photo of a {}.",
+    "a blurry photo of the {}.",
+    "a photo of the {}.",
+    "a good photo of the {}.",
+    "a rendering of the {}.",
+    "a {} in a video game.",
+    "a photo of one {}.",
+    "a doodle of a {}.",
+    "a close-up photo of the {}.",
+    "a photo of a {}.",
+    "the origami {}.",
+    "the {} in a video game.",
+    "a sketch of a {}.",
+    "a doodle of the {}.",
+    "a origami {}.",
+    "a low resolution photo of a {}.",
+    "the toy {}.",
+    "a rendition of the {}.",
+    "a photo of the clean {}.",
+    "a photo of a large {}.",
+    "a rendition of a {}.",
+    "a photo of a nice {}.",
+    "a photo of a weird {}.",
+    "a blurry photo of a {}.",
+    "a cartoon {}.",
+    "art of a {}.",
+    "a sketch of the {}.",
+    "a embroidered {}.",
+    "a pixelated photo of a {}.",
+    "itap of the {}.",
+    "a jpeg corrupted photo of the {}.",
+    "a good photo of a {}.",
+    "a plushie {}.",
+    "a photo of the nice {}.",
+    "a photo of the small {}.",
+    "a photo of the weird {}.",
+    "the cartoon {}.",
+    "art of the {}.",
+    "a drawing of the {}.",
+    "a photo of the large {}.",
+    "a black and white photo of a {}.",
+    "the plushie {}.",
+    "a dark photo of a {}.",
+    "itap of a {}.",
+    "graffiti of the {}.",
+    "a toy {}.",
+    "itap of my {}.",
+    "a photo of a cool {}.",
+    "a photo of a small {}.",
+    "a tattoo of the {}.",
+)
+
+OPENAI_CIFAR10_TEMPLATES = (
+    "a photo of a {}.",
+    "a blurry photo of a {}.",
+    "a black and white photo of a {}.",
+    "a low contrast photo of a {}.",
+    "a high contrast photo of a {}.",
+    "a bad photo of a {}.",
+    "a good photo of a {}.",
+    "a photo of a small {}.",
+    "a photo of a big {}.",
+    "a photo of the {}.",
+    "a blurry photo of the {}.",
+    "a black and white photo of the {}.",
+    "a low contrast photo of the {}.",
+    "a high contrast photo of the {}.",
+    "a bad photo of the {}.",
+    "a good photo of the {}.",
+    "a photo of the small {}.",
+    "a photo of the big {}.",
+)
+
+OPENAI_CALTECH101_TEMPLATES = (
+    "a photo of a {}.",
+    "a painting of a {}.",
+    "a plastic {}.",
+    "a sculpture of a {}.",
+    "a sketch of a {}.",
+    "a tattoo of a {}.",
+    "a toy {}.",
+    "a rendition of a {}.",
+    "a embroidered {}.",
+    "a cartoon {}.",
+    "a {} in a video game.",
+    "a plushie {}.",
+    "a origami {}.",
+    "art of a {}.",
+    "graffiti of a {}.",
+    "a drawing of a {}.",
+    "a doodle of a {}.",
+    "a photo of the {}.",
+    "a painting of the {}.",
+    "the plastic {}.",
+    "a sculpture of the {}.",
+    "a sketch of the {}.",
+    "a tattoo of the {}.",
+    "the toy {}.",
+    "a rendition of the {}.",
+    "the embroidered {}.",
+    "the cartoon {}.",
+    "the {} in a video game.",
+    "the plushie {}.",
+    "the origami {}.",
+    "art of the {}.",
+    "graffiti of the {}.",
+    "a drawing of the {}.",
+    "a doodle of the {}.",
+)
+
+OPENAI_COUNTRY211_TEMPLATES = (
+    "a photo i took in {}.",
+    "a photo i took while visiting {}.",
+    "a photo from my home country of {}.",
+    "a photo from my visit to {}.",
+    "a photo showing the country of {}.",
+)
+
+OPENAI_EUROSAT_TEMPLATES = (
+    "a centered satellite photo of {}.",
+    "a centered satellite photo of a {}.",
+    "a centered satellite photo of the {}.",
+)
+
+OPENAI_FGVC_AIRCRAFT_TEMPLATES = (
+    "a photo of a {}, a type of aircraft.",
+    "a photo of the {}, a type of aircraft.",
+)
+
+OPENAI_KINETICS_TEMPLATES = (
+    "a photo of {}.",
+    "a photo of a person {}.",
+    "a photo of a person using {}.",
+    "a photo of a person doing {}.",
+    "a photo of a person during {}.",
+    "a photo of a person performing {}.",
+    "a photo of a person practicing {}.",
+    "a video of {}.",
+    "a video of a person {}.",
+    "a video of a person using {}.",
+    "a video of a person doing {}.",
+    "a video of a person during {}.",
+    "a video of a person performing {}.",
+    "a video of a person practicing {}.",
+    "a example of {}.",
+    "a example of a person {}.",
+    "a example of a person using {}.",
+    "a example of a person doing {}.",
+    "a example of a person during {}.",
+    "a example of a person performing {}.",
+    "a example of a person practicing {}.",
+    "a demonstration of {}.",
+    "a demonstration of a person {}.",
+    "a demonstration of a person using {}.",
+    "a demonstration of a person doing {}.",
+    "a demonstration of a person during {}.",
+    "a demonstration of a person performing {}.",
+    "a demonstration of a person practicing {}.",
+)
+
+TEMPLATE_SETS = {
+    "default": DEFAULT_TEMPLATES,
+    "identity": IDENTITY_TEMPLATES,
+    "simple_imagenet": SIMPLE_IMAGENET_TEMPLATES,
+    "openai_imagenet": OPENAI_IMAGENET_TEMPLATES,
+    "openai_cifar10": OPENAI_CIFAR10_TEMPLATES,
+    "openai_caltech101": OPENAI_CALTECH101_TEMPLATES,
+    "openai_country211": OPENAI_COUNTRY211_TEMPLATES,
+    "openai_eurosat": OPENAI_EUROSAT_TEMPLATES,
+    "openai_fgvc_aircraft": OPENAI_FGVC_AIRCRAFT_TEMPLATES,
+    "openai_kinetics": OPENAI_KINETICS_TEMPLATES,
+}

birder_clip/model_registry/__init__.py

@@ -0,0 +1,9 @@
+from birder_clip.model_registry.model_registry import Task
+from birder_clip.model_registry.model_registry import list_pretrained_models
+from birder_clip.model_registry.model_registry import registry
+
+__all__ = [
+    "Task",
+    "list_pretrained_models",
+    "registry",
+]

birder_clip/model_registry/model_registry.py

@@ -0,0 +1,133 @@
+import fnmatch
+import warnings
+from enum import StrEnum
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Literal
+from typing import Optional
+from typing import TypeAlias
+
+from birder.model_registry.model_registry import group_sort
+
+if TYPE_CHECKING is True:
+    from birder_clip.net.base import BaseNet  # pylint: disable=cyclic-import
+    from birder_clip.net.text.base import TextBaseNet  # pylint: disable=cyclic-import
+
+    NetType: TypeAlias = type[BaseNet] | type[TextBaseNet]
+
+
+class Task(StrEnum):
+    IMAGE_TEXT = "image_text"
+    TEXT = "text"
+
+
+class ModelRegistry:
+    def __init__(self) -> None:
+        self.registered_configs: dict[str, "NetType"] = {}
+        self._image_text_nets: dict[str, type["BaseNet"]] = {}
+        self._text_nets: dict[str, type[TextBaseNet]] = {}
+        self._pretrained_nets: dict[str, dict[str, Any]] = {}
+
+    @property
+    def all_nets(self) -> dict[str, "NetType"]:
+        return {**self._image_text_nets, **self._text_nets}
+
+    def _get_models_for_task(self, task: Task) -> dict[str, "NetType"]:
+        if task == Task.IMAGE_TEXT:
+            nets: dict[str, "NetType"] = self._image_text_nets
+        elif task == Task.TEXT:
+            nets = self._text_nets
+        else:
+            raise ValueError(f"Unsupported model task: {task}")
+
+        return nets
+
+    def _register_model(self, name: str, net_type: "NetType") -> None:
+        name_key = name.lower()
+        task = Task(net_type.task)
+        nets = self._get_models_for_task(task)
+        if name_key in self.all_nets and name_key not in nets:
+            raise ValueError(f"Registered model name '{name}' collides with an existing registered model name")
+        if name_key in nets:
+            warnings.warn(f"Network '{name}' is already registered and will be overwritten", UserWarning)
+
+        nets[name_key] = net_type
+
+    def register_model_config(self, name: str, net_type: "NetType", *, config: Optional[dict[str, Any]] = None) -> None:
+        name_key = name.lower()
+        registered_net_type = type(name, (net_type,), {"config": config})
+        self._register_model(name_key, registered_net_type)
+        if name_key in self.registered_configs:
+            warnings.warn(f"Registered config '{name}' is already registered and will be overwritten", UserWarning)
+
+        self.registered_configs[name_key] = registered_net_type
+
+    def register_weights(self, name: str, weights_info: dict[str, Any]) -> None:
+        if name in self._pretrained_nets:
+            warnings.warn(f"Weights '{name}' are already registered and will be overwritten", UserWarning)
+
+        self._pretrained_nets[name] = weights_info
+
+    def list_models(
+        self,
+        include_filter: Optional[str] = None,
+        *,
+        task: Optional[Task] = None,
+        net_type: Optional[type | tuple[type, ...]] = None,
+        net_type_op: Literal["AND", "OR"] = "AND",
+    ) -> list[str]:
+        nets = self.all_nets
+        if task is not None:
+            nets = self._get_models_for_task(task)
+
+        if net_type is not None:
+            if not isinstance(net_type, tuple):
+                net_type = (net_type,)
+
+            if net_type_op == "OR":
+                nets = {name: t for name, t in nets.items() if issubclass(t, net_type) is True}
+            elif net_type_op == "AND":
+                nets = {name: t for name, t in nets.items() if all(issubclass(t, nt) for nt in net_type)}
+            else:
+                raise ValueError(f"Unknown op {net_type_op}")
+
+        model_list = list(nets.keys())
+        if include_filter is not None:
+            model_list = fnmatch.filter(model_list, include_filter)
+
+        return group_sort(model_list)
+
+    def list_pretrained_models(self, include_filter: Optional[str] = None) -> list[str]:
+        model_list = list(self._pretrained_nets.keys())
+        if include_filter is not None:
+            model_list = fnmatch.filter(model_list, include_filter)
+
+        return group_sort(model_list)
+
+    def exists(self, name: str, task: Optional[Task] = None, net_type: Optional[type] = None) -> bool:
+        nets = self.all_nets
+        if task is not None:
+            nets = self._get_models_for_task(task)
+
+        if net_type is not None:
+            nets = {name: t for name, t in nets.items() if issubclass(t, net_type) is True}
+
+        return name.lower() in nets
+
+    def pretrained_exists(self, name: str) -> bool:
+        return name in self._pretrained_nets
+
+    def get_pretrained_metadata(self, name: str) -> dict[str, Any]:
+        return self._pretrained_nets[name]
+
+    def text_factory(self, name: str, *, config: Optional[dict[str, Any]] = None) -> "TextBaseNet":
+        name_key = name.lower()
+        return self._text_nets[name_key](config=config)
+
+    def net_factory(self, name: str, *, config: Optional[dict[str, Any]] = None) -> "BaseNet":
+        name_key = name.lower()
+        return self._image_text_nets[name_key](config=config)
+
+
+registry = ModelRegistry()
+list_pretrained_models = registry.list_pretrained_models

birder_clip/net/__init__.py

@@ -0,0 +1,5 @@
+from birder_clip.net.clip import CLIP
+
+__all__ = [
+    "CLIP",
+]

birder_clip/net/base.py

@@ -0,0 +1,44 @@
+import copy
+from typing import Any
+from typing import Optional
+
+import torch
+from birder.net.base import BaseNet as ImageEncoder
+from torch import nn
+
+from birder_clip.model_registry import Task
+
+
+class BaseNet(nn.Module):
+    task = str(Task.IMAGE_TEXT)
+
+    def __init__(self, *, config: Optional[dict[str, Any]] = None) -> None:
+        super().__init__()
+        if hasattr(self, "config") is False:
+            self.config = config
+        else:
+            if self.config is not None:
+                self.config = copy.deepcopy(self.config)  # Avoid mutating registered configs
+
+        if config is not None:
+            assert self.config is not None
+            self.config.update(config)  # Override with custom config
+
+        self.image_encoder: ImageEncoder
+        self.embedding_size: int
+        self.tokenizer_name: str
+
+    def encode_image(self, image: torch.Tensor, normalize: bool = False) -> torch.Tensor:
+        raise NotImplementedError
+
+    def encode_text(self, text: torch.Tensor, normalize: bool = False) -> torch.Tensor:
+        raise NotImplementedError
+
+    def forward_logits(self, image_features: torch.Tensor, text_features: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+    def forward(self, image: torch.Tensor, text: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+    def adjust_image_size(self, new_size: tuple[int, int]) -> None:
+        self.image_encoder.adjust_size(new_size)