|
import os
|
|
import sys
|
|
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
|
|
import time
|
|
from abc import ABC, abstractmethod
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from utils import configs
|
|
from utils.functional import check_data_type_variable, get_device, image_augmentations
|
|
|
|
|
|
class BaseModelImageEmbeddings(ABC):
|
|
def __init__(
|
|
self,
|
|
name_model: str,
|
|
freeze_model: bool,
|
|
pretrained_model: bool,
|
|
support_set_method: str,
|
|
):
|
|
self.name_model = name_model
|
|
self.freeze_model = freeze_model
|
|
self.pretrained_model = pretrained_model
|
|
self.support_set_method = support_set_method
|
|
self.model = None
|
|
self.device = get_device()
|
|
|
|
self.check_arguments()
|
|
|
|
def check_arguments(self):
|
|
check_data_type_variable(self.name_model, str)
|
|
check_data_type_variable(self.freeze_model, bool)
|
|
check_data_type_variable(self.pretrained_model, bool)
|
|
check_data_type_variable(self.support_set_method, str)
|
|
|
|
old_name_model = self.name_model
|
|
if self.name_model == configs.CLIP_NAME_MODEL:
|
|
old_name_model = self.name_model
|
|
self.name_model = "clip"
|
|
if self.name_model not in tuple(configs.NAME_MODELS.keys()):
|
|
raise ValueError(f"Model {self.name_model} not supported")
|
|
if self.support_set_method not in configs.SUPPORT_SET_METHODS:
|
|
raise ValueError(
|
|
f"Support set method {self.support_set_method} not supported"
|
|
)
|
|
self.name_model = old_name_model
|
|
|
|
@abstractmethod
|
|
def init_model(self):
|
|
pass
|
|
|
|
def get_embeddings(self, image: np.ndarray) -> dict:
|
|
image_input = image_augmentations()(image=image)["image"]
|
|
image_input = image_input.unsqueeze(axis=0).to(self.device)
|
|
with torch.no_grad():
|
|
start_time = time.perf_counter()
|
|
embeddings = self.model(image_input)
|
|
end_time = time.perf_counter() - start_time
|
|
|
|
embeddings = embeddings.detach().cpu().numpy()
|
|
return {
|
|
"embeddings": embeddings,
|
|
"inference_time": end_time,
|
|
}
|
|
|
