hloc/extractors/dir.py

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,
        }