vismatch 1.1.1__py3-none-any.whl

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

@@ -0,0 +1,33 @@
+import time
+
+import numpy as np
+import torch
+
+
+def benchmark(model, data, device, r=100):
+    timings = np.zeros((r, 1))
+    if device.type == "cuda":
+        starter = torch.cuda.Event(enable_timing=True)
+        ender = torch.cuda.Event(enable_timing=True)
+    # warmup
+    for _ in range(10):
+        _ = model(data)
+    # measurements
+    with torch.no_grad():
+        for rep in range(r):
+            if device.type == "cuda":
+                starter.record()
+                _ = model(data)
+                ender.record()
+                # sync gpu
+                torch.cuda.synchronize()
+                curr_time = starter.elapsed_time(ender)
+            else:
+                start = time.perf_counter()
+                _ = model(data)
+                curr_time = (time.perf_counter() - start) * 1e3
+            timings[rep] = curr_time
+
+    mean_syn = np.sum(timings) / r
+    std_syn = np.std(timings)
+    return {"mean": mean_syn, "std": std_syn}

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

@@ -0,0 +1,134 @@
+"""
+A set of utilities to manage and load checkpoints of training experiments.
+
+Author: Paul-Edouard Sarlin (skydes)
+"""
+
+import logging
+import os
+import re
+import shutil
+from pathlib import Path
+
+import torch
+from omegaconf import OmegaConf
+
+from ..models import get_model
+from ..settings import TRAINING_PATH
+
+logger = logging.getLogger(__name__)
+
+
+def list_checkpoints(dir_):
+    """List all valid checkpoints in a given directory."""
+    checkpoints = []
+    for p in dir_.glob("checkpoint_*.tar"):
+        numbers = re.findall(r"(\d+)", p.name)
+        assert len(numbers) <= 2
+        if len(numbers) == 0:
+            continue
+        if len(numbers) == 1:
+            checkpoints.append((int(numbers[0]), p))
+        else:
+            checkpoints.append((int(numbers[1]), p))
+    return checkpoints
+
+
+def get_last_checkpoint(exper, allow_interrupted=True):
+    """Get the last saved checkpoint for a given experiment name."""
+    ckpts = list_checkpoints(Path(TRAINING_PATH, exper))
+    if not allow_interrupted:
+        ckpts = [(n, p) for (n, p) in ckpts if "_interrupted" not in p.name]
+    assert len(ckpts) > 0
+    return sorted(ckpts)[-1][1]
+
+
+def get_best_checkpoint(exper):
+    """Get the checkpoint with the best loss, for a given experiment name."""
+    p = Path(TRAINING_PATH, exper, "checkpoint_best.tar")
+    return p
+
+
+def delete_old_checkpoints(dir_, num_keep):
+    """Delete all but the num_keep last saved checkpoints."""
+    ckpts = list_checkpoints(dir_)
+    ckpts = sorted(ckpts)[::-1]
+    kept = 0
+    for ckpt in ckpts:
+        if ("_interrupted" in str(ckpt[1]) and kept > 0) or kept >= num_keep:
+            logger.info(f"Deleting checkpoint {ckpt[1].name}")
+            ckpt[1].unlink()
+        else:
+            kept += 1
+
+
+def load_experiment(exper, conf={}, get_last=False, ckpt=None):
+    """Load and return the model of a given experiment."""
+    exper = Path(exper)
+    if exper.suffix != ".tar":
+        if get_last:
+            ckpt = get_last_checkpoint(exper)
+        else:
+            ckpt = get_best_checkpoint(exper)
+    else:
+        ckpt = exper
+    logger.info(f"Loading checkpoint {ckpt.name}")
+    ckpt = torch.load(str(ckpt), map_location="cpu")
+
+    loaded_conf = OmegaConf.create(ckpt["conf"])
+    OmegaConf.set_struct(loaded_conf, False)
+    conf = OmegaConf.merge(loaded_conf.model, OmegaConf.create(conf))
+    model = get_model(conf.name)(conf).eval()
+
+    state_dict = ckpt["model"]
+    dict_params = set(state_dict.keys())
+    model_params = set(map(lambda n: n[0], model.named_parameters()))
+    diff = model_params - dict_params
+    if len(diff) > 0:
+        subs = os.path.commonprefix(list(diff)).rstrip(".")
+        logger.warning(f"Missing {len(diff)} parameters in {subs}")
+    model.load_state_dict(state_dict, strict=False)
+    return model
+
+
+# @TODO: also copy the respective module scripts (i.e. the code)
+def save_experiment(
+    model,
+    optimizer,
+    lr_scheduler,
+    conf,
+    losses,
+    results,
+    best_eval,
+    epoch,
+    iter_i,
+    output_dir,
+    stop=False,
+    distributed=False,
+    cp_name=None,
+):
+    """Save the current model to a checkpoint
+    and return the best result so far."""
+    state = (model.module if distributed else model).state_dict()
+    checkpoint = {
+        "model": state,
+        "optimizer": optimizer.state_dict(),
+        "lr_scheduler": lr_scheduler.state_dict(),
+        "conf": OmegaConf.to_container(conf, resolve=True),
+        "epoch": epoch,
+        "losses": losses,
+        "eval": results,
+    }
+    if cp_name is None:
+        cp_name = (
+            f"checkpoint_{epoch}_{iter_i}" + ("_interrupted" if stop else "") + ".tar"
+        )
+    logger.info(f"Saving checkpoint {cp_name}")
+    cp_path = str(output_dir / cp_name)
+    torch.save(checkpoint, cp_path)
+    if cp_name != "checkpoint_best.tar" and results[conf.train.best_key] < best_eval:
+        best_eval = results[conf.train.best_key]
+        logger.info(f"New best val: {conf.train.best_key}={best_eval}")
+        shutil.copy(cp_path, str(output_dir / "checkpoint_best.tar"))
+    delete_old_checkpoints(output_dir, conf.train.keep_last_checkpoints)
+    return best_eval

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

