from typing import cast
from ...lazy_import import OptionalImport
from ...typing import FloatNumpyArray, ImageInput
from ..backbones import BackboneWithHead
with OptionalImport(package="torch", extra="nn") as _torch_import:
import torch
from torchvision import transforms
_torch_import.check()
[docs]
class BCESiameseNetwork(BackboneWithHead):
"""
Siamese network that classifies image pairs, proposed in
`Koch, G., Zemel, R., & Salakhutdinov, R. (2015). Siamese Neural Networks
for One-shot Image Recognition`.
For more information, see the documentation:
``https://mechacritter.github.io/Python-Visual-Similarity/neural_networks/bce_siamese/bce_siamese.html``.
NOTE
----
The score is a *learned probability*, not a geometric similarity: it is
symmetric in its inputs (the L1 distance is), lives in ``(0, 1)``, and for
two identical images equals ``sigmoid(b)`` -- the learned bias sets the
operating point, so a perfect match does not score exactly ``1``.
References:
===========
[1] Koch, G., Zemel, R., & Salakhutdinov, R. (2015). Siamese Neural Networks
for One-shot Image Recognition. ICML Deep Learning Workshop.
https://www.cs.cmu.edu/~rsalakhu/papers/oneshot1.pdf
:param backbone: name of feature-extraction network.
See
``https://mechacritter.github.io/Python-Visual-Similarity/neural_networks/backbones/backbones.html``.
:param embedding_dim: Dimensionality of the twin feature vectors that the
scoring layer compares.
:param transform: processing transform applied to every input image. If
``None``, the ImageNet preprocessing matching the backbone is
used.
:param device: Device on which the model is placed.
:param pretrained_backbone: Whether to use a backbone pretrained on
ImageNet. If you are loading the ``BCESiameseNetwork`` from a
checkpoint, set this to ``False`` to avoid downloading the weights again.
:param batch_size: Maximum number of images processed in a single batch.
Set to ``-1`` to process all images as a single batch.
:raises ValueError: If ``embedding_dim`` is not a positive integer or if
``backbone`` is not a supported backbone name.
"""
def __init__(
self,
backbone: str = "resnet18",
embedding_dim: int = 128,
transform: transforms.Compose | None = None,
device: str | torch.device = "cpu",
pretrained_backbone: bool = True,
*,
batch_size: int = 16,
):
super().__init__(
backbone=backbone,
embedding_dim=embedding_dim,
transform=transform,
pretrained_backbone=pretrained_backbone,
batch_size=batch_size,
)
self._scorer: torch.nn.Module = torch.nn.Linear(embedding_dim, 1)
self.to(torch.device(device))
def _forward_once(self, x: torch.Tensor) -> torch.Tensor:
"""
Computes sigmoid-activated feature vectors for a batch of images.
Unlike the contrastive variant, the features are *not* L2-normalized;
each component is squashed into ``(0, 1)`` so the component-wise L1
distances fed to the scoring layer are bounded.
:param x: Preprocessed image tensor of shape (batch, channels, H, W).
:return: Feature tensor of shape (batch, embedding_dim) with values
in ``(0, 1)``.
"""
features = self._backbone(x)
return torch.sigmoid(self._head(features))
[docs]
def embed(
self,
images: ImageInput,
*,
dims: str = "HWC",
value_range: tuple[float, float] = (0.0, 255.0),
) -> FloatNumpyArray:
"""Not implemented for this class. Please do not use!"""
raise NotImplementedError(
f"{type(self).__name__} does not learn to generate embeddings. "
"Use the ContrastiveSiameseNetwork for that purpose."
)
[docs]
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
"""
Computes same-class logits for a batch of aligned image pairs.
The i-th logit scores the pair ``(x1[i], x2[i])``; apply
:func:`torch.sigmoid` to obtain probabilities, or feed the logits
directly to :class:`torch.nn.BCEWithLogitsLoss` during training.
:param x1: First preprocessed image batch, shape (batch, channels, H, W).
:param x2: Second preprocessed image batch of the same shape.
:return: Logit tensor of shape (batch,).
:raises ValueError: If the two batches differ in shape.
"""
if x1.shape != x2.shape:
raise ValueError(
f"Input batches must have the same shape, got "
f"{tuple(x1.shape)} vs {tuple(x2.shape)}."
)
features1 = self._forward_once(x1)
features2 = self._forward_once(x2)
return self._score_distances(torch.abs(features1 - features2))
def _score_distances(self, distances: torch.Tensor) -> torch.Tensor:
"""
Maps component-wise L1 distance vectors to same-class logits.
:param distances: Tensor of shape (..., embedding_dim) holding
``|h_1 - h_2|`` for each pair.
:return: Logit tensor of shape (...,).
"""
return cast(torch.Tensor, self._scorer(distances).squeeze(-1))
[docs]
@torch.no_grad()
def similarity_score(
self,
images1: ImageInput,
images2: ImageInput,
*,
dims: str = "HWC",
value_range: tuple[float, float] = (0.0, 255.0),
) -> FloatNumpyArray:
features1 = self._embed_images(images1, dims=dims, value_range=value_range)
features2 = self._embed_images(images2, dims=dims, value_range=value_range)
distances = torch.abs(features1.unsqueeze(1) - features2.unsqueeze(0))
probabilities = torch.sigmoid(self._score_distances(distances))
return cast(FloatNumpyArray, probabilities.cpu().numpy())
@property
def scorer(self) -> torch.nn.Module:
"""The learned layer mapping L1 distances to same-class logits."""
return self._scorer