vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/gim/dkm/models/encoders.py

@@ -0,0 +1,148 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as tvm
+
+class ResNet18(nn.Module):
+    def __init__(self, pretrained=False) -> None:
+        super().__init__()
+        self.net = tvm.resnet18(pretrained=pretrained)
+    def forward(self, x):
+        self = self.net
+        x1 = x
+        x = self.conv1(x1)
+        x = self.bn1(x)
+        x2 = self.relu(x)
+        x = self.maxpool(x2)
+        x4 = self.layer1(x)
+        x8 = self.layer2(x4)
+        x16 = self.layer3(x8)
+        x32 = self.layer4(x16)
+        return {32:x32,16:x16,8:x8,4:x4,2:x2,1:x1}
+
+    def train(self, mode=True):
+        super().train(mode)
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm2d):
+                m.eval()
+            pass
+
+class ResNet50(nn.Module):
+    def __init__(self, pretrained=False, high_res = False, weights = None, dilation = None, freeze_bn = True, anti_aliased = False) -> None:
+        super().__init__()
+        if dilation is None:
+            dilation = [False,False,False]
+        if anti_aliased:
+            pass
+        else:
+            if weights is not None:
+                self.net = tvm.resnet50(weights = weights,replace_stride_with_dilation=dilation)
+            else:
+                self.net = tvm.resnet50(pretrained=pretrained,replace_stride_with_dilation=dilation)
+
+        del self.net.fc
+        self.high_res = high_res
+        self.freeze_bn = freeze_bn
+    def forward(self, x):
+        net = self.net
+        feats = {1:x}
+        x = net.conv1(x)
+        x = net.bn1(x)
+        x = net.relu(x)
+        feats[2] = x 
+        x = net.maxpool(x)
+        x = net.layer1(x)
+        feats[4] = x 
+        x = net.layer2(x)
+        feats[8] = x  
+        x = net.layer3(x)
+        feats[16] = x
+        x = net.layer4(x)
+        feats[32] = x
+        return feats
+
+    def train(self, mode=True):
+        super().train(mode)
+        if self.freeze_bn:
+            for m in self.modules():
+                if isinstance(m, nn.BatchNorm2d):
+                    m.eval()
+                pass
+
+
+
+
+class ResNet101(nn.Module):
+    def __init__(self, pretrained=False, high_res = False, weights = None) -> None:
+        super().__init__()
+        if weights is not None:
+            self.net = tvm.resnet101(weights = weights)
+        else:
+            self.net = tvm.resnet101(pretrained=pretrained)
+        self.high_res = high_res
+        self.scale_factor = 1 if not high_res else 1.5
+    def forward(self, x):
+        net = self.net
+        feats = {1:x}
+        sf = self.scale_factor
+        if self.high_res:
+            x = F.interpolate(x, scale_factor=sf, align_corners=False, mode="bicubic")
+        x = net.conv1(x)
+        x = net.bn1(x)
+        x = net.relu(x)
+        feats[2] = x if not self.high_res else F.interpolate(x,scale_factor=1/sf,align_corners=False, mode="bilinear")    
+        x = net.maxpool(x)
+        x = net.layer1(x)
+        feats[4] = x if not self.high_res else F.interpolate(x,scale_factor=1/sf,align_corners=False, mode="bilinear")    
+        x = net.layer2(x)
+        feats[8] = x if not self.high_res else F.interpolate(x,scale_factor=1/sf,align_corners=False, mode="bilinear")    
+        x = net.layer3(x)
+        feats[16] = x if not self.high_res else F.interpolate(x,scale_factor=1/sf,align_corners=False, mode="bilinear")    
+        x = net.layer4(x)
+        feats[32] = x if not self.high_res else F.interpolate(x,scale_factor=1/sf,align_corners=False, mode="bilinear")    
+        return feats
+
+    def train(self, mode=True):
+        super().train(mode)
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm2d):
+                m.eval()
+            pass
+
+
+class WideResNet50(nn.Module):
+    def __init__(self, pretrained=False, high_res = False, weights = None) -> None:
+        super().__init__()
+        if weights is not None:
+            self.net = tvm.wide_resnet50_2(weights = weights)
+        else:
+            self.net = tvm.wide_resnet50_2(pretrained=pretrained)
+        self.high_res = high_res
+        self.scale_factor = 1 if not high_res else 1.5
+    def forward(self, x):
+        net = self.net
+        feats = {1:x}
+        sf = self.scale_factor
+        if self.high_res:
+            x = F.interpolate(x, scale_factor=sf, align_corners=False, mode="bicubic")
+        x = net.conv1(x)
+        x = net.bn1(x)
+        x = net.relu(x)
+        feats[2] = x if not self.high_res else F.interpolate(x,scale_factor=1/sf,align_corners=False, mode="bilinear")    
+        x = net.maxpool(x)
+        x = net.layer1(x)
+        feats[4] = x if not self.high_res else F.interpolate(x,scale_factor=1/sf,align_corners=False, mode="bilinear")    
+        x = net.layer2(x)
+        feats[8] = x if not self.high_res else F.interpolate(x,scale_factor=1/sf,align_corners=False, mode="bilinear")    
+        x = net.layer3(x)
+        feats[16] = x if not self.high_res else F.interpolate(x,scale_factor=1/sf,align_corners=False, mode="bilinear")    
+        x = net.layer4(x)
+        feats[32] = x if not self.high_res else F.interpolate(x,scale_factor=1/sf,align_corners=False, mode="bilinear")    
+        return feats
+
+    def train(self, mode=True):
+        super().train(mode)
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm2d):
+                m.eval()
+            pass