@@ -0,0 +1,81 @@
+"""
+Export the predictions of a model for a given dataloader (e.g. ImageFolder).
+Use a standalone script with `python3 -m dsfm.scipts.export_predictions dir`
+or call from another script.
+"""
+
+from pathlib import Path
+
+import h5py
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from .tensor import batch_to_device
+
+
+@torch.no_grad()
+def export_predictions(
+    loader,
+    model,
+    output_file,
+    as_half=False,
+    keys="*",
+    callback_fn=None,
+    optional_keys=[],
+):
+    assert keys == "*" or isinstance(keys, (tuple, list))
+    Path(output_file).parent.mkdir(exist_ok=True, parents=True)
+    hfile = h5py.File(str(output_file), "w")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device).eval()
+    for data_ in tqdm(loader):
+        data = batch_to_device(data_, device, non_blocking=True)
+        pred = model(data)
+        if callback_fn is not None:
+            pred = {**callback_fn(pred, data), **pred}
+        if keys != "*":
+            if len(set(keys) - set(pred.keys())) > 0:
+                raise ValueError(f"Missing key {set(keys) - set(pred.keys())}")
+            pred = {k: v for k, v in pred.items() if k in keys + optional_keys}
+        assert len(pred) > 0
+
+        # renormalization
+        for k in pred.keys():
+            if k.startswith("keypoints"):
+                idx = k.replace("keypoints", "")
+                scales = 1.0 / (
+                    data["scales"] if len(idx) == 0 else data[f"view{idx}"]["scales"]
+                )
+                pred[k] = pred[k] * scales[None]
+            if k.startswith("lines"):
+                idx = k.replace("lines", "")
+                scales = 1.0 / (
+                    data["scales"] if len(idx) == 0 else data[f"view{idx}"]["scales"]
+                )
+                pred[k] = pred[k] * scales[None]
+            if k.startswith("orig_lines"):
+                idx = k.replace("orig_lines", "")
+                scales = 1.0 / (
+                    data["scales"] if len(idx) == 0 else data[f"view{idx}"]["scales"]
+                )
+                pred[k] = pred[k] * scales[None]
+
+        pred = {k: v[0].cpu().numpy() for k, v in pred.items()}
+
+        if as_half:
+            for k in pred:
+                dt = pred[k].dtype
+                if (dt == np.float32) and (dt != np.float16):
+                    pred[k] = pred[k].astype(np.float16)
+        try:
+            name = data["name"][0]
+            grp = hfile.create_group(name)
+            for k, v in pred.items():
+                grp.create_dataset(k, data=v)
+        except RuntimeError:
+            continue
+
+        del pred
+    hfile.close()
+    return output_file

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

