vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/gim/gluefactory/datasets/megadepth.py

@@ -0,0 +1,514 @@
+import argparse
+import logging
+import shutil
+import tarfile
+from collections.abc import Iterable
+from pathlib import Path
+
+import h5py
+import matplotlib.pyplot as plt
+import numpy as np
+import PIL.Image
+import torch
+from omegaconf import OmegaConf
+
+from ..geometry.wrappers import Camera, Pose
+from ..models.cache_loader import CacheLoader
+from ..settings import DATA_PATH
+from ..utils.image import ImagePreprocessor, load_image
+from ..utils.tools import fork_rng
+from ..visualization.viz2d import plot_heatmaps, plot_image_grid
+from .base_dataset import BaseDataset
+from .utils import rotate_intrinsics, rotate_pose_inplane, scale_intrinsics
+
+logger = logging.getLogger(__name__)
+scene_lists_path = Path(__file__).parent / "megadepth_scene_lists"
+
+
+def sample_n(data, num, seed=None):
+    if len(data) > num:
+        selected = np.random.RandomState(seed).choice(len(data), num, replace=False)
+        return data[selected]
+    else:
+        return data
+
+
+class MegaDepth(BaseDataset):
+    default_conf = {
+        # paths
+        "data_dir": "megadepth/",
+        "depth_subpath": "depth_undistorted/",
+        "image_subpath": "Undistorted_SfM/",
+        "info_dir": "scene_info/",  # @TODO: intrinsics problem?
+        # Training
+        "train_split": "train_scenes_clean.txt",
+        "train_num_per_scene": 500,
+        # Validation
+        "val_split": "valid_scenes_clean.txt",
+        "val_num_per_scene": None,
+        "val_pairs": None,
+        # Test
+        "test_split": "test_scenes_clean.txt",
+        "test_num_per_scene": None,
+        "test_pairs": None,
+        # data sampling
+        "views": 2,
+        "min_overlap": 0.3,  # only with D2-Net format
+        "max_overlap": 1.0,  # only with D2-Net format
+        "num_overlap_bins": 1,
+        "sort_by_overlap": False,
+        "triplet_enforce_overlap": False,  # only with views==3
+        # image options
+        "read_depth": True,
+        "read_image": True,
+        "grayscale": False,
+        "preprocessing": ImagePreprocessor.default_conf,
+        "p_rotate": 0.0,  # probability to rotate image by +/- 90°
+        "reseed": False,
+        "seed": 0,
+        # features from cache
+        "load_features": {
+            "do": False,
+            **CacheLoader.default_conf,
+            "collate": False,
+        },
+    }
+
+    def _init(self, conf):
+        if not (DATA_PATH / conf.data_dir).exists():
+            logger.info("Downloading the MegaDepth dataset.")
+            self.download()
+
+    def download(self):
+        data_dir = DATA_PATH / self.conf.data_dir
+        tmp_dir = data_dir.parent / "megadepth_tmp"
+        if tmp_dir.exists():  # The previous download failed.
+            shutil.rmtree(tmp_dir)
+        tmp_dir.mkdir(exist_ok=True, parents=True)
+        url_base = "https://cvg-data.inf.ethz.ch/megadepth/"
+        for tar_name, out_name in (
+            ("Undistorted_SfM.tar.gz", self.conf.image_subpath),
+            ("depth_undistorted.tar.gz", self.conf.depth_subpath),
+            ("scene_info.tar.gz", self.conf.info_dir),
+        ):
+            tar_path = tmp_dir / tar_name
+            torch.hub.download_url_to_file(url_base + tar_name, tar_path)
+            with tarfile.open(tar_path) as tar:
+                tar.extractall(path=tmp_dir)
+            tar_path.unlink()
+            shutil.move(tmp_dir / tar_name.split(".")[0], tmp_dir / out_name)
+        shutil.move(tmp_dir, data_dir)
+
+    def get_dataset(self, split):
+        assert self.conf.views in [1, 2, 3]
+        if self.conf.views == 3:
+            return _TripletDataset(self.conf, split)
+        else:
+            return _PairDataset(self.conf, split)
+
+
+class _PairDataset(torch.utils.data.Dataset):
+    def __init__(self, conf, split, load_sample=True):
+        self.root = DATA_PATH / conf.data_dir
+        assert self.root.exists(), self.root
+        self.split = split
+        self.conf = conf
+
+        split_conf = conf[split + "_split"]
+        if isinstance(split_conf, (str, Path)):
+            scenes_path = scene_lists_path / split_conf
+            scenes = scenes_path.read_text().rstrip("\n").split("\n")
+        elif isinstance(split_conf, Iterable):
+            scenes = list(split_conf)
+        else:
+            raise ValueError(f"Unknown split configuration: {split_conf}.")
+        scenes = sorted(set(scenes))
+
+        if conf.load_features.do:
+            self.feature_loader = CacheLoader(conf.load_features)
+
+        self.preprocessor = ImagePreprocessor(conf.preprocessing)
+
+        self.images = {}
+        self.depths = {}
+        self.poses = {}
+        self.intrinsics = {}
+        self.valid = {}
+
+        # load metadata
+        self.info_dir = self.root / self.conf.info_dir
+        self.scenes = []
+        for scene in scenes:
+            path = self.info_dir / (scene + ".npz")
+            try:
+                info = np.load(str(path), allow_pickle=True)
+            except Exception:
+                logger.warning(
+                    "Cannot load scene info for scene %s at %s.", scene, path
+                )
+                continue
+            self.images[scene] = info["image_paths"]
+            self.depths[scene] = info["depth_paths"]
+            self.poses[scene] = info["poses"]
+            self.intrinsics[scene] = info["intrinsics"]
+            self.scenes.append(scene)
+
+        if load_sample:
+            self.sample_new_items(conf.seed)
+            assert len(self.items) > 0
+
+    def sample_new_items(self, seed):
+        logger.info("Sampling new %s data with seed %d.", self.split, seed)
+        self.items = []
+        split = self.split
+        num_per_scene = self.conf[self.split + "_num_per_scene"]
+        if isinstance(num_per_scene, Iterable):
+            num_pos, num_neg = num_per_scene
+        else:
+            num_pos = num_per_scene
+            num_neg = None
+        if split != "train" and self.conf[split + "_pairs"] is not None:
+            # Fixed validation or test pairs
+            assert num_pos is None
+            assert num_neg is None
+            assert self.conf.views == 2
+            pairs_path = scene_lists_path / self.conf[split + "_pairs"]
+            for line in pairs_path.read_text().rstrip("\n").split("\n"):
+                im0, im1 = line.split(" ")
+                scene = im0.split("/")[0]
+                assert im1.split("/")[0] == scene
+                im0, im1 = [self.conf.image_subpath + im for im in [im0, im1]]
+                assert im0 in self.images[scene]
+                assert im1 in self.images[scene]
+                idx0 = np.where(self.images[scene] == im0)[0][0]
+                idx1 = np.where(self.images[scene] == im1)[0][0]
+                self.items.append((scene, idx0, idx1, 1.0))
+        elif self.conf.views == 1:
+            for scene in self.scenes:
+                if scene not in self.images:
+                    continue
+                valid = (self.images[scene] != None) | (  # noqa: E711
+                    self.depths[scene] != None  # noqa: E711
+                )
+                ids = np.where(valid)[0]
+                if num_pos and len(ids) > num_pos:
+                    ids = np.random.RandomState(seed).choice(
+                        ids, num_pos, replace=False
+                    )
+                ids = [(scene, i) for i in ids]
+                self.items.extend(ids)
+        else:
+            for scene in self.scenes:
+                path = self.info_dir / (scene + ".npz")
+                assert path.exists(), path
+                info = np.load(str(path), allow_pickle=True)
+                valid = (self.images[scene] != None) & (  # noqa: E711
+                    self.depths[scene] != None  # noqa: E711
+                )
+                ind = np.where(valid)[0]
+                mat = info["overlap_matrix"][valid][:, valid]
+
+                if num_pos is not None:
+                    # Sample a subset of pairs, binned by overlap.
+                    num_bins = self.conf.num_overlap_bins
+                    assert num_bins > 0
+                    bin_width = (
+                        self.conf.max_overlap - self.conf.min_overlap
+                    ) / num_bins
+                    num_per_bin = num_pos // num_bins
+                    pairs_all = []
+                    for k in range(num_bins):
+                        bin_min = self.conf.min_overlap + k * bin_width
+                        bin_max = bin_min + bin_width
+                        pairs_bin = (mat > bin_min) & (mat <= bin_max)
+                        pairs_bin = np.stack(np.where(pairs_bin), -1)
+                        pairs_all.append(pairs_bin)
+                    # Skip bins with too few samples
+                    has_enough_samples = [len(p) >= num_per_bin * 2 for p in pairs_all]
+                    num_per_bin_2 = num_pos // max(1, sum(has_enough_samples))
+                    pairs = []
+                    for pairs_bin, keep in zip(pairs_all, has_enough_samples):
+                        if keep:
+                            pairs.append(sample_n(pairs_bin, num_per_bin_2, seed))
+                    pairs = np.concatenate(pairs, 0)
+                else:
+                    pairs = (mat > self.conf.min_overlap) & (
+                        mat <= self.conf.max_overlap
+                    )
+                    pairs = np.stack(np.where(pairs), -1)
+
+                pairs = [(scene, ind[i], ind[j], mat[i, j]) for i, j in pairs]
+                if num_neg is not None:
+                    neg_pairs = np.stack(np.where(mat <= 0.0), -1)
+                    neg_pairs = sample_n(neg_pairs, num_neg, seed)
+                    pairs += [(scene, ind[i], ind[j], mat[i, j]) for i, j in neg_pairs]
+                self.items.extend(pairs)
+        if self.conf.views == 2 and self.conf.sort_by_overlap:
+            self.items.sort(key=lambda i: i[-1], reverse=True)
+        else:
+            np.random.RandomState(seed).shuffle(self.items)
+
+    def _read_view(self, scene, idx):
+        path = self.root / self.images[scene][idx]
+
+        # read pose data
+        K = self.intrinsics[scene][idx].astype(np.float32, copy=False)
+        T = self.poses[scene][idx].astype(np.float32, copy=False)
+
+        # read image
+        if self.conf.read_image:
+            img = load_image(self.root / self.images[scene][idx], self.conf.grayscale)
+        else:
+            size = PIL.Image.open(path).size[::-1]
+            img = torch.zeros(
+                [3 - 2 * int(self.conf.grayscale), size[0], size[1]]
+            ).float()
+
+        # read depth
+        if self.conf.read_depth:
+            # depth_path = (
+            #     self.root / self.conf.depth_subpath / scene / (path.stem + ".h5")
+            # )
+            depth_subpath = self.depths[scene][idx]
+            depth_id = depth_subpath.split('/')[-1][:-3]
+            assert depth_id == path.stem
+            depth_path = self.root / depth_subpath
+            with h5py.File(str(depth_path), "r") as f:
+                depth = f["/depth"].__array__().astype(np.float32, copy=False)
+                depth = torch.Tensor(depth)[None]
+            assert depth.shape[-2:] == img.shape[-2:]
+        else:
+            depth = None
+
+        # add random rotations
+        do_rotate = self.conf.p_rotate > 0.0 and self.split == "train"
+        if do_rotate:
+            p = self.conf.p_rotate
+            k = 0
+            if np.random.rand() < p:
+                k = np.random.choice(2, 1, replace=False)[0] * 2 - 1
+                img = np.rot90(img, k=-k, axes=(-2, -1))
+                if self.conf.read_depth:
+                    depth = np.rot90(depth, k=-k, axes=(-2, -1)).copy()
+                K = rotate_intrinsics(K, img.shape, k + 2)
+                T = rotate_pose_inplane(T, k + 2)
+
+        name = path.name
+
+        data = self.preprocessor(img)
+        if depth is not None:
+            data["depth"] = self.preprocessor(depth, interpolation="nearest")["image"][
+                0
+            ]
+        K = scale_intrinsics(K, data["scales"])
+
+        data = {
+            "name": name,
+            "scene": scene,
+            "T_w2cam": Pose.from_4x4mat(T),
+            "depth": depth,
+            "camera": Camera.from_calibration_matrix(K).float(),
+            **data,
+        }
+
+        if self.conf.load_features.do:
+            features = self.feature_loader({k: [v] for k, v in data.items()})
+            if do_rotate and k != 0:
+                # ang = np.deg2rad(k * 90.)
+                kpts = features["keypoints"].copy()
+                x, y = kpts[:, 0].copy(), kpts[:, 1].copy()
+                w, h = data["image_size"]
+                if k == 1:
+                    kpts[:, 0] = w - y
+                    kpts[:, 1] = x
+                elif k == -1:
+                    kpts[:, 0] = y
+                    kpts[:, 1] = h - x
+
+                else:
+                    raise ValueError
+                features["keypoints"] = kpts
+
+            data = {"cache": features, **data}
+        return data
+
+    def __getitem__(self, idx):
+        if self.conf.reseed:
+            with fork_rng(self.conf.seed + idx, False):
+                return self.getitem(idx)
+        else:
+            return self.getitem(idx)
+
+    def getitem(self, idx):
+        if self.conf.views == 2:
+            if isinstance(idx, list):
+                scene, idx0, idx1, overlap = idx
+            else:
+                scene, idx0, idx1, overlap = self.items[idx]
+            data0 = self._read_view(scene, idx0)
+            data1 = self._read_view(scene, idx1)
+            data = {
+                "view0": data0,
+                "view1": data1,
+            }
+            data["T_0to1"] = data1["T_w2cam"] @ data0["T_w2cam"].inv()
+            data["T_1to0"] = data0["T_w2cam"] @ data1["T_w2cam"].inv()
+            data["overlap_0to1"] = overlap
+            data["name"] = f"{scene}/{data0['name']}_{data1['name']}"
+        else:
+            assert self.conf.views == 1
+            scene, idx0 = self.items[idx]
+            data = self._read_view(scene, idx0)
+        data["scene"] = scene
+        data["idx"] = idx
+        return data
+
+    def __len__(self):
+        return len(self.items)
+
+
+class _TripletDataset(_PairDataset):
+    def sample_new_items(self, seed):
+        logging.info("Sampling new triplets with seed %d", seed)
+        self.items = []
+        split = self.split
+        num = self.conf[self.split + "_num_per_scene"]
+        if split != "train" and self.conf[split + "_pairs"] is not None:
+            if Path(self.conf[split + "_pairs"]).exists():
+                pairs_path = Path(self.conf[split + "_pairs"])
+            else:
+                pairs_path = DATA_PATH / "configs" / self.conf[split + "_pairs"]
+            for line in pairs_path.read_text().rstrip("\n").split("\n"):
+                im0, im1, im2 = line.split(" ")
+                assert im0[:4] == im1[:4]
+                scene = im1[:4]
+                idx0 = np.where(self.images[scene] == im0)
+                idx1 = np.where(self.images[scene] == im1)
+                idx2 = np.where(self.images[scene] == im2)
+                self.items.append((scene, idx0, idx1, idx2, 1.0, 1.0, 1.0))
+        else:
+            for scene in self.scenes:
+                path = self.info_dir / (scene + ".npz")
+                assert path.exists(), path
+                info = np.load(str(path), allow_pickle=True)
+                if self.conf.num_overlap_bins > 1:
+                    raise NotImplementedError("TODO")
+                valid = (self.images[scene] != None) & (  # noqa: E711
+                    self.depth[scene] != None  # noqa: E711
+                )
+                ind = np.where(valid)[0]
+                mat = info["overlap_matrix"][valid][:, valid]
+                good = (mat > self.conf.min_overlap) & (mat <= self.conf.max_overlap)
+                triplets = []
+                if self.conf.triplet_enforce_overlap:
+                    pairs = np.stack(np.where(good), -1)
+                    for i0, i1 in pairs:
+                        for i2 in pairs[pairs[:, 0] == i0, 1]:
+                            if good[i1, i2]:
+                                triplets.append((i0, i1, i2))
+                    if len(triplets) > num:
+                        selected = np.random.RandomState(seed).choice(
+                            len(triplets), num, replace=False
+                        )
+                        selected = range(num)
+                        triplets = np.array(triplets)[selected]
+                else:
+                    # we first enforce that each row has >1 pairs
+                    non_unique = good.sum(-1) > 1
+                    ind_r = np.where(non_unique)[0]
+                    good = good[non_unique]
+                    pairs = np.stack(np.where(good), -1)
+                    if len(pairs) > num:
+                        selected = np.random.RandomState(seed).choice(
+                            len(pairs), num, replace=False
+                        )
+                        pairs = pairs[selected]
+                    for idx, (k, i) in enumerate(pairs):
+                        # We now sample a j from row k s.t. i != j
+                        possible_j = np.where(good[k])[0]
+                        possible_j = possible_j[possible_j != i]
+                        selected = np.random.RandomState(seed + idx).choice(
+                            len(possible_j), 1, replace=False
+                        )[0]
+                        triplets.append((ind_r[k], i, possible_j[selected]))
+                    triplets = [
+                        (scene, ind[k], ind[i], ind[j], mat[k, i], mat[k, j], mat[i, j])
+                        for k, i, j in triplets
+                    ]
+                    self.items.extend(triplets)
+        np.random.RandomState(seed).shuffle(self.items)
+
+    def __getitem__(self, idx):
+        scene, idx0, idx1, idx2, overlap01, overlap02, overlap12 = self.items[idx]
+        data0 = self._read_view(scene, idx0)
+        data1 = self._read_view(scene, idx1)
+        data2 = self._read_view(scene, idx2)
+        data = {
+            "view0": data0,
+            "view1": data1,
+            "view2": data2,
+        }
+        data["T_0to1"] = data1["T_w2cam"] @ data0["T_w2cam"].inv()
+        data["T_0to2"] = data2["T_w2cam"] @ data0["T_w2cam"].inv()
+        data["T_1to2"] = data2["T_w2cam"] @ data1["T_w2cam"].inv()
+        data["T_1to0"] = data0["T_w2cam"] @ data1["T_w2cam"].inv()
+        data["T_2to0"] = data0["T_w2cam"] @ data2["T_w2cam"].inv()
+        data["T_2to1"] = data1["T_w2cam"] @ data2["T_w2cam"].inv()
+
+        data["overlap_0to1"] = overlap01
+        data["overlap_0to2"] = overlap02
+        data["overlap_1to2"] = overlap12
+        data["scene"] = scene
+        data["name"] = f"{scene}/{data0['name']}_{data1['name']}_{data2['name']}"
+        return data
+
+    def __len__(self):
+        return len(self.items)
+
+
+def visualize(args):
+    conf = {
+        "min_overlap": 0.1,
+        "max_overlap": 0.7,
+        "num_overlap_bins": 3,
+        "sort_by_overlap": False,
+        "train_num_per_scene": 5,
+        "batch_size": 1,
+        "num_workers": 0,
+        "prefetch_factor": None,
+        "val_num_per_scene": None,
+    }
+    conf = OmegaConf.merge(conf, OmegaConf.from_cli(args.dotlist))
+    dataset = MegaDepth(conf)
+    loader = dataset.get_data_loader(args.split)
+    logger.info("The dataset has elements.", len(loader))
+
+    with fork_rng(seed=dataset.conf.seed):
+        images, depths = [], []
+        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(dataset.conf.views)
+                ]
+            )
+            depths.append(
+                [data[f"view{i}"]["depth"][0] for i in range(dataset.conf.views)]
+            )
+
+    axes = plot_image_grid(images, dpi=args.dpi)
+    for i in range(len(images)):
+        plot_heatmaps(depths[i], axes=axes[i])
+    plt.show()
+
+
+if __name__ == "__main__":
+    from .. import logger  # overwrite the logger
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--split", type=str, default="val")
+    parser.add_argument("--num_items", type=int, default=4)
+    parser.add_argument("--dpi", type=int, default=100)
+    parser.add_argument("dotlist", nargs="*")
+    args = parser.parse_intermixed_args()
+    visualize(args)