vismatch/third_party/gim/dkm/models/model_zoo/DKMv3.py

@@ -0,0 +1,148 @@
+import torch
+
+from torch import nn
+from dkm.models.dkm import *
+from dkm.models.encoders import *
+
+
+def DKMv3(weights, h, w, symmetric = True, sample_mode= "threshold_balanced", **kwargs):
+    gp_dim = 256
+    dfn_dim = 384
+    feat_dim = 256
+    coordinate_decoder = DFN(
+        internal_dim=dfn_dim,
+        feat_input_modules=nn.ModuleDict(
+            {
+                "32": nn.Conv2d(512, feat_dim, 1, 1),
+                "16": nn.Conv2d(512, feat_dim, 1, 1),
+            }
+        ),
+        pred_input_modules=nn.ModuleDict(
+            {
+                "32": nn.Identity(),
+                "16": nn.Identity(),
+            }
+        ),
+        rrb_d_dict=nn.ModuleDict(
+            {
+                "32": RRB(gp_dim + feat_dim, dfn_dim),
+                "16": RRB(gp_dim + feat_dim, dfn_dim),
+            }
+        ),
+        cab_dict=nn.ModuleDict(
+            {
+                "32": CAB(2 * dfn_dim, dfn_dim),
+                "16": CAB(2 * dfn_dim, dfn_dim),
+            }
+        ),
+        rrb_u_dict=nn.ModuleDict(
+            {
+                "32": RRB(dfn_dim, dfn_dim),
+                "16": RRB(dfn_dim, dfn_dim),
+            }
+        ),
+        terminal_module=nn.ModuleDict(
+            {
+                "32": nn.Conv2d(dfn_dim, 3, 1, 1, 0),
+                "16": nn.Conv2d(dfn_dim, 3, 1, 1, 0),
+            }
+        ),
+    )
+    dw = True
+    hidden_blocks = 8
+    kernel_size = 5
+    displacement_emb = "linear"
+    conv_refiner = nn.ModuleDict(
+        {
+            "16": ConvRefiner(
+                2 * 512+128+(2*7+1)**2,
+                2 * 512+128+(2*7+1)**2,
+                3,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=128,
+                local_corr_radius = 7,
+                corr_in_other = True,
+            ),
+            "8": ConvRefiner(
+                2 * 512+64+(2*3+1)**2,
+                2 * 512+64+(2*3+1)**2,
+                3,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=64,
+                local_corr_radius = 3,
+                corr_in_other = True,
+            ),
+            "4": ConvRefiner(
+                2 * 256+32+(2*2+1)**2,
+                2 * 256+32+(2*2+1)**2,
+                3,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=32,
+                local_corr_radius = 2,
+                corr_in_other = True,
+            ),
+            "2": ConvRefiner(
+                2 * 64+16,
+                128+16,
+                3,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=16,
+            ),
+            "1": ConvRefiner(
+                2 * 3+6,
+                24,
+                3,
+                kernel_size=kernel_size,
+                dw=dw,
+                hidden_blocks=hidden_blocks,
+                displacement_emb=displacement_emb,
+                displacement_emb_dim=6,
+            ),
+        }
+    )
+    kernel_temperature = 0.2
+    learn_temperature = False
+    no_cov = True
+    kernel = CosKernel
+    only_attention = False
+    basis = "fourier"
+    gp32 = GP(
+        kernel,
+        T=kernel_temperature,
+        learn_temperature=learn_temperature,
+        only_attention=only_attention,
+        gp_dim=gp_dim,
+        basis=basis,
+        no_cov=no_cov,
+    )
+    gp16 = GP(
+        kernel,
+        T=kernel_temperature,
+        learn_temperature=learn_temperature,
+        only_attention=only_attention,
+        gp_dim=gp_dim,
+        basis=basis,
+        no_cov=no_cov,
+    )
+    gps = nn.ModuleDict({"32": gp32, "16": gp16})
+    proj = nn.ModuleDict(
+        {"16": nn.Conv2d(1024, 512, 1, 1), "32": nn.Conv2d(2048, 512, 1, 1)}
+    )
+    decoder = Decoder(coordinate_decoder, gps, proj, conv_refiner, detach=True)
+
+    encoder = ResNet50(pretrained = False, high_res = False, freeze_bn=False)
+    matcher = RegressionMatcher(encoder, decoder, h=h, w=w, name = "DKMv3", sample_mode=sample_mode, symmetric = symmetric, **kwargs)
+    # res = matcher.load_state_dict(weights)
+    return matcher

