|
|
from typing import * |
|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
from torchvision import transforms |
|
|
from transformers import DINOv3ViTModel |
|
|
import numpy as np |
|
|
from PIL import Image |
|
|
|
|
|
|
|
|
class DinoV2FeatureExtractor: |
|
|
""" |
|
|
Feature extractor for DINOv2 models. |
|
|
""" |
|
|
def __init__(self, model_name: str): |
|
|
self.model_name = model_name |
|
|
self.model = torch.hub.load('facebookresearch/dinov2', model_name, pretrained=True) |
|
|
self.model.eval() |
|
|
self.transform = transforms.Compose([ |
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), |
|
|
]) |
|
|
|
|
|
def to(self, device): |
|
|
self.model.to(device) |
|
|
|
|
|
def cuda(self): |
|
|
self.model.cuda() |
|
|
|
|
|
def cpu(self): |
|
|
self.model.cpu() |
|
|
|
|
|
@torch.no_grad() |
|
|
def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor: |
|
|
""" |
|
|
Extract features from the image. |
|
|
|
|
|
Args: |
|
|
image: A batch of images as a tensor of shape (B, C, H, W) or a list of PIL images. |
|
|
|
|
|
Returns: |
|
|
A tensor of shape (B, N, D) where N is the number of patches and D is the feature dimension. |
|
|
""" |
|
|
if isinstance(image, torch.Tensor): |
|
|
assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)" |
|
|
elif isinstance(image, list): |
|
|
assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images" |
|
|
image = [i.resize((518, 518), Image.LANCZOS) for i in image] |
|
|
image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] |
|
|
image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] |
|
|
image = torch.stack(image).cuda() |
|
|
else: |
|
|
raise ValueError(f"Unsupported type of image: {type(image)}") |
|
|
|
|
|
image = self.transform(image).cuda() |
|
|
features = self.model(image, is_training=True)['x_prenorm'] |
|
|
patchtokens = F.layer_norm(features, features.shape[-1:]) |
|
|
return patchtokens |
|
|
|
|
|
|
|
|
class DinoV3FeatureExtractor: |
|
|
""" |
|
|
Feature extractor for DINOv3 models. |
|
|
""" |
|
|
def __init__(self, model_name: str, image_size=512): |
|
|
self.model_name = model_name |
|
|
self.model = DINOv3ViTModel.from_pretrained(model_name) |
|
|
self.model.eval() |
|
|
self.image_size = image_size |
|
|
self.transform = transforms.Compose([ |
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), |
|
|
]) |
|
|
|
|
|
def to(self, device): |
|
|
self.model.to(device) |
|
|
|
|
|
def cuda(self): |
|
|
self.model.cuda() |
|
|
|
|
|
def cpu(self): |
|
|
self.model.cpu() |
|
|
|
|
|
def extract_features(self, image: torch.Tensor) -> torch.Tensor: |
|
|
image = image.to(self.model.embeddings.patch_embeddings.weight.dtype) |
|
|
hidden_states = self.model.embeddings(image, bool_masked_pos=None) |
|
|
position_embeddings = self.model.rope_embeddings(image) |
|
|
|
|
|
for i, layer_module in enumerate(self.model.layer): |
|
|
hidden_states = layer_module( |
|
|
hidden_states, |
|
|
position_embeddings=position_embeddings, |
|
|
) |
|
|
|
|
|
return F.layer_norm(hidden_states, hidden_states.shape[-1:]) |
|
|
|
|
|
@torch.no_grad() |
|
|
def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor: |
|
|
""" |
|
|
Extract features from the image. |
|
|
|
|
|
Args: |
|
|
image: A batch of images as a tensor of shape (B, C, H, W) or a list of PIL images. |
|
|
|
|
|
Returns: |
|
|
A tensor of shape (B, N, D) where N is the number of patches and D is the feature dimension. |
|
|
""" |
|
|
if isinstance(image, torch.Tensor): |
|
|
assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)" |
|
|
elif isinstance(image, list): |
|
|
assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images" |
|
|
image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image] |
|
|
image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] |
|
|
image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] |
|
|
image = torch.stack(image).cuda() |
|
|
else: |
|
|
raise ValueError(f"Unsupported type of image: {type(image)}") |
|
|
|
|
|
image = self.transform(image).cuda() |
|
|
features = self.extract_features(image) |
|
|
return features |
|
|
|
