hloc/extractors/fire.py

import logging
import subprocess
import sys
from pathlib import Path

import torch
import torchvision.transforms as tvf

from ..utils.base_model import BaseModel

logger = logging.getLogger(__name__)
fire_path = Path(__file__).parent / "../../third_party/fire"
sys.path.append(str(fire_path))


import fire_network


class FIRe(BaseModel):
    default_conf = {
        "global": True,
        "asmk": False,
        "model_name": "fire_SfM_120k.pth",
        "scales": [2.0, 1.414, 1.0, 0.707, 0.5, 0.353, 0.25],  # default params
        "features_num": 1000,  # TODO:not supported now
        "asmk_name": "asmk_codebook.bin",  # TODO:not supported now
        "config_name": "eval_fire.yml",
    }
    required_inputs = ["image"]

    # Models exported using
    fire_models = {
        "fire_SfM_120k.pth": "http://download.europe.naverlabs.com/ComputerVision/FIRe/official/fire.pth",
        "fire_imagenet.pth": "http://download.europe.naverlabs.com/ComputerVision/FIRe/pretraining/fire_imagenet.pth",
    }

    def _init(self, conf):
        assert conf["model_name"] in self.fire_models.keys()
        # Config paths
        model_path = fire_path / "model" / conf["model_name"]

        # Download the model.
        if not model_path.exists():
            model_path.parent.mkdir(exist_ok=True)
            link = self.fire_models[conf["model_name"]]
            cmd = ["wget", "--quiet", link, "-O", str(model_path)]
            logger.info(f"Downloading the FIRe model with `{cmd}`.")
            subprocess.run(cmd, check=True)

        logger.info("Loading fire model...")

        # Load net
        state = torch.load(model_path)
        state["net_params"]["pretrained"] = None
        net = fire_network.init_network(**state["net_params"])
        net.load_state_dict(state["state_dict"])
        self.net = net

        self.norm_rgb = tvf.Normalize(
            **dict(zip(["mean", "std"], net.runtime["mean_std"]))
        )

        # params
        self.scales = conf["scales"]

    def _forward(self, data):
        image = self.norm_rgb(data["image"])

        # Feature extraction.
        desc = self.net.forward_global(image, scales=self.scales)

        return {"global_descriptor": desc}