vismatch/third_party/gim/dkm/models/model_zoo/__init__.py

@@ -0,0 +1,39 @@
+weight_urls = {
+    "DKMv3": {
+        "outdoor": "https://github.com/Parskatt/storage/releases/download/dkmv3/DKMv3_outdoor.pth",
+        "indoor": "https://github.com/Parskatt/storage/releases/download/dkmv3/DKMv3_indoor.pth",
+    },
+}
+import torch
+from .DKMv3 import DKMv3
+
+
+def DKMv3_outdoor(path_to_weights = None, device=None):
+    """
+    Loads DKMv3 outdoor weights, uses internal resolution of (540, 720) by default
+    resolution can be changed by setting model.h_resized, model.w_resized later.
+    Additionally upsamples preds to fixed resolution of (864, 1152),
+    can be turned off by model.upsample_preds = False
+    """
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if path_to_weights is not None:
+        weights = torch.load(path_to_weights, map_location=device)
+    else:
+        weights = torch.hub.load_state_dict_from_url(weight_urls["DKMv3"]["outdoor"],
+                                                     map_location=device)
+    return DKMv3(weights, 540, 720, upsample_preds = True, device=device)
+
+def DKMv3_indoor(path_to_weights = None, device=None):
+    """
+    Loads DKMv3 indoor weights, uses internal resolution of (480, 640) by default
+    Resolution can be changed by setting model.h_resized, model.w_resized later.
+    """
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if path_to_weights is not None:
+        weights = torch.load(path_to_weights, map_location=device)
+    else:
+        weights = torch.hub.load_state_dict_from_url(weight_urls["DKMv3"]["indoor"],
+                                                     map_location=device)
+    return DKMv3(weights, 480, 640, upsample_preds = False, device=device)

vismatch/third_party/gim/dkm/train/__init__.py

@@ -0,0 +1 @@
+from .train import train_k_epochs

vismatch/third_party/gim/dkm/train/train.py

