import torch from torch import nn, Tensor import torch.nn.functional as F from transformers import Dinov2Model, Dinov2Config from torchvision.transforms import v2 from code import interact import json import os from PIL import Image import numpy as np from typing import Union transforms = v2.Compose([ v2.ToImage(), v2.ToDtype(torch.float32, scale=True), v2.Resize((224, 224)), v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) class CrossAttention(nn.Module): def __init__(self, d_model:int, *args, **kwargs): super().__init__(*args, **kwargs) self.Wq = nn.Linear(d_model, d_model) self.Wk = nn.Linear(d_model, d_model) self.Wv = nn.Linear(d_model, d_model) def forward(self, queries, candidates): Q = self.Wk(candidates) # (B, num_candidates, d_model) K = self.Wq(queries) # (B, num_queries, d_model) V = self.Wv(queries) # (B, num_queries, d_model) attn_out = F.scaled_dot_product_attention(Q, K, V) # (B, num_candidates, d_model) return attn_out class JointTransformer(nn.Module): def __init__( self, d_model=384, nhead=4, num_layers=4, *args, **kwargs ): super().__init__(*args, **kwargs) # Transformer encoder encoder_layer = nn.TransformerEncoderLayer( d_model=d_model, nhead=nhead, dim_feedforward=4 * d_model, batch_first=True, dropout=0.0 ) self.transformer = nn.TransformerEncoder(encoder_layer, num_layers) def forward(self, query: Tensor, candidates: Tensor) -> Tensor : Q = query.size(1) assert Q == 1 x = torch.cat((query, candidates), dim=1) # (B, Q+C, D) x = self.transformer(x) # (B, Q+C, D) query = x[:,:Q,:] # (B, Q, D) candidates = x[:, Q:, :] # (B, C, D) return query, candidates class MLP(nn.Module): def __init__(self, emb_dim, expand_factor, *args, **kwargs): super().__init__(*args, **kwargs) self.lin1 = nn.Linear(emb_dim, emb_dim*expand_factor) self.gelu = nn.GELU("tanh") self.lin2 = nn.Linear(emb_dim*expand_factor, emb_dim) def forward(self, x:Tensor) -> Tensor: x = self.lin1(x) x = self.gelu(x) x = self.lin2(x) return x class Matcher(nn.Module): def __init__(self, max_candidates, num_layers, dino_dir, *args, **kwargs): super().__init__(*args, **kwargs) # -------------- Pre-trained Encoder (frozen) ----------------- assert isinstance(dino_dir, str) with open(os.path.join(dino_dir, "config.json"), "r") as f: dino_cfg = json.load(f) self.encoder = Dinov2Model.from_pretrained(dino_dir, config = Dinov2Config(**dino_cfg)) self.freeze_encoder() # ----------------- Embeddings to distinguish queries and candidates --------------------- self.query_image_embed = nn.Parameter(torch.randn(1, 1, dino_cfg["hidden_size"])) self.candidates_image_embed = nn.Embedding(max_candidates, dino_cfg["hidden_size"]) self.null_candidate = nn.Parameter(torch.randn(1, 1, dino_cfg["hidden_size"])) # null candidate embedding # ---------------- Joint transformer (trained) ---------------------- self.max_candidates = max_candidates self.num_layers = num_layers self.joint_transformer = JointTransformer( d_model = dino_cfg["hidden_size"], nhead = dino_cfg["num_attention_heads"], num_layers = num_layers, ) self.lnormq = nn.LayerNorm(dino_cfg["hidden_size"], ) self.lnormc = nn.LayerNorm(dino_cfg["hidden_size"], ) # ------------------------ Final operation --------------------------- self.cross_attn = CrossAttention(dino_cfg["hidden_size"]) self.lnormc2 = nn.LayerNorm(dino_cfg["hidden_size"]) self.classification_layer = nn.Linear(dino_cfg["hidden_size"], 1) def freeze_encoder(self) -> None: for p in self.encoder.parameters(): p.requires_grad_(False) def pre_process_img(self, image:Union[Image.Image, np.ndarray, str]): if isinstance(image, str): image = Image.open(image) return transforms(image) @torch.inference_mode() def predict(self, query_crop: np.ndarray, candidate_crops: list[np.ndarray]): query = transforms(query_crop)[None, None, ...] candidates = torch.stack([transforms(candidate_crop) for candidate_crop in candidate_crops]).unsqueeze(0) probs = self.forward(query, candidates).softmax(dim=-1) return probs.numpy() def forward(self, query: Tensor, candidates: Tensor) -> Tensor : # query (B,1,3,H,W), candidates (B,C,3,H,W) B, C, _, H, W = candidates.shape query = self.encoder( query.view(B, 3, H, W) )['last_hidden_state'] # (B, T, D) # pick the CLS_TOKEN query = query[:,0,:].view(B, 1, -1) # (B, 1, D) candidates = self.encoder( candidates.view(B*C, 3, H, W) )['last_hidden_state'] # (B*C, T, D) # pick the CLS_TOKEN candidates = candidates[:,0,:].view(B, C, -1) # (B, C, D) # Add embeddings query = query + self.query_image_embed.repeat(B, 1, 1) # (B, 1, D) candidate_ids = torch.arange(C, device=query.device).view(1, C) candidates = candidates + self.candidates_image_embed(candidate_ids) # (B, C, D) candidates = torch.cat( ( candidates, self.null_candidate.repeat(B, 1, 1) ), dim=1) # (B, C+1, D) # Joint transformer, candidate and query tokens attend to each other q, c = self.joint_transformer(query, candidates) # skip connections query = self.lnormq(query + q) candidates = self.lnormc(candidates + c) # Cross attention, query attends to candidates c = self.cross_attn(query, candidates) # (B, C+1, D) candidates = self.lnormc2(candidates + c) candidates = candidates + c logits = self.classification_layer(candidates) # (B, C+1, 1) return logits.squeeze(-1) if __name__ == "__main__": import random B, H, W = 1, 224, 224 max_candidates = 10 num_layers = 4 query = torch.randn((B, 1, 3, H, W)) candidates = torch.randn((B, random.randint(2, max_candidates), 3, H, W)) matcher = Matcher(max_candidates, num_layers, "DINOv2_base") out = matcher(query, candidates) interact(local=locals())