vismatch/third_party/gim/gluefactory/datasets/utils.py

@@ -0,0 +1,131 @@
+import cv2
+import numpy as np
+import torch
+
+
+def read_image(path, grayscale=False):
+    """Read an image from path as RGB or grayscale"""
+    mode = cv2.IMREAD_GRAYSCALE if grayscale else cv2.IMREAD_COLOR
+    image = cv2.imread(str(path), mode)
+    if image is None:
+        raise IOError(f"Could not read image at {path}.")
+    if not grayscale:
+        image = image[..., ::-1]
+    return image
+
+
+def numpy_image_to_torch(image):
+    """Normalize the image tensor and reorder the dimensions."""
+    if image.ndim == 3:
+        image = image.transpose((2, 0, 1))  # HxWxC to CxHxW
+    elif image.ndim == 2:
+        image = image[None]  # add channel axis
+    else:
+        raise ValueError(f"Not an image: {image.shape}")
+    return torch.tensor(image / 255.0, dtype=torch.float)
+
+
+def rotate_intrinsics(K, image_shape, rot):
+    """image_shape is the shape of the image after rotation"""
+    assert rot <= 3
+    h, w = image_shape[:2][:: -1 if (rot % 2) else 1]
+    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
+    rot = rot % 4
+    if rot == 1:
+        return np.array(
+            [[fy, 0.0, cy], [0.0, fx, w - cx], [0.0, 0.0, 1.0]], dtype=K.dtype
+        )
+    elif rot == 2:
+        return np.array(
+            [[fx, 0.0, w - cx], [0.0, fy, h - cy], [0.0, 0.0, 1.0]],
+            dtype=K.dtype,
+        )
+    else:  # if rot == 3:
+        return np.array(
+            [[fy, 0.0, h - cy], [0.0, fx, cx], [0.0, 0.0, 1.0]], dtype=K.dtype
+        )
+
+
+def rotate_pose_inplane(i_T_w, rot):
+    rotation_matrices = [
+        np.array(
+            [
+                [np.cos(r), -np.sin(r), 0.0, 0.0],
+                [np.sin(r), np.cos(r), 0.0, 0.0],
+                [0.0, 0.0, 1.0, 0.0],
+                [0.0, 0.0, 0.0, 1.0],
+            ],
+            dtype=np.float32,
+        )
+        for r in [np.deg2rad(d) for d in (0, 270, 180, 90)]
+    ]
+    return np.dot(rotation_matrices[rot], i_T_w)
+
+
+def scale_intrinsics(K, scales):
+    """Scale intrinsics after resizing the corresponding image."""
+    scales = np.diag(np.concatenate([scales, [1.0]]))
+    return np.dot(scales.astype(K.dtype, copy=False), K)
+
+
+def get_divisible_wh(w, h, df=None):
+    if df is not None:
+        w_new, h_new = map(lambda x: int(x // df * df), [w, h])
+    else:
+        w_new, h_new = w, h
+    return w_new, h_new
+
+
+def resize(image, size, fn=None, interp="linear", df=None):
+    """Resize an image to a fixed size, or according to max or min edge."""
+    h, w = image.shape[:2]
+    if isinstance(size, int):
+        scale = size / fn(h, w)
+        h_new, w_new = int(round(h * scale)), int(round(w * scale))
+        w_new, h_new = get_divisible_wh(w_new, h_new, df)
+        scale = (w_new / w, h_new / h)
+    elif isinstance(size, (tuple, list)):
+        h_new, w_new = size
+        scale = (w_new / w, h_new / h)
+    else:
+        raise ValueError(f"Incorrect new size: {size}")
+    mode = {
+        "linear": cv2.INTER_LINEAR,
+        "cubic": cv2.INTER_CUBIC,
+        "nearest": cv2.INTER_NEAREST,
+        "area": cv2.INTER_AREA,
+    }[interp]
+    return cv2.resize(image, (w_new, h_new), interpolation=mode), scale
+
+
+def crop(image, size, random=True, other=None, K=None, return_bbox=False):
+    """Random or deterministic crop of an image, adjust depth and intrinsics."""
+    h, w = image.shape[:2]
+    h_new, w_new = (size, size) if isinstance(size, int) else size
+    top = np.random.randint(0, h - h_new + 1) if random else 0
+    left = np.random.randint(0, w - w_new + 1) if random else 0
+    image = image[top : top + h_new, left : left + w_new]
+    ret = [image]
+    if other is not None:
+        ret += [other[top : top + h_new, left : left + w_new]]
+    if K is not None:
+        K[0, 2] -= left
+        K[1, 2] -= top
+        ret += [K]
+    if return_bbox:
+        ret += [(top, top + h_new, left, left + w_new)]
+    return ret
+
+
+def zero_pad(size, *images):
+    """zero pad images to size x size"""
+    ret = []
+    for image in images:
+        if image is None:
+            ret.append(None)
+            continue
+        h, w = image.shape[:2]
+        padded = np.zeros((size, size) + image.shape[2:], dtype=image.dtype)
+        padded[:h, :w] = image
+        ret.append(padded)
+    return ret

vismatch/third_party/gim/gluefactory/eval/__init__.py

@@ -0,0 +1,20 @@
+import torch
+
+from ..utils.tools import get_class
+from .eval_pipeline import EvalPipeline
+
+
+def get_benchmark(benchmark):
+    return get_class(f"{__name__}.{benchmark}", EvalPipeline)
+
+
+@torch.no_grad()
+def run_benchmark(benchmark, eval_conf, experiment_dir, model=None):
+    """This overwrites existing benchmarks"""
+    experiment_dir.mkdir(exist_ok=True, parents=True)
+    bm = get_benchmark(benchmark)
+
+    pipeline = bm(eval_conf)
+    return pipeline.run(
+        experiment_dir, model=model, overwrite=True, overwrite_eval=True
+    )