| import torch | |
| import torchvision.transforms as tvf | |
| from ..utils.base_model import BaseModel | |
| class OpenIBL(BaseModel): | |
| default_conf = { | |
| "model_name": "vgg16_netvlad", | |
| } | |
| required_inputs = ["image"] | |
| def _init(self, conf): | |
| self.net = torch.hub.load( | |
| "yxgeee/OpenIBL", conf["model_name"], pretrained=True | |
| ).eval() | |
| mean = [0.48501960784313836, 0.4579568627450961, 0.4076039215686255] | |
| std = [0.00392156862745098, 0.00392156862745098, 0.00392156862745098] | |
| self.norm_rgb = tvf.Normalize(mean=mean, std=std) | |
| def _forward(self, data): | |
| image = self.norm_rgb(data["image"]) | |
| desc = self.net(image) | |
| return { | |
| "global_descriptor": desc, | |
| } | |