Spaces:
Paused
Paused
File size: 2,188 Bytes
6642f4e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import torch
from torch.utils.data import DataLoader
from starvector.metrics.base_metric import BaseMetric
from tqdm import tqdm
from transformers import AutoModel, AutoImageProcessor
from PIL import Image
import torch.nn as nn
class DINOScoreCalculator(BaseMetric):
def __init__(self, config=None, device='cuda'):
super().__init__()
self.class_name = self.__class__.__name__
self.config = config
self.model, self.processor = self.get_DINOv2_model("base")
self.model = self.model.to(device)
self.device = device
self.metric = self.calculate_DINOv2_similarity_score
def get_DINOv2_model(self, model_size):
if model_size == "small":
model_size = "facebook/dinov2-small"
elif model_size == "base":
model_size = "facebook/dinov2-base"
elif model_size == "large":
model_size = "facebook/dinov2-large"
else:
raise ValueError(f"model_size should be either 'small', 'base' or 'large', got {model_size}")
return AutoModel.from_pretrained(model_size), AutoImageProcessor.from_pretrained(model_size)
def process_input(self, image, processor):
if isinstance(image, str):
image = Image.open(image)
if isinstance(image, Image.Image):
with torch.no_grad():
inputs = processor(images=image, return_tensors="pt").to(self.device)
outputs = self.model(**inputs)
features = outputs.last_hidden_state.mean(dim=1)
elif isinstance(image, torch.Tensor):
features = image.unsqueeze(0) if image.dim() == 1 else image
else:
raise ValueError("Input must be a file path, PIL Image, or tensor of features")
return features
def calculate_DINOv2_similarity_score(self, **kwargs):
image1 = kwargs.get('gt_im')
image2 = kwargs.get('gen_im')
features1 = self.process_input(image1, self.processor)
features2 = self.process_input(image2, self.processor)
cos = nn.CosineSimilarity(dim=1)
sim = cos(features1, features2).item()
sim = (sim + 1) / 2
return sim
|