Spaces:
Running
Running
| import sys | |
| from pathlib import Path | |
| import torch | |
| from zipfile import ZipFile | |
| import os | |
| import sklearn | |
| import gdown | |
| from ..utils.base_model import BaseModel | |
| sys.path.append(str(Path(__file__).parent / "../../third_party/deep-image-retrieval")) | |
| os.environ["DB_ROOT"] = "" # required by dirtorch | |
| from dirtorch.utils import common # noqa: E402 | |
| from dirtorch.extract_features import load_model # noqa: E402 | |
| # The DIR model checkpoints (pickle files) include sklearn.decomposition.pca, | |
| # which has been deprecated in sklearn v0.24 | |
| # and must be explicitly imported with `from sklearn.decomposition import PCA`. | |
| # This is a hacky workaround to maintain forward compatibility. | |
| sys.modules["sklearn.decomposition.pca"] = sklearn.decomposition._pca | |
| class DIR(BaseModel): | |
| default_conf = { | |
| "model_name": "Resnet-101-AP-GeM", | |
| "whiten_name": "Landmarks_clean", | |
| "whiten_params": { | |
| "whitenp": 0.25, | |
| "whitenv": None, | |
| "whitenm": 1.0, | |
| }, | |
| "pooling": "gem", | |
| "gemp": 3, | |
| } | |
| required_inputs = ["image"] | |
| dir_models = { | |
| "Resnet-101-AP-GeM": "https://docs.google.com/uc?export=download&id=1UWJGDuHtzaQdFhSMojoYVQjmCXhIwVvy", | |
| } | |
| def _init(self, conf): | |
| checkpoint = Path(torch.hub.get_dir(), "dirtorch", conf["model_name"] + ".pt") | |
| if not checkpoint.exists(): | |
| checkpoint.parent.mkdir(exist_ok=True, parents=True) | |
| link = self.dir_models[conf["model_name"]] | |
| gdown.download(str(link), str(checkpoint) + ".zip", quiet=False) | |
| zf = ZipFile(str(checkpoint) + ".zip", "r") | |
| zf.extractall(checkpoint.parent) | |
| zf.close() | |
| os.remove(str(checkpoint) + ".zip") | |
| self.net = load_model(checkpoint, False) # first load on CPU | |
| if conf["whiten_name"]: | |
| assert conf["whiten_name"] in self.net.pca | |
| def _forward(self, data): | |
| image = data["image"] | |
| assert image.shape[1] == 3 | |
| mean = self.net.preprocess["mean"] | |
| std = self.net.preprocess["std"] | |
| image = image - image.new_tensor(mean)[:, None, None] | |
| image = image / image.new_tensor(std)[:, None, None] | |
| desc = self.net(image) | |
| desc = desc.unsqueeze(0) # batch dimension | |
| if self.conf["whiten_name"]: | |
| pca = self.net.pca[self.conf["whiten_name"]] | |
| desc = common.whiten_features( | |
| desc.cpu().numpy(), pca, **self.conf["whiten_params"] | |
| ) | |
| desc = torch.from_numpy(desc) | |
| return { | |
| "global_descriptor": desc, | |
| } | |