@@ -0,0 +1,67 @@
+from tqdm import tqdm
+from dkm.utils.utils import to_cuda
+
+
+def train_step(train_batch, model, objective, optimizer, **kwargs):
+    optimizer.zero_grad()
+    out = model(train_batch)
+    l = objective(out, train_batch)
+    l.backward()
+    optimizer.step()
+    return {"train_out": out, "train_loss": l.item()}
+
+
+def train_k_steps(
+    n_0, k, dataloader, model, objective, optimizer, lr_scheduler, progress_bar=True
+):
+    for n in tqdm(range(n_0, n_0 + k), disable=not progress_bar):
+        batch = next(dataloader)
+        model.train(True)
+        batch = to_cuda(batch)
+        train_step(
+            train_batch=batch,
+            model=model,
+            objective=objective,
+            optimizer=optimizer,
+            lr_scheduler=lr_scheduler,
+            n=n,
+        )
+        lr_scheduler.step()
+
+
+def train_epoch(
+    dataloader=None,
+    model=None,
+    objective=None,
+    optimizer=None,
+    lr_scheduler=None,
+    epoch=None,
+):
+    model.train(True)
+    print(f"At epoch {epoch}")
+    for batch in tqdm(dataloader, mininterval=5.0):
+        batch = to_cuda(batch)
+        train_step(
+            train_batch=batch, model=model, objective=objective, optimizer=optimizer
+        )
+    lr_scheduler.step()
+    return {
+        "model": model,
+        "optimizer": optimizer,
+        "lr_scheduler": lr_scheduler,
+        "epoch": epoch,
+    }
+
+
+def train_k_epochs(
+    start_epoch, end_epoch, dataloader, model, objective, optimizer, lr_scheduler
+):
+    for epoch in range(start_epoch, end_epoch + 1):
+        train_epoch(
+            dataloader=dataloader,
+            model=model,
+            objective=objective,
+            optimizer=optimizer,
+            lr_scheduler=lr_scheduler,
+            epoch=epoch,
+        )

vismatch/third_party/gim/dkm/utils/__init__.py

@@ -0,0 +1,13 @@
+from .utils import (
+    pose_auc,
+    get_pose,
+    compute_relative_pose,
+    compute_pose_error,
+    estimate_pose,
+    rotate_intrinsic,
+    get_tuple_transform_ops,
+    get_depth_tuple_transform_ops,
+    warp_kpts,
+    numpy_to_pil,
+    tensor_to_pil,
+)

vismatch/third_party/gim/dkm/utils/kde.py

