basketball_analysis/matcherBeta.py

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())