| import warnings |
| warnings.filterwarnings("ignore") |
|
|
| import torch |
| import numpy as np |
| from PIL import Image |
|
|
|
|
| MODELS = { |
| "CLIP ViT-B/32": "clip", |
| "SigLIP 2 Base": "siglip2", |
| "X-CLIP Base": "xclip", |
| } |
|
|
|
|
| class ZeroShotVideoClassifier: |
| def __init__(self, model_key: str = "CLIP ViT-B/32", device: str = None): |
| self.model_key = model_key |
| self.backend = MODELS[model_key] |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| self._load_model() |
|
|
| def _load_model(self): |
| if self.backend == "clip": |
| import open_clip |
| self.model, _, self.preprocess = open_clip.create_model_and_transforms( |
| "ViT-B-32", pretrained="openai" |
| ) |
| self.tokenizer = open_clip.get_tokenizer("ViT-B-32") |
| self.model.to(self.device).eval() |
|
|
| elif self.backend == "siglip2": |
| from transformers import AutoModel, AutoProcessor |
| self.model = AutoModel.from_pretrained("google/siglip2-base-patch16-224").eval() |
| self.processor = AutoProcessor.from_pretrained("google/siglip2-base-patch16-224") |
| self.model.to(self.device) |
|
|
| elif self.backend == "xclip": |
| from transformers import XCLIPModel, XCLIPProcessor |
| self.model = XCLIPModel.from_pretrained("microsoft/xclip-base-patch16").eval() |
| self.processor = XCLIPProcessor.from_pretrained("microsoft/xclip-base-patch16") |
| self.model.to(self.device) |
|
|
| def _encode_text_clip(self, labels: list[str]) -> torch.Tensor: |
| texts = self.tokenizer(["a video of {}".format(l) for l in labels]).to(self.device) |
| with torch.no_grad(): |
| embs = self.model.encode_text(texts) |
| return embs / embs.norm(dim=-1, keepdim=True) |
|
|
| def _encode_frames_clip(self, frames: list[Image.Image]) -> torch.Tensor: |
| tensors = torch.stack([self.preprocess(f) for f in frames]).to(self.device) |
| with torch.no_grad(): |
| embs = self.model.encode_image(tensors) |
| embs = embs / embs.norm(dim=-1, keepdim=True) |
| return embs.mean(dim=0, keepdim=True) |
|
|
| def _encode_text_siglip2(self, labels: list[str]) -> torch.Tensor: |
| inputs = self.processor( |
| text=["a video of {}".format(l) for l in labels], |
| padding="max_length", |
| truncation=True, |
| max_length=64, |
| return_tensors="pt" |
| ) |
| with torch.no_grad(): |
| out = self.model.get_text_features(**inputs) |
| embs = out.pooler_output if hasattr(out, "pooler_output") else out[0][:, 0, :] |
| return embs / embs.norm(dim=-1, keepdim=True) |
|
|
| def _encode_frames_siglip2(self, frames: list[Image.Image]) -> torch.Tensor: |
| inputs = self.processor(images=frames, return_tensors="pt") |
| with torch.no_grad(): |
| out = self.model.get_image_features(**inputs) |
| embs = out.pooler_output if hasattr(out, "pooler_output") else out[0][:, 0, :] |
| embs = embs / embs.norm(dim=-1, keepdim=True) |
| video_emb = embs.mean(dim=0, keepdim=True) |
| return video_emb / video_emb.norm(dim=-1, keepdim=True) |
|
|
| def _encode_text_xclip(self, labels: list[str]) -> torch.Tensor: |
| inputs = self.processor.tokenizer( |
| ["a video of {}".format(l) for l in labels], |
| return_tensors="pt", |
| padding=True, |
| truncation=True |
| ) |
| with torch.no_grad(): |
| out = self.model.get_text_features(**inputs) |
| embs = out.pooler_output |
| return embs / embs.norm(dim=-1, keepdim=True) |
|
|
| def _encode_frames_xclip(self, frames: list[Image.Image]) -> torch.Tensor: |
| frames_np = [np.array(f) for f in frames] |
| inputs = self.processor.image_processor(images=frames_np, return_tensors="pt") |
| with torch.no_grad(): |
| out = self.model.get_video_features(pixel_values=inputs["pixel_values"]) |
| embs = out.pooler_output |
| return embs / embs.norm(dim=-1, keepdim=True) |
|
|
| def classify(self, frames: list[Image.Image], labels: list[str], top_k: int = 5) -> list[dict]: |
| if self.backend == "clip": |
| text_embs = self._encode_text_clip(labels) |
| video_emb = self._encode_frames_clip(frames) |
| scores = (video_emb @ text_embs.T).squeeze().softmax(dim=-1).cpu().numpy() |
|
|
| elif self.backend == "siglip2": |
| text_embs = self._encode_text_siglip2(labels) |
| video_emb = self._encode_frames_siglip2(frames) |
| logit_scale = self.model.logit_scale.exp() |
| logit_bias = self.model.logit_bias |
| logits = video_emb @ text_embs.T * logit_scale + logit_bias |
| scores = torch.sigmoid(logits).detach().squeeze().cpu().numpy() |
|
|
| elif self.backend == "xclip": |
| text_embs = self._encode_text_xclip(labels) |
| video_emb = self._encode_frames_xclip(frames) |
| scores = (video_emb @ text_embs.T).detach().squeeze().cpu().numpy() |
|
|
| top_indices = np.argsort(scores)[::-1][:top_k] |
| return [{"label": labels[i], "score": float(scores[i])} for i in top_indices] |
|
|
