| import os | |
| import sys | |
| sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) | |
| import time | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from transformers import CLIPModel as CLIPTransformersModel | |
| from transformers import CLIPProcessor | |
| from models.base_model import BaseModelMainModel | |
| class CLIPModel(BaseModelMainModel): | |
| def __init__( | |
| self, | |
| name_model: str, | |
| freeze_model: bool, | |
| pretrained_model: bool, | |
| support_set_method: str, | |
| ): | |
| super().__init__(name_model, freeze_model, pretrained_model, support_set_method) | |
| self.init_model() | |
| def init_model(self): | |
| self.model = CLIPTransformersModel.from_pretrained(self.name_model) | |
| for layer in self.model.children(): | |
| if hasattr(layer, "reset_parameters") and not self.pretrained_model: | |
| layer.reset_parameters() | |
| for param in self.model.parameters(): | |
| param.required_grad = False if not self.freeze_model else True | |
| self.model.to(self.device) | |
| self.model.eval() | |
| self.processor = CLIPProcessor.from_pretrained(self.name_model) | |
| def predict(self, image: np.ndarray, list_class: tuple) -> dict: | |
| image = Image.fromarray(image) | |
| with torch.no_grad(): | |
| inputs = self.processor( | |
| text=list_class, images=image, return_tensors="pt", padding=True | |
| ) | |
| start_time = time.perf_counter() | |
| outputs = self.model(**inputs) | |
| end_time = time.perf_counter() - start_time | |
| logits_per_image = outputs.logits_per_image | |
| probs = logits_per_image.softmax(dim=1) | |
| argmax_probs = probs.argmax(dim=1) | |
| result = list_class[argmax_probs[0]] | |
| return { | |
| "class": result, | |
| "confidence": float(probs[0, argmax_probs[0]]), | |
| "inference_time": end_time, | |
| } | |