@@ -0,0 +1,130 @@
+import collections.abc as collections
+from pathlib import Path
+from typing import Optional, Tuple
+
+import cv2
+import kornia
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+
+
+class ImagePreprocessor:
+    default_conf = {
+        "resize": None,  # target edge length, None for no resizing
+        "edge_divisible_by": None,
+        "side": "long",
+        "interpolation": "bilinear",
+        "align_corners": None,
+        "antialias": True,
+        "square_pad": False,
+        "add_padding_mask": False,
+    }
+
+    def __init__(self, conf) -> None:
+        super().__init__()
+        default_conf = OmegaConf.create(self.default_conf)
+        OmegaConf.set_struct(default_conf, True)
+        self.conf = OmegaConf.merge(default_conf, conf)
+
+    def __call__(self, img: torch.Tensor, interpolation: Optional[str] = None) -> dict:
+        """Resize and preprocess an image, return image and resize scale"""
+        h, w = img.shape[-2:]
+        size = h, w
+        if self.conf.resize is not None:
+            if interpolation is None:
+                interpolation = self.conf.interpolation
+            size = self.get_new_image_size(h, w)
+            img = kornia.geometry.transform.resize(
+                img,
+                size,
+                side=self.conf.side,
+                antialias=self.conf.antialias,
+                align_corners=self.conf.align_corners,
+                interpolation=interpolation,
+            )
+        scale = torch.Tensor([img.shape[-1] / w, img.shape[-2] / h]).to(img)
+        T = np.diag([scale[0], scale[1], 1])
+
+        data = {
+            "scales": scale,
+            "image_size": np.array(size[::-1]),
+            "transform": T,
+            "original_image_size": np.array([w, h]),
+        }
+        if self.conf.square_pad:
+            sl = max(img.shape[-2:])
+            data["image"] = torch.zeros(
+                *img.shape[:-2], sl, sl, device=img.device, dtype=img.dtype
+            )
+            data["image"][:, : img.shape[-2], : img.shape[-1]] = img
+            if self.conf.add_padding_mask:
+                data["padding_mask"] = torch.zeros(
+                    *img.shape[:-3], 1, sl, sl, device=img.device, dtype=torch.bool
+                )
+                data["padding_mask"][:, : img.shape[-2], : img.shape[-1]] = True
+
+        else:
+            data["image"] = img
+        return data
+
+    def load_image(self, image_path: Path) -> dict:
+        return self(load_image(image_path))
+
+    def get_new_image_size(
+        self,
+        h: int,
+        w: int,
+    ) -> Tuple[int, int]:
+        side = self.conf.side
+        if isinstance(self.conf.resize, collections.Iterable):
+            assert len(self.conf.resize) == 2
+            return tuple(self.conf.resize)
+        side_size = self.conf.resize
+        aspect_ratio = w / h
+        if side not in ("short", "long", "vert", "horz"):
+            raise ValueError(
+                f"side can be one of 'short', 'long', 'vert', and 'horz'. Got '{side}'"
+            )
+        if side == "vert":
+            size = side_size, int(side_size * aspect_ratio)
+        elif side == "horz":
+            size = int(side_size / aspect_ratio), side_size
+        elif (side == "short") ^ (aspect_ratio < 1.0):
+            size = side_size, int(side_size * aspect_ratio)
+        else:
+            size = int(side_size / aspect_ratio), side_size
+
+        if self.conf.edge_divisible_by is not None:
+            df = self.conf.edge_divisible_by
+            size = list(map(lambda x: int(x // df * df), size))
+        return size
+
+
+def read_image(path: Path, grayscale: bool = False) -> np.ndarray:
+    """Read an image from path as RGB or grayscale"""
+    if not Path(path).exists():
+        raise FileNotFoundError(f"No image at path {path}.")
+    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: np.ndarray) -> torch.Tensor:
+    """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 load_image(path: Path, grayscale=False) -> torch.Tensor:
+    image = read_image(path, grayscale=grayscale)
+    return numpy_image_to_torch(image)

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

@@ -0,0 +1,44 @@
+import torch
+
+
+def to_view(data, i):
+    return {k + i: v for k, v in data.items()}
+
+
+def get_view(data, i):
+    data_g = {k: v for k, v in data.items() if not k[-1].isnumeric()}
+    data_i = {k[:-1]: v for k, v in data.items() if k[-1] == i}
+    return {**data_g, **data_i}
+
+
+def get_twoview(data, idx):
+    li = idx[0]
+    ri = idx[-1]
+    assert idx == f"{li}to{ri}"
+    data_lr = {k[:-4] + "0to1": v for k, v in data.items() if k[-4:] == f"{li}to{ri}"}
+    data_rl = {k[:-4] + "1to0": v for k, v in data.items() if k[-4:] == f"{ri}ito{li}"}
+    data_l = {
+        k[:-1] + "0": v for k, v in data.items() if k[-1:] == li and k[-3:-1] != "to"
+    }
+    data_r = {
+        k[:-1] + "1": v for k, v in data.items() if k[-1:] == ri and k[-3:-1] != "to"
+    }
+    return {**data_lr, **data_rl, **data_l, **data_r}
+
+
+def stack_twoviews(data, indices=["0to1", "0to2", "1to2"]):
+    idx0 = indices[0]
+    m_data = data[idx0] if idx0 in data else get_twoview(data, idx0)
+    # stack on dim=0
+    for idx in indices[1:]:
+        data_i = data[idx] if idx in data else get_twoview(data, idx)
+        for k, v in data_i.items():
+            m_data[k] = torch.cat([m_data[k], v], dim=0)
+    return m_data
+
+
+def unstack_twoviews(data, B, indices=["0to1", "0to2", "1to2"]):
+    out = {}
+    for i, idx in enumerate(indices):
+        out[idx] = {k: v[i * B : (i + 1) * B] for k, v in data.items()}
+    return out

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

@@ -0,0 +1,50 @@
+import torch
+
+
+def extract_patches(
+    tensor: torch.Tensor,
+    required_corners: torch.Tensor,
+    ps: int,
+) -> torch.Tensor:
+    c, h, w = tensor.shape
+    corner = required_corners.long()
+    corner[:, 0] = corner[:, 0].clamp(min=0, max=w - 1 - ps)
+    corner[:, 1] = corner[:, 1].clamp(min=0, max=h - 1 - ps)
+    offset = torch.arange(0, ps)
+
+    kw = {"indexing": "ij"} if torch.__version__ >= "1.10" else {}
+    x, y = torch.meshgrid(offset, offset, **kw)
+    patches = torch.stack((x, y)).permute(2, 1, 0).unsqueeze(2)
+    patches = patches.to(corner) + corner[None, None]
+    pts = patches.reshape(-1, 2)
+    sampled = tensor.permute(1, 2, 0)[tuple(pts.T)[::-1]]
+    sampled = sampled.reshape(ps, ps, -1, c)
+    assert sampled.shape[:3] == patches.shape[:3]
+    return sampled.permute(2, 3, 0, 1), corner.float()
+
+
+def batch_extract_patches(tensor: torch.Tensor, kpts: torch.Tensor, ps: int):
+    b, c, h, w = tensor.shape
+    b, n, _ = kpts.shape
+    out = torch.zeros((b, n, c, ps, ps), dtype=tensor.dtype, device=tensor.device)
+    corners = torch.zeros((b, n, 2), dtype=tensor.dtype, device=tensor.device)
+    for i in range(b):
+        out[i], corners[i] = extract_patches(tensor[i], kpts[i] - ps / 2 - 1, ps)
+    return out, corners
+
+
+def draw_image_patches(img, patches, corners):
+    b, c, h, w = img.shape
+    b, n, c, p, p = patches.shape
+    b, n, _ = corners.shape
+    for i in range(b):
+        for k in range(n):
+            y, x = corners[i, k]
+            img[i, :, x : x + p, y : y + p] = patches[i, k]
+
+
+def build_heatmap(img, patches, corners):
+    hmap = torch.zeros_like(img)
+    draw_image_patches(hmap, patches, corners.long())
+    hmap = hmap.squeeze(1)
+    return hmap, (hmap > 0.0).float()  # bxhxw

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

@@ -0,0 +1,134 @@
+"""
+Based on sacred/stdout_capturing.py in project Sacred
+https://github.com/IDSIA/sacred
+
+Author: Paul-Edouard Sarlin (skydes)
+"""
+
+from __future__ import division, print_function, unicode_literals
+
+import os
+import subprocess
+import sys
+from contextlib import contextmanager
+from threading import Timer
+
+
+def apply_backspaces_and_linefeeds(text):
+    """
+    Interpret backspaces and linefeeds in text like a terminal would.
+    Interpret text like a terminal by removing backspace and linefeed
+    characters and applying them line by line.
+    If final line ends with a carriage it keeps it to be concatenable with next
+    output chunk.
+    """
+    orig_lines = text.split("\n")
+    orig_lines_len = len(orig_lines)
+    new_lines = []
+    for orig_line_idx, orig_line in enumerate(orig_lines):
+        chars, cursor = [], 0
+        orig_line_len = len(orig_line)
+        for orig_char_idx, orig_char in enumerate(orig_line):
+            if orig_char == "\r" and (
+                orig_char_idx != orig_line_len - 1
+                or orig_line_idx != orig_lines_len - 1
+            ):
+                cursor = 0
+            elif orig_char == "\b":
+                cursor = max(0, cursor - 1)
+            else:
+                if (
+                    orig_char == "\r"
+                    and orig_char_idx == orig_line_len - 1
+                    and orig_line_idx == orig_lines_len - 1
+                ):
+                    cursor = len(chars)
+                if cursor == len(chars):
+                    chars.append(orig_char)
+                else:
+                    chars[cursor] = orig_char
+                cursor += 1
+        new_lines.append("".join(chars))
+    return "\n".join(new_lines)
+
+
+def flush():
+    """Try to flush all stdio buffers, both from python and from C."""
+    try:
+        sys.stdout.flush()
+        sys.stderr.flush()
+    except (AttributeError, ValueError, IOError):
+        pass  # unsupported
+
+
+# Duplicate stdout and stderr to a file. Inspired by:
+# http://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/
+# http://stackoverflow.com/a/651718/1388435
+# http://stackoverflow.com/a/22434262/1388435
+@contextmanager
+def capture_outputs(filename):
+    """Duplicate stdout and stderr to a file on the file descriptor level."""
+    with open(str(filename), "a+") as target:
+        original_stdout_fd = 1
+        original_stderr_fd = 2
+        target_fd = target.fileno()
+
+        # Save a copy of the original stdout and stderr file descriptors
+        saved_stdout_fd = os.dup(original_stdout_fd)
+        saved_stderr_fd = os.dup(original_stderr_fd)
+
+        tee_stdout = subprocess.Popen(
+            ["tee", "-a", "-i", "/dev/stderr"],
+            start_new_session=True,
+            stdin=subprocess.PIPE,
+            stderr=target_fd,
+            stdout=1,
+        )
+        tee_stderr = subprocess.Popen(
+            ["tee", "-a", "-i", "/dev/stderr"],
+            start_new_session=True,
+            stdin=subprocess.PIPE,
+            stderr=target_fd,
+            stdout=2,
+        )
+
+        flush()
+        os.dup2(tee_stdout.stdin.fileno(), original_stdout_fd)
+        os.dup2(tee_stderr.stdin.fileno(), original_stderr_fd)
+
+        try:
+            yield
+        finally:
+            flush()
+
+            # then redirect stdout back to the saved fd
+            tee_stdout.stdin.close()
+            tee_stderr.stdin.close()
+
+            # restore original fds
+            os.dup2(saved_stdout_fd, original_stdout_fd)
+            os.dup2(saved_stderr_fd, original_stderr_fd)
+
+            # wait for completion of the tee processes with timeout
+            # implemented using a timer because timeout support is py3 only
+            def kill_tees():
+                tee_stdout.kill()
+                tee_stderr.kill()
+
+            tee_timer = Timer(1, kill_tees)
+            try:
+                tee_timer.start()
+                tee_stdout.wait()
+                tee_stderr.wait()
+            finally:
+                tee_timer.cancel()
+
+            os.close(saved_stdout_fd)
+            os.close(saved_stderr_fd)
+
+    # Cleanup log file
+    with open(str(filename), "r") as target:
+        text = target.read()
+    text = apply_backspaces_and_linefeeds(text)
+    with open(str(filename), "w") as target:
+        target.write(text)

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

@@ -0,0 +1,48 @@
+"""
+Author: Paul-Edouard Sarlin (skydes)
+"""
+
+import collections.abc as collections
+
+import numpy as np
+import torch
+
+string_classes = (str, bytes)
+
+
+def map_tensor(input_, func):
+    if isinstance(input_, string_classes):
+        return input_
+    elif isinstance(input_, collections.Mapping):
+        return {k: map_tensor(sample, func) for k, sample in input_.items()}
+    elif isinstance(input_, collections.Sequence):
+        return [map_tensor(sample, func) for sample in input_]
+    elif input_ is None:
+        return None
+    else:
+        return func(input_)
+
+
+def batch_to_numpy(batch):
+    return map_tensor(batch, lambda tensor: tensor.cpu().numpy())
+
+
+def batch_to_device(batch, device, non_blocking=True):
+    def _func(tensor):
+        return tensor.to(device=device, non_blocking=non_blocking)
+
+    return map_tensor(batch, _func)
+
+
+def rbd(data: dict) -> dict:
+    """Remove batch dimension from elements in data"""
+    return {
+        k: v[0] if isinstance(v, (torch.Tensor, np.ndarray, list)) else v
+        for k, v in data.items()
+    }
+
+
+def index_batch(tensor_dict):
+    batch_size = len(next(iter(tensor_dict.values())))
+    for i in range(batch_size):
+        yield map_tensor(tensor_dict, lambda t: t[i])