| import os | |
| import sys | |
| sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) | |
| import numpy as np | |
| import torch.nn as nn | |
| from PIL import Image | |
| from models.base_model import BaseModelImageEmbeddings | |
| from utils import configs | |
| from .lightning_module import ImageClassificationLightningModule | |
| class ImageEmbeddings(BaseModelImageEmbeddings): | |
| 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 = ImageClassificationLightningModule.load_from_checkpoint( | |
| os.path.join( | |
| configs.WEIGHTS_PATH, | |
| self.name_model, | |
| self.support_set_method, | |
| "best.ckpt", | |
| ), | |
| name_model=self.name_model, | |
| freeze_model=self.freeze_model, | |
| pretrained_model=self.pretrained_model, | |
| ) | |
| self.model = self.model.model | |
| self.model.classifier = nn.Identity() | |
| 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() | |
| if __name__ == "__main__": | |
| model = ImageEmbeddings("resnet50", True, True, "1_shot") | |
| image = np.array( | |
| Image.open( | |
| "../../assets/example_images/gon/306e5d35-b301-4299-8022-0c89dc0b7690.png" | |
| ).convert("RGB") | |
| ) | |
| result = model.get_embeddings(image) | |
| print(result) | |