Spaces:
Sleeping
Sleeping
| import torch | |
| from torchvision import transforms | |
| from torchvision import models | |
| ENCODING_DEVICE = 'cuda' | |
| class Extractor: | |
| """ | |
| Features extractor boilerplate | |
| """ | |
| def __init__(self,): | |
| self.name='' | |
| self.input_shape=None | |
| self.output_shape=None | |
| self.model = None | |
| self.preprocess = None | |
| self.device = None | |
| class ResnetEncoder(Extractor): | |
| def __init__(self): | |
| super().__init__() | |
| pass | |
| """ | |
| Legacy code | |
| """ | |
| # class ResnetEncoder(Extractor): | |
| # def __init__(self,): | |
| # super().__init__() | |
| # self.device = ENCODING_DEVICE | |
| # self.model = models.resnet18(weights='ResNet18_Weights.DEFAULT') | |
| # self.model = torch.nn.Sequential(*(list(self.model.children())[:-1])).to(self.device) #remove classifier | |
| # self.model.eval() | |
| # self.preprocess = transforms.Compose([ | |
| # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| # ]) | |
| # def encode(self,images:torch.Tensor)->torch.Tensor: | |
| # """ | |
| # args: | |
| # image(torch.Tensor): Batch of images with shape (B,3,H,W) with uint8 values | |
| # return: | |
| # output(torch.Tensor): Batch of encoded images (images features) with shape(B,512) | |
| # """ | |
| # assert images.dim()==4 | |
| # assert images.shape[1]==3 | |
| # with torch.inference_mode(): | |
| # output = images.clone().to(self.device) | |
| # output = self.preprocess(output) | |
| # output = self.model(output) | |
| # output = torch.flatten(output,start_dim=-3,end_dim=-1) | |
| # return output | |