glue-factory-custom/gluefactory/datasets/hpatches.py

"""
Simply load images from a folder or nested folders (does not have any split).
"""
import argparse
import logging
import tarfile

import matplotlib.pyplot as plt
import numpy as np
import torch
from omegaconf import OmegaConf

from ..settings import DATA_PATH
from ..utils.image import ImagePreprocessor, load_image
from ..utils.tools import fork_rng
from ..visualization.viz2d import plot_image_grid
from .base_dataset import BaseDataset

logger = logging.getLogger(__name__)


def read_homography(path):
    with open(path) as f:
        result = []
        for line in f.readlines():
            while "  " in line:  # Remove double spaces
                line = line.replace("  ", " ")
            line = line.replace(" \n", "").replace("\n", "")
            # Split and discard empty strings
            elements = list(filter(lambda s: s, line.split(" ")))
            if elements:
                result.append(elements)
        return np.array(result).astype(float)


class HPatches(BaseDataset, torch.utils.data.Dataset):
    default_conf = {
        "preprocessing": ImagePreprocessor.default_conf,
        "data_dir": "hpatches-sequences-release",
        "subset": None,
        "ignore_large_images": True,
        "grayscale": False,
    }

    # Large images that were ignored in previous papers
    ignored_scenes = (
        "i_contruction",
        "i_crownnight",
        "i_dc",
        "i_pencils",
        "i_whitebuilding",
        "v_artisans",
        "v_astronautis",
        "v_talent",
    )
    url = "http://icvl.ee.ic.ac.uk/vbalnt/hpatches/hpatches-sequences-release.tar.gz"

    def _init(self, conf):
        assert conf.batch_size == 1
        self.preprocessor = ImagePreprocessor(conf.preprocessing)

        self.root = DATA_PATH / conf.data_dir
        if not self.root.exists():
            logger.info("Downloading the HPatches dataset.")
            self.download()
        self.sequences = sorted([x.name for x in self.root.iterdir()])
        if not self.sequences:
            raise ValueError("No image found!")
        self.items = []  # (seq, q_idx, is_illu)
        for seq in self.sequences:
            if conf.ignore_large_images and seq in self.ignored_scenes:
                continue
            if conf.subset is not None and conf.subset != seq[0]:
                continue
            for i in range(2, 7):
                self.items.append((seq, i, seq[0] == "i"))

    def download(self):
        data_dir = self.root.parent
        data_dir.mkdir(exist_ok=True, parents=True)
        tar_path = data_dir / self.url.rsplit("/", 1)[-1]
        torch.hub.download_url_to_file(self.url, tar_path)
        with tarfile.open(tar_path) as tar:
            tar.extractall(data_dir)
        tar_path.unlink()

    def get_dataset(self, split):
        assert split in ["val", "test"]
        return self

    def _read_image(self, seq: str, idx: int) -> dict:
        img = load_image(self.root / seq / f"{idx}.ppm", self.conf.grayscale)
        return self.preprocessor(img)

    def __getitem__(self, idx):
        seq, q_idx, is_illu = self.items[idx]
        data0 = self._read_image(seq, 1)
        data1 = self._read_image(seq, q_idx)
        H = read_homography(self.root / seq / f"H_1_{q_idx}")
        H = data1["transform"] @ H @ np.linalg.inv(data0["transform"])
        return {
            "H_0to1": H.astype(np.float32),
            "scene": seq,
            "idx": idx,
            "is_illu": is_illu,
            "name": f"{seq}/{idx}.ppm",
            "view0": data0,
            "view1": data1,
        }

    def __len__(self):
        return len(self.items)


def visualize(args):
    conf = {
        "batch_size": 1,
        "num_workers": 8,
        "prefetch_factor": 1,
    }
    conf = OmegaConf.merge(conf, OmegaConf.from_cli(args.dotlist))
    dataset = HPatches(conf)
    loader = dataset.get_data_loader("test")
    logger.info("The dataset has %d elements.", len(loader))

    with fork_rng(seed=dataset.conf.seed):
        images = []
        for _, data in zip(range(args.num_items), loader):
            images.append(
                (data[f"view{i}"]["image"][0].permute(1, 2, 0) for i in range(2))
            )
    plot_image_grid(images, dpi=args.dpi)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    from .. import logger  # overwrite the logger

    parser = argparse.ArgumentParser()
    parser.add_argument("--num_items", type=int, default=8)
    parser.add_argument("--dpi", type=int, default=100)
    parser.add_argument("dotlist", nargs="*")
    args = parser.parse_intermixed_args()
    visualize(args)