@@ -0,0 +1,26 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+def fast_kde(x, std = 0.1, kernel_size = 9, dilation = 3, padding = 9//2, stride = 1):
+    raise NotImplementedError("WIP, use at your own risk.")
+    # Note: when doing symmetric matching this might not be very exact, since we only check neighbours on the grid
+    x = x.permute(0,3,1,2)
+    B,C,H,W = x.shape
+    K = kernel_size ** 2
+    unfolded_x = F.unfold(x,kernel_size=kernel_size, dilation = dilation, padding = padding, stride = stride).reshape(B, C, K, H, W)
+    scores = (-(unfolded_x - x[:,:,None]).sum(dim=1)**2/(2*std**2)).exp()
+    density = scores.sum(dim=1)
+    return density
+
+
+def kde(x, std = 0.1, device=None):
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if isinstance(x, np.ndarray):
+        x = torch.from_numpy(x)
+    # use a gaussian kernel to estimate density
+    x = x.to(device)
+    scores = (-torch.cdist(x,x)**2/(2*std**2)).exp()
+    density = scores.sum(dim=-1)
+    return density

vismatch/third_party/gim/dkm/utils/local_correlation.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn.functional as F
+
+
+def local_correlation(
+    feature0,
+    feature1,
+    local_radius,
+    padding_mode="zeros",
+    flow = None
+):
+    device = feature0.device
+    b, c, h, w = feature0.size()
+    if flow is None:
+        # If flow is None, assume feature0 and feature1 are aligned
+        coords = torch.meshgrid(
+                (
+                    torch.linspace(-1 + 1 / h, 1 - 1 / h, h, device=device),
+                    torch.linspace(-1 + 1 / w, 1 - 1 / w, w, device=device),
+                ))
+        coords = torch.stack((coords[1], coords[0]), dim=-1)[
+            None
+        ].expand(b, h, w, 2)
+    else:
+        coords = flow.permute(0,2,3,1) # If using flow, sample around flow target.
+    r = local_radius
+    local_window = torch.meshgrid(
+                (
+                    torch.linspace(-2*local_radius/h, 2*local_radius/h, 2*r+1, device=device),
+                    torch.linspace(-2*local_radius/w, 2*local_radius/w, 2*r+1, device=device),
+                ))
+    local_window = torch.stack((local_window[1], local_window[0]), dim=-1)[
+            None
+        ].expand(b, 2*r+1, 2*r+1, 2).reshape(b, (2*r+1)**2, 2)
+    coords = (coords[:,:,:,None]+local_window[:,None,None]).reshape(b,h,w*(2*r+1)**2,2)
+    window_feature = F.grid_sample(
+        feature1, coords, padding_mode=padding_mode, align_corners=False
+    )[...,None].reshape(b,c,h,w,(2*r+1)**2)
+    corr = torch.einsum("bchw, bchwk -> bkhw", feature0, window_feature)/(c**.5)
+    return corr

vismatch/third_party/gim/dkm/utils/transforms.py

@@ -0,0 +1,104 @@
+from typing import Dict
+import numpy as np
+import torch
+import kornia.augmentation as K
+from kornia.geometry.transform import warp_perspective
+
+# Adapted from Kornia
+class GeometricSequential:
+    def __init__(self, *transforms, align_corners=True) -> None:
+        self.transforms = transforms
+        self.align_corners = align_corners
+
+    def __call__(self, x, mode="bilinear"):
+        b, c, h, w = x.shape
+        M = torch.eye(3, device=x.device)[None].expand(b, 3, 3)
+        for t in self.transforms:
+            if np.random.rand() < t.p:
+                M = M.matmul(
+                    t.compute_transformation(x, t.generate_parameters((b, c, h, w)))
+                )
+        return (
+            warp_perspective(
+                x, M, dsize=(h, w), mode=mode, align_corners=self.align_corners
+            ),
+            M,
+        )
+
+    def apply_transform(self, x, M, mode="bilinear"):
+        b, c, h, w = x.shape
+        return warp_perspective(
+            x, M, dsize=(h, w), align_corners=self.align_corners, mode=mode
+        )
+
+
+class RandomPerspective(K.RandomPerspective):
+    def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
+        distortion_scale = torch.as_tensor(
+            self.distortion_scale, device=self._device, dtype=self._dtype
+        )
+        return self.random_perspective_generator(
+            batch_shape[0],
+            batch_shape[-2],
+            batch_shape[-1],
+            distortion_scale,
+            self.same_on_batch,
+            self.device,
+            self.dtype,
+        )
+
+    def random_perspective_generator(
+        self,
+        batch_size: int,
+        height: int,
+        width: int,
+        distortion_scale: torch.Tensor,
+        same_on_batch: bool = False,
+        device: torch.device = torch.device("cpu"),
+        dtype: torch.dtype = torch.float32,
+    ) -> Dict[str, torch.Tensor]:
+        r"""Get parameters for ``perspective`` for a random perspective transform.
+
+        Args:
+            batch_size (int): the tensor batch size.
+            height (int) : height of the image.
+            width (int): width of the image.
+            distortion_scale (torch.Tensor): it controls the degree of distortion and ranges from 0 to 1.
+            same_on_batch (bool): apply the same transformation across the batch. Default: False.
+            device (torch.device): the device on which the random numbers will be generated. Default: cpu.
+            dtype (torch.dtype): the data type of the generated random numbers. Default: float32.
+
+        Returns:
+            params Dict[str, torch.Tensor]: parameters to be passed for transformation.
+                - start_points (torch.Tensor): element-wise perspective source areas with a shape of (B, 4, 2).
+                - end_points (torch.Tensor): element-wise perspective target areas with a shape of (B, 4, 2).
+
+        Note:
+            The generated random numbers are not reproducible across different devices and dtypes.
+        """
+        if not (distortion_scale.dim() == 0 and 0 <= distortion_scale <= 1):
+            raise AssertionError(
+                f"'distortion_scale' must be a scalar within [0, 1]. Got {distortion_scale}."
+            )
+        if not (
+            type(height) is int and height > 0 and type(width) is int and width > 0
+        ):
+            raise AssertionError(
+                f"'height' and 'width' must be integers. Got {height}, {width}."
+            )
+
+        start_points: torch.Tensor = torch.tensor(
+            [[[0.0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]]],
+            device=distortion_scale.device,
+            dtype=distortion_scale.dtype,
+        ).expand(batch_size, -1, -1)
+
+        # generate random offset not larger than half of the image
+        fx = distortion_scale * width / 2
+        fy = distortion_scale * height / 2
+
+        factor = torch.stack([fx, fy], dim=0).view(-1, 1, 2)
+        offset = (torch.rand_like(start_points) - 0.5) * 2
+        end_points = start_points + factor * offset
+
+        return dict(start_points=start_points, end_points=end_points)