vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/rdd/training/utils.py

@@ -0,0 +1,246 @@
+import torch
+import numpy as np
+import pdb
+from copy import deepcopy
+import cv2
+debug_cnt = -1
+from RDD.matchers import DualSoftmaxMatcher
+
+matcher = DualSoftmaxMatcher(inv_temperature = 1)
+
+def make_batch(augmentor, difficulty = 0.3, train = True):
+    Hs = []
+    img_list = augmentor.train if train else augmentor.test
+    dev = augmentor.device
+    batch_images = []
+
+    with torch.no_grad(): # we dont require grads in the augmentation
+        for b in range(augmentor.batch_size):
+            rdidx = np.random.randint(len(img_list))
+            img = torch.tensor(img_list[rdidx], dtype=torch.float32).permute(2,0,1).to(augmentor.device).unsqueeze(0)
+            batch_images.append(img)
+
+        batch_images = torch.cat(batch_images)
+
+        p1, H1 = augmentor(batch_images, difficulty)
+        p2, H2 = augmentor(batch_images, difficulty, TPS = True, prob_deformation = 0.7)
+
+    return p1, p2, H1, H2
+
+def plot_corrs(p1, p2, src_pts, tgt_pts):
+    import matplotlib.pyplot as plt
+    p1 = p1.cpu()
+    p2 = p2.cpu()
+    src_pts = src_pts.cpu() ; tgt_pts = tgt_pts.cpu()
+    rnd_idx = np.random.randint(len(src_pts), size=200)
+    src_pts = src_pts[rnd_idx, ...]
+    tgt_pts = tgt_pts[rnd_idx, ...]
+
+    #Plot ground-truth correspondences
+    fig, ax = plt.subplots(1,2,figsize=(18, 12))
+    colors = np.random.uniform(size=(len(tgt_pts),3))
+    #Src image
+    img = p1
+    for i, p in enumerate(src_pts):
+        ax[0].scatter(p[0],p[1],color=colors[i])
+    ax[0].imshow(img.permute(1,2,0).numpy()[...,::-1])
+
+    #Target img
+    img2 = p2
+    for i, p in enumerate(tgt_pts):
+        ax[1].scatter(p[0],p[1],color=colors[i])
+    ax[1].imshow(img2.permute(1,2,0).numpy()[...,::-1])
+    plt.show()
+
+
+def get_corresponding_pts(p1, p2, H, H2, augmentor, h, w, crop = None):
+    '''
+        Get dense corresponding points
+    '''
+    global debug_cnt
+    negatives, positives = [], []
+
+    with torch.no_grad():
+        #real input res of samples
+        rh, rw = p1.shape[-2:]
+        ratio = torch.tensor([rw/w, rh/h], device = p1.device)
+
+        (H, mask1) = H
+        (H2, src, W, A, mask2) = H2
+
+        #Generate meshgrid of target pts
+        x, y = torch.meshgrid(torch.arange(w, device=p1.device), torch.arange(h, device=p1.device), indexing ='xy')
+        mesh = torch.cat([x.unsqueeze(-1), y.unsqueeze(-1)], dim=-1)
+        target_pts = mesh.view(-1, 2) * ratio
+
+        #Pack all transformations into T
+        for batch_idx in range(len(p1)):
+            with torch.no_grad():
+                T = (H[batch_idx], H2[batch_idx], 
+                    src[batch_idx].unsqueeze(0), W[batch_idx].unsqueeze(0), A[batch_idx].unsqueeze(0))
+                #We now warp the target points to src image
+                src_pts = (augmentor.get_correspondences(target_pts, T) ) #target to src 
+                tgt_pts = (target_pts)
+            
+                #Check out of bounds points
+                mask_valid = (src_pts[:, 0] >=0) & (src_pts[:, 1] >=0) & \
+                            (src_pts[:, 0] < rw) & (src_pts[:, 1] < rh)
+
+                negatives.append( tgt_pts[~mask_valid] )            
+                tgt_pts = tgt_pts[mask_valid]
+                src_pts = src_pts[mask_valid]
+
+
+                #Remove invalid pixels
+                mask_valid =    mask1[batch_idx, src_pts[:,1].long(), src_pts[:,0].long()]  & \
+                                mask2[batch_idx, tgt_pts[:,1].long(), tgt_pts[:,0].long()]
+                tgt_pts = tgt_pts[mask_valid]
+                src_pts = src_pts[mask_valid]
+
+                # limit nb of matches if desired
+                if crop is not None:
+                    rnd_idx = torch.randperm(len(src_pts), device=src_pts.device)[:crop]
+                    src_pts = src_pts[rnd_idx]
+                    tgt_pts = tgt_pts[rnd_idx]
+
+                if debug_cnt >=0 and debug_cnt < 4:
+                    plot_corrs(p1[batch_idx], p2[batch_idx], src_pts , tgt_pts )
+                    debug_cnt +=1
+
+                src_pts = (src_pts / ratio)
+                tgt_pts = (tgt_pts / ratio)
+
+                #Check out of bounds points
+                padto = 10 if crop is not None else 2
+                mask_valid1 = (src_pts[:, 0] >= (0 + padto)) & (src_pts[:, 1] >= (0 + padto)) & \
+                             (src_pts[:, 0] < (w - padto)) & (src_pts[:, 1] < (h - padto))
+                mask_valid2 = (tgt_pts[:, 0] >= (0 + padto)) & (tgt_pts[:, 1] >= (0 + padto)) & \
+                             (tgt_pts[:, 0] < (w - padto)) & (tgt_pts[:, 1] < (h - padto))
+                mask_valid = mask_valid1 & mask_valid2
+                tgt_pts = tgt_pts[mask_valid]
+                src_pts = src_pts[mask_valid]         
+
+                #Remove repeated correspondences
+                lut_mat = torch.ones((h, w, 4), device = src_pts.device, dtype = src_pts.dtype) * -1
+                # src_pts_np = src_pts.cpu().numpy()
+                # tgt_pts_np = tgt_pts.cpu().numpy()
+                try:
+                    lut_mat[src_pts[:,1].long(), src_pts[:,0].long()] = torch.cat([src_pts, tgt_pts], dim=1)
+                    mask_valid = torch.all(lut_mat >= 0, dim=-1)
+                    points = lut_mat[mask_valid]
+                    positives.append(points)
+                except:
+                    pdb.set_trace()
+                    print('..')
+
+    return negatives, positives
+
+
+def crop_patches(tensor, coords, size = 7):
+    '''
+        Crop [size x size] patches around 2D coordinates from a tensor.
+    '''
+    B, C, H, W = tensor.shape
+
+    x, y = coords[:, 0], coords[:, 1]
+    y = y.view(-1, 1, 1)
+    x = x.view(-1, 1, 1)
+    halfsize = size // 2
+    # Create meshgrid for indexing
+    x_offset, y_offset = torch.meshgrid(torch.arange(-halfsize, halfsize+1), torch.arange(-halfsize, halfsize+1), indexing='xy')
+    y_offset = y_offset.to(tensor.device)
+    x_offset = x_offset.to(tensor.device)
+
+    # Compute indices around each coordinate
+    y_indices = (y + y_offset.view(1, size, size)).squeeze(0) + halfsize
+    x_indices = (x + x_offset.view(1, size, size)).squeeze(0) + halfsize
+
+    # Handle out-of-boundary indices with padding
+    tensor_padded = torch.nn.functional.pad(tensor, (halfsize, halfsize, halfsize, halfsize), mode='constant')
+
+    # Index tensor to get patches
+    patches = tensor_padded[:, :, y_indices, x_indices] # [B, C, N, H, W]
+    return patches
+
+def subpix_softmax2d(heatmaps, temp = 0.25):
+    N, H, W = heatmaps.shape
+    heatmaps = torch.softmax(temp * heatmaps.view(-1, H*W), -1).view(-1, H, W)
+    x, y = torch.meshgrid(torch.arange(W, device =  heatmaps.device ), torch.arange(H, device =  heatmaps.device ), indexing = 'xy')
+    x = x - (W//2)
+    y = y - (H//2)
+    #pdb.set_trace()
+    coords_x = (x[None, ...] * heatmaps)
+    coords_y = (y[None, ...] * heatmaps)
+    coords = torch.cat([coords_x[..., None], coords_y[..., None]], -1).view(N, H*W, 2)
+    coords = coords.sum(1)
+
+    return coords
+
+
+def check_accuracy(X, Y, pts1 = None, pts2 = None):
+    with torch.no_grad():
+        #dist_mat = torch.cdist(X,Y)
+        dist_mat = X @ Y.t()
+        nn = torch.argmax(dist_mat, dim=1)
+        #nn = torch.argmin(dist_mat, dim=1)
+        correct = nn == torch.arange(len(X), device = X.device)
+
+        acc = correct.sum().item() / len(X)
+        return acc
+
+def get_nb_trainable_params(model):
+	model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+	nb_params = sum([np.prod(p.size()) for p in model_parameters])
+ 
+	print('Number of trainable parameters: {:d}'.format(nb_params))
+ 
+ 
+def plot_keypoints(image, kpts, radius=2, color=(255, 0, 0)):
+    image = image.cpu().detach().numpy() if isinstance(image, torch.Tensor) else image
+    kpts = kpts.cpu().detach().numpy() if isinstance(kpts, torch.Tensor) else kpts
+
+    if image.dtype is not np.dtype('uint8'):
+        image = image * 255
+        image = image.astype(np.uint8)
+
+    if len(image.shape) == 2 or image.shape[2] == 1:
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+
+    out = np.ascontiguousarray(deepcopy(image))
+    kpts = np.round(kpts).astype(int)
+
+    for kpt in kpts:
+        y0, x0 = kpt
+        cv2.drawMarker(out, (x0, y0), color, cv2.MARKER_CROSS, radius)
+
+        # cv2.circle(out, (x0, y0), radius, color, -1, lineType=cv2.LINE_4)
+    return out
+
+def save_image_in_actual_size(image, name):
+    import matplotlib.pyplot as plt
+
+    dpi = 100
+    height, width = image.shape[:2]
+
+    # What size does the figure need to be in inches to fit the image?
+    figsize = width / float(dpi), height / float(dpi)
+
+    # Create a figure of the right size with one axes that takes up the full figure
+    fig = plt.figure(figsize=figsize)
+    ax = fig.add_axes([0, 0, 1, 1])
+
+    # Hide spines, ticks, etc.
+    ax.axis('off')
+
+    # Display the image.
+    if len(image.shape) == 3:
+        ax.imshow(image, cmap='gray')
+    elif len(image.shape) == 2:
+        if image.dtype == np.uint8:
+            ax.imshow(image, cmap='gray')
+        else:
+            ax.imshow(image)
+            ax.text(20, 20, f"Range: {image.min():g}~{image.max():g}", color='red')
+    
+    # save the image
+    plt.savefig(name, dpi=dpi)

vismatch/utils.py

@@ -0,0 +1,390 @@
+import importlib
+import logging
+from pathlib import Path
+import numpy as np
+import torch
+from PIL import Image
+import torchvision.transforms as tfm
+from yacs.config import CfgNode as CN
+import sys
+
+logger = logging.getLogger(__name__)
+logger.setLevel(31)  # Avoid printing useless low-level logs
+
+
+def get_image_pairs_paths(inputs: list[Path] | Path) -> list[tuple[Path, Path]]:
+    """process input to produce a list of image pairs paths
+
+    Args:
+        inputs (list[Path] | Path): input path, which could be one of:
+            (1) two image paths
+            (2) dir with two images
+            (3) dir with dirs with image pairs
+            (4) txt file with two image paths per line
+
+    Returns:
+        list[tuple[Path, Path]]: list of pairs of image paths
+    """
+    if len(inputs) > 2:
+        raise ValueError(f"--input should be one or two paths, not {len(inputs)} paths like {inputs}")
+
+    if len(inputs) == 2:
+        # --input is two paths of images
+        if not inputs[0].is_file() or not inputs[1].is_file():
+            raise ValueError(f"If --input is two paths, it should be two images, not {inputs}")
+        return [inputs]
+
+    assert len(inputs) == 1
+    inputs = Path(inputs[0])
+
+    if not inputs.exists():
+        raise ValueError(f"{inputs} does not exist")
+
+    if inputs.is_file():
+        # --input is a file with pairs of images paths
+        with open(inputs) as file:
+            lines = file.read().splitlines()
+        pairs_of_paths = [line.strip().split(" ") for line in lines]
+        for pair in pairs_of_paths:
+            if len(pair) != 2:
+                raise ValueError(f"{pair} should be a pair of paths")
+        return [(Path(path0.strip()), Path(path1.strip())) for path0, path1 in pairs_of_paths]
+    elif inputs.is_dir():
+        inner_files = sorted(Path(inputs).glob("*"))
+        if len(inner_files) == 2 and inner_files[0].is_file() and inner_files[1].is_file():
+            # --input is a dir with a pair of images
+            return [inner_files]
+        else:
+            # --input is a dir of subdirs, where each subdir has a pair of images
+            pairs_of_paths = [list(pair_dir.glob("*")) for pair_dir in inner_files]
+            for pair in pairs_of_paths:
+                if len(pair) != 2:
+                    raise ValueError(f"{pair} should be a pair of paths")
+            return pairs_of_paths
+    else:
+        print(f"Could not parse inputs: {inputs}")
+
+
+def to_numpy(x: torch.Tensor | np.ndarray | dict | list) -> np.ndarray:
+    """convert item or container of items to numpy
+
+    Args:
+        x (torch.Tensor | np.ndarray | dict | list): input
+
+    Returns:
+        np.ndarray: numpy array of input
+    """
+    if isinstance(x, list):
+        return np.array([to_numpy(i) for i in x])
+    if isinstance(x, dict):
+        for k, v in x.items():
+            x[k] = to_numpy(v)
+    if isinstance(x, torch.Tensor):
+        return x.cpu().numpy()
+    if isinstance(x, np.ndarray):
+        return x
+    if x is None:
+        return
+    raise NotImplementedError(f"to_numpy not implemented for data type {type(x)}")
+
+
+def to_tensor(x: np.ndarray | torch.Tensor, device: str = None) -> torch.Tensor:
+    """Convert to tensor and place on device
+
+    Args:
+        x (np.ndarray | torch.Tensor): item to convert to tensor
+        device (str, optional): device to place tensor on. Defaults to None.
+
+    Returns:
+        torch.Tensor: tensor with data from `x` on device `device`
+    """
+    if isinstance(x, torch.Tensor):
+        pass
+    elif isinstance(x, np.ndarray):
+        x = torch.from_numpy(x)
+
+    if device is not None:
+        return x.to(device)
+    else:
+        return x
+
+
+def to_device(data: torch.Tensor | dict | list, device: str = "cuda"):
+    """Recursively move tensors in nested data structures to `device`."""
+    if isinstance(data, torch.Tensor):
+        return data.to(device)
+    elif isinstance(data, dict):
+        return {k: to_device(v, device) for k, v in data.items()}
+    elif isinstance(data, list):
+        return [to_device(item, device) for item in data]
+    else:
+        return data
+
+
+def to_normalized_coords(pts: np.ndarray | torch.Tensor, height: int, width: int):
+    """normalize kpt coords from px space to [0,1]
+    Assumes pts are in x, y order in array/tensor shape (N, 2)
+
+    Args:
+        pts (np.ndarray | torch.Tensor): array of kpts, must be shape (N, 2)
+        height (int): height of img
+        width (int): width of img
+
+    Returns:
+        np.array: kpts in normalized [0,1] coords
+    """
+    # normalize kpt coords from px space to [0,1]
+    # assume pts are in x,y order
+    assert pts.shape[-1] == 2, f"input to `to_normalized_coords` should be shape (N, 2), input is shape {pts.shape}"
+    pts = to_numpy(pts).astype(float)
+    pts[:, 0] /= width
+    pts[:, 1] /= height
+
+    return pts
+
+
+def to_px_coords(pts: np.ndarray | torch.Tensor, height: int, width: int) -> np.ndarray:
+    """unnormalized kpt coords from [0,1] to px space
+    Assumes pts are in x, y order
+
+    Args:
+        pts (np.ndarray | torch.Tensor): array of kpts, must be shape (N, 2)
+        height (int): height of img
+        width (int): width of img
+
+    Returns:
+        np.array: kpts in normalized [0,1] coords
+    """
+    assert pts.shape[-1] == 2, f"input to `to_px_coords` should be shape (N, 2), input is shape {pts.shape}"
+    pts = to_numpy(pts)
+    pts[:, 0] *= width
+    pts[:, 1] *= height
+
+    return pts
+
+
+def pad_images_to_same_shape(img0: torch.Tensor, img1: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pad two image tensors to the same spatial dimensions (right/bottom zero-padding)."""
+    h0, w0 = img0.shape[-2:]
+    h1, w1 = img1.shape[-2:]
+    if h0 == h1 and w0 == w1:
+        return img0, img1
+    max_h = max(h0, h1)
+    max_w = max(w0, w1)
+    img0 = torch.nn.functional.pad(img0, (0, max_w - w0, 0, max_h - h0))
+    img1 = torch.nn.functional.pad(img1, (0, max_w - w1, 0, max_h - h1))
+    return img0, img1
+
+
+def resize_to_divisible(img: torch.Tensor, divisible_by: int = 14) -> torch.Tensor:
+    """Resize to be divisible by a factor. Useful for ViT based models.
+
+    Args:
+        img (torch.Tensor): img as tensor, in (*, H, W) order
+        divisible_by (int, optional): factor to make sure img is divisible by. Defaults to 14.
+
+    Returns:
+        torch.Tensor: img tensor with divisible shape
+    """
+    h, w = img.shape[-2:]
+
+    divisible_h = round(h / divisible_by) * divisible_by
+    divisible_w = round(w / divisible_by) * divisible_by
+    img = tfm.functional.resize(img, [divisible_h, divisible_w], antialias=True)
+
+    return img
+
+
+def lower_config(yacs_cfg: CN) -> dict:
+    """Convert yacs config to lower-case dict recursively."""
+    if not isinstance(yacs_cfg, CN):
+        return yacs_cfg
+    return {k.lower(): lower_config(v) for k, v in yacs_cfg.items()}
+
+
+def load_module(module_name: str, module_path: Path | str) -> None:
+    """Load module from `module_path` into the interpreter with the namespace given by module_name.
+
+    Note that `module_path` is usually the path to an `__init__.py` file.
+
+    Args:
+        module_name (str): module name (will be used to import from later, as in `from module_name import my_function`)
+        module_path (Path | str): path to module (usually an __init__.py file)
+    """
+    # load gluefactory into namespace
+    # module_name = 'gluefactory'
+    spec = importlib.util.spec_from_file_location(module_name, module_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+
+
+def add_to_path(path: str | Path, insert: int | None = None) -> None:
+    """Add input path to sys.path ($PATH), allowing for imports from the specified path
+
+    Args:
+        path (str | Path): path to add to sys.path
+        insert (int | None, optional): insert location / order. Defaults to None, which inserts at end of sys.path
+    """
+    path = str(path)
+    if path in sys.path:
+        sys.path.remove(path)
+    if insert is None:
+        sys.path.append(path)
+    else:
+        sys.path.insert(insert, path)
+
+
+def get_default_device() -> str:
+    """get best available device for torch: cuda, mps (mac), else cpu
+
+    Returns:
+        str: best available device as str
+    """
+    # default device is cpu
+    device = "cpu"
+
+    # test for mac device (darwin) and mps availability
+    if sys.platform == "darwin" and torch.backends.mps.is_available():
+        device = "mps"
+
+    # check cuda availability
+    elif torch.cuda.is_available():
+        device = "cuda"
+
+    return device
+
+
+def flow_to_matches(
+    flow: np.ndarray,
+    covisibility: np.ndarray,
+    num_samples: int = 1000,
+    min_confidence: float = 0.0,
+    method: str = "probabilistic",
+    rng: np.random.RandomState | np.random.Generator = None,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Convert a dense optical flow + covisibility map to sparse keypoint matches.
+
+    Args:
+        flow (np.ndarray): shape (2, H, W) or (H, W, 2). Interpreted as (dx, dy) per pixel.
+        covisibility (np.ndarray): shape (H, W) with confidence in [0, 1] (or any non-negative scores).
+        num_samples (int, optional): max number of matches to return. Defaults to 1000.
+        min_confidence (float, optional): ignore pixels with covisibility <= min_confidence. Defaults to 0.0.
+        method (str, optional): sampling method, one of "probabilistic", "topk", or "grid". Defaults to "probabilistic".
+        rng (np.random.RandomState | np.random.Generator, optional): for reproducibility. Defaults to None.
+
+    Returns:
+        tuple: (matches0, matches1, confidences) where:
+            - matches0 (np.ndarray): (N, 2) source keypoints as (x, y) (float32)
+            - matches1 (np.ndarray): (N, 2) target keypoints as (x, y) = source + flow (float32)
+            - confidences (np.ndarray): (N,) covisibility/confidence values (float32)
+    """
+    if rng is None:
+        rng = np.random
+
+    # Normalize flow shape -> (2, H, W)
+    if flow.ndim == 3 and flow.shape[2] == 2:
+        flow_xy = flow.transpose(2, 0, 1)  # (2, H, W)
+    elif flow.ndim == 3 and flow.shape[0] == 2:
+        flow_xy = flow
+    else:
+        raise ValueError("flow must have shape (2,H,W) or (H,W,2)")
+
+    H, W = covisibility.shape
+    if flow_xy.shape[1:] != (H, W):
+        raise ValueError(
+            f"flow and covisibility spatial dims mismatch: flow {flow_xy.shape[1:]}, covisibility {covisibility.shape}"
+        )
+
+    # Flatten grids
+    xs = np.arange(W, dtype=np.float32)
+    ys = np.arange(H, dtype=np.float32)
+    gx, gy = np.meshgrid(xs, ys)  # gx: (H,W) x coords, gy: (H,W) y coords
+
+    flat_conf = covisibility.ravel().astype(np.float64)
+    valid_mask = flat_conf > min_confidence
+    if valid_mask.sum() == 0:
+        return np.zeros((0, 2), dtype=np.float32), np.zeros((0, 2), dtype=np.float32), np.zeros((0,), dtype=np.float32)
+
+    valid_idxs = np.nonzero(valid_mask)[0]
+
+    if method == "probabilistic":
+        scores = flat_conf[valid_mask].astype(np.float64)
+        # avoid degenerate all-zero
+        if scores.sum() <= 0:
+            probs = None
+        else:
+            probs = scores / scores.sum()
+        k = min(num_samples, len(valid_idxs))
+        # if probs is None or degenerate, fallback to uniform
+        chosen = rng.choice(valid_idxs, size=k, replace=False, p=probs)
+    elif method == "topk":
+        k = min(num_samples, len(valid_idxs))
+        topk_local = np.argsort(-flat_conf[valid_mask])[:k]
+        chosen = valid_idxs[topk_local]
+    elif method == "grid":
+        # choose roughly sqrt grid
+        n = max(1, int(np.sqrt(num_samples)))
+        xs_idx = np.linspace(0, W - 1, n, dtype=int)
+        ys_idx = np.linspace(0, H - 1, int(np.ceil(num_samples / n)), dtype=int)
+        gx_idx, gy_idx = np.meshgrid(xs_idx, ys_idx)
+        chosen_coords = np.stack([gy_idx.ravel(), gx_idx.ravel()], axis=1)
+        chosen_coords = chosen_coords[:num_samples]
+        chosen = chosen_coords[:, 0] * W + chosen_coords[:, 1]
+        # mask by min_confidence
+        keep = flat_conf[chosen] > min_confidence
+        chosen = chosen[keep]
+    else:
+        raise ValueError("method must be one of 'probabilistic','topk','grid'")
+
+    # gather coordinates and flows
+    gy_flat = gy.ravel().astype(np.float32)
+    gx_flat = gx.ravel().astype(np.float32)
+    dx_flat = flow_xy[0].ravel().astype(np.float32)
+    dy_flat = flow_xy[1].ravel().astype(np.float32)
+
+    src_x = gx_flat[chosen]
+    src_y = gy_flat[chosen]
+    dx = dx_flat[chosen]
+    dy = dy_flat[chosen]
+    confs = flat_conf[chosen].astype(np.float32)
+
+    matches0 = np.stack([src_x, src_y], axis=1).astype(np.float32)
+    matches1 = (matches0 + np.stack([dx, dy], axis=1)).astype(np.float32)
+
+    return matches0, matches1, confs
+
+
+def _load_image(path: str | Path, resize: int | tuple = None, rot_angle: float = 0) -> torch.Tensor:
+    """load image from filesystem and return as tensor. Optionally rotate and resize.
+
+    Args:
+        path (str | Path): path to image on filesystem
+        resize (int | tuple, optional): size to resize img, either single value for square resize or tuple of (H, W). Defaults to None.
+        rot_angle (float, optional): CCW rotation angle in degrees. Defaults to 0.
+
+    Returns:
+        torch.Tensor: image as tensor (C x H x W)
+    """
+    if isinstance(resize, int):
+        resize = (resize, resize)
+    img = tfm.ToTensor()(Image.open(path).convert("RGB"))
+    if resize is not None:
+        img = tfm.Resize(resize, antialias=True)(img)
+    img = tfm.functional.rotate(img, rot_angle)
+    return img
+
+
+def to_tensor_image(img):
+    if isinstance(img, (str, Path)):
+        img = _load_image(img)
+    elif isinstance(img, Image.Image):
+        img = tfm.ToTensor()(img.convert("RGB"))
+    elif isinstance(img, np.ndarray):
+        img = torch.from_numpy(img)
+    assert isinstance(img, torch.Tensor), "img should be a torch.Tensor, a path, or a PIL Image"
+    assert img.ndim == 3 and img.shape[0] == 3, f"img should have shape (3, H, W), got {img.shape}"
+    # Small tolerance of 0.2 because images after bicubic resizing can slightly exceed the [0, 1] range
+    # This is expected, not a bug, see https://github.com/opencv/opencv/issues/7195
+    assert -0.2 <= img.min() and img.max() <= 1.2, f"img should be in [0, 1], got [{img.min()}, {img.max()}]"
+    return img