Spaces:
Runtime error
Runtime error
| import numpy as np | |
| from PIL import Image | |
| import torch | |
| import torch.nn.functional as F | |
| from typing import List | |
| from torchvision.models import vit_b_16 | |
| import torchvision.transforms as transforms | |
| ROUND_DIGIT=3 | |
| NUM_ASPECT=5 | |
| DINO_POINT_HIGH=0.97 | |
| DINO_POINT_MID=0.9 | |
| DINO_POINT_LOW=0.8 | |
| class MetricDINO_sim(): | |
| def __init__(self, device="cuda") -> None: | |
| """ | |
| Initialize a class MetricDINO_sim with the specified device for testing temporal consistency of a given video. | |
| Args: | |
| device (str, optional): The device on which the model will run. Defaults to "cuda". | |
| """ | |
| self.device = device | |
| self.model = vit_b_16(pretrained=True) | |
| self.model.to(self.device).eval() | |
| self.preprocess = transforms.Compose([ | |
| transforms.Resize(256), | |
| transforms.CenterCrop(224), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| ]) | |
| def evaluate(self, frame_list:List[Image.Image]): | |
| """ | |
| Calculate the cosine similarity between the DINO features of adjacent frames of a given video to test temporal consistency, | |
| then quantize the orginal output based on some predefined thresholds. | |
| Args: | |
| frame_list:List[Image.Image], frames of the video used in calculation. | |
| Returns: | |
| dino_frame_score: float, the computed DINO feature cosine similarity between each adjacent pair of frames and then averaged among all the pairs. | |
| quantized_ans: int, the quantized value of the above avg DINO-Sim scores based on pre-defined thresholds. | |
| """ | |
| device = self.device | |
| frame_sim_list=[] | |
| for f_idx in range(len(frame_list)-1): | |
| frame_1=frame_list[f_idx] | |
| frame_2=frame_list[f_idx+1] | |
| frame_tensor_1 = self.preprocess(frame_1).unsqueeze(0).to(device) | |
| frame_tensor_2 = self.preprocess(frame_2).unsqueeze(0).to(device) | |
| with torch.no_grad(): | |
| feat_1 = self.model(frame_tensor_1).flatten() | |
| feat_2 = self.model(frame_tensor_2).flatten() | |
| cos_sim=F.cosine_similarity(feat_1, feat_2, dim=0).item() | |
| frame_sim_list.append(cos_sim) | |
| dino_frame_score = np.mean(frame_sim_list) | |
| quantized_ans=0 | |
| if dino_frame_score >= DINO_POINT_HIGH: | |
| quantized_ans=4 | |
| elif dino_frame_score < DINO_POINT_HIGH and dino_frame_score >= DINO_POINT_MID: | |
| quantized_ans=3 | |
| elif dino_frame_score < DINO_POINT_MID and dino_frame_score >= DINO_POINT_LOW: | |
| quantized_ans=2 | |
| else: | |
| quantized_ans=1 | |
| return dino_frame_score, quantized_ans | |