vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/src/loftr/utils/geometry.py

@@ -0,0 +1,298 @@
+import torch
+from src.utils.homography_utils import warp_points_torch
+
+def get_unique_indices(input_tensor):
+    if input_tensor.shape[0] > 1:
+        unique, inverse = torch.unique(input_tensor, sorted=True, return_inverse=True, dim=0)
+        perm = torch.arange(inverse.size(0), dtype=inverse.dtype, device=inverse.device)
+        inverse, perm = inverse.flip([0]), perm.flip([0])
+        perm = inverse.new_empty(unique.size(0)).scatter_(0, inverse, perm)
+    else:
+        perm = torch.zeros((input_tensor.shape[0],), dtype=torch.long, device=input_tensor.device) 
+    return perm
+
+
+@torch.no_grad()
+def warp_kpts(kpts0, depth0, depth1, T_0to1, K0, K1, consistency_thr=0.2, cycle_proj_distance_thr=3.0):
+    """ Warp kpts0 from I0 to I1 with depth, K and Rt
+    Also check covisibility and depth consistency.
+    Depth is consistent if relative error < 0.2 (hard-coded).
+    
+    Args:
+        kpts0 (torch.Tensor): [N, L, 2] - <x, y>,
+        depth0 (torch.Tensor): [N, H, W],
+        depth1 (torch.Tensor): [N, H, W],
+        T_0to1 (torch.Tensor): [N, 3, 4],
+        K0 (torch.Tensor): [N, 3, 3],
+        K1 (torch.Tensor): [N, 3, 3],
+    Returns:
+        calculable_mask (torch.Tensor): [N, L]
+        warped_keypoints0 (torch.Tensor): [N, L, 2] <x0_hat, y1_hat>
+    """
+    kpts0_long = kpts0.round().long()
+
+    # Sample depth, get calculable_mask on depth != 0
+    kpts0_depth = torch.stack(
+        [depth0[i, kpts0_long[i, :, 1], kpts0_long[i, :, 0]] for i in range(kpts0.shape[0])], dim=0
+    )  # (N, L)
+    nonzero_mask = kpts0_depth != 0
+
+    # Unproject
+    kpts0_h = torch.cat([kpts0, torch.ones_like(kpts0[:, :, [0]])], dim=-1) * kpts0_depth[..., None]  # (N, L, 3)
+    kpts0_cam = K0.inverse() @ kpts0_h.transpose(2, 1)  # (N, 3, L)
+
+    # Rigid Transform
+    w_kpts0_cam = T_0to1[:, :3, :3] @ kpts0_cam + T_0to1[:, :3, [3]]    # (N, 3, L)
+    w_kpts0_depth_computed = w_kpts0_cam[:, 2, :]
+
+    # Project
+    w_kpts0_h = (K1 @ w_kpts0_cam).transpose(2, 1)  # (N, L, 3)
+    w_kpts0 = w_kpts0_h[:, :, :2] / (w_kpts0_h[:, :, [2]] + 1e-4)  # (N, L, 2), +1e-4 to avoid zero depth
+
+    # Covisible Check
+    h, w = depth1.shape[1:3]
+    covisible_mask = (w_kpts0[:, :, 0] > 0) * (w_kpts0[:, :, 0] < w-1) * \
+        (w_kpts0[:, :, 1] > 0) * (w_kpts0[:, :, 1] < h-1)
+    w_kpts0_long = w_kpts0.long()
+    w_kpts0_long[~covisible_mask, :] = 0
+
+    w_kpts0_depth = torch.stack(
+        [depth1[i, w_kpts0_long[i, :, 1], w_kpts0_long[i, :, 0]] for i in range(w_kpts0_long.shape[0])], dim=0
+    )  # (N, L)
+    consistent_mask = ((w_kpts0_depth - w_kpts0_depth_computed) / w_kpts0_depth).abs() < consistency_thr
+
+    # Cycle Consistency Check
+    dst_pts_h = torch.cat([w_kpts0, torch.ones_like(w_kpts0[..., [0]], device=w_kpts0.device)], dim=-1) * w_kpts0_depth[..., None] # B * N_dst * N_pts * 3
+    dst_pts_cam = K1.inverse() @ dst_pts_h.transpose(2, 1) # (N, 3, L)
+    dst_pose = T_0to1.inverse()
+    world_points_cycle_back = dst_pose[:, :3, :3] @ dst_pts_cam + dst_pose[:, :3, [3]]
+    src_warp_back_h = (K0 @ world_points_cycle_back).transpose(2, 1) # (N, L, 3)
+    src_back_proj_pts = src_warp_back_h[..., :2] / (src_warp_back_h[..., [2]] + 1e-4)
+    cycle_reproj_distance_mask = torch.linalg.norm(src_back_proj_pts - kpts0[:, None], dim=-1) < cycle_proj_distance_thr
+
+    valid_mask = nonzero_mask * covisible_mask * consistent_mask * cycle_reproj_distance_mask
+
+    return valid_mask, w_kpts0
+
+@torch.no_grad()
+def warp_kpts_by_sparse_gt_matches_batches(kpts0, gt_matches, dist_thr):
+    B, n_pts = kpts0.shape[0], kpts0.shape[1]
+    if n_pts > 20 * 10000:
+        all_kpts_valid_mask, all_kpts_warpped = [], []
+        for b_id in range(B):
+            kpts_valid_mask, kpts_warpped = warp_kpts_by_sparse_gt_matches(kpts0[[b_id]], gt_matches[[b_id]], dist_thr[[b_id]])
+            all_kpts_valid_mask.append(kpts_valid_mask)
+            all_kpts_warpped.append(kpts_warpped)
+        return torch.cat(all_kpts_valid_mask, dim=0), torch.cat(all_kpts_warpped, dim=0)
+    else:
+        return warp_kpts_by_sparse_gt_matches(kpts0, gt_matches, dist_thr)
+
+@torch.no_grad()
+def warp_kpts_by_sparse_gt_matches(kpts0, gt_matches, dist_thr):
+    kpts_warpped = torch.zeros_like(kpts0)
+    kpts_valid_mask = torch.zeros_like(kpts0[..., 0], dtype=torch.bool)
+    gt_matches_non_padding_mask = gt_matches.sum(-1) > 0
+
+    dist_matrix = torch.cdist(kpts0, gt_matches[..., :2]) # B * N * M
+    if dist_thr is not None:
+        mask = dist_matrix < dist_thr[:, None, None]
+    else:
+        mask = torch.ones_like(dist_matrix, dtype=torch.bool)
+    # Mutual-Nearest check:
+    mask = mask \
+        * (dist_matrix == dist_matrix.min(dim=2, keepdim=True)[0]) \
+        * (dist_matrix == dist_matrix.min(dim=1, keepdim=True)[0])
+
+    mask_v, all_j_ids = mask.max(dim=2)
+    b_ids, i_ids = torch.where(mask_v)
+    j_ids = all_j_ids[b_ids, i_ids]
+
+    j_uq_indices = get_unique_indices(torch.stack([b_ids, j_ids], dim=-1))
+    b_ids, i_ids, j_ids = map(lambda x: x[j_uq_indices], [b_ids, i_ids, j_ids])
+
+    i_uq_indices = get_unique_indices(torch.stack([b_ids, i_ids], dim=-1))
+    b_ids, i_ids, j_ids = map(lambda x: x[i_uq_indices], [b_ids, i_ids, j_ids])
+
+    kpts_valid_mask[b_ids, i_ids] = gt_matches_non_padding_mask[b_ids, j_ids]
+    kpts_warpped[b_ids, i_ids] = gt_matches[..., 2:][b_ids, j_ids]
+
+    return kpts_valid_mask, kpts_warpped
+
+@torch.no_grad()
+def warp_kpts_by_sparse_gt_matches_fine_chunks(kpts0, gt_matches, dist_thr):
+    B, n_pts = kpts0.shape[0], kpts0.shape[1]
+    chunk_n = 500
+    all_kpts_valid_mask, all_kpts_warpped = [], []
+    for b_id in range(0, B, chunk_n):
+        kpts_valid_mask, kpts_warpped = warp_kpts_by_sparse_gt_matches_fine(kpts0[b_id : b_id+chunk_n], gt_matches, dist_thr)
+        all_kpts_valid_mask.append(kpts_valid_mask)
+        all_kpts_warpped.append(kpts_warpped)
+    return torch.cat(all_kpts_valid_mask, dim=0), torch.cat(all_kpts_warpped, dim=0)
+
+@torch.no_grad()
+def warp_kpts_by_sparse_gt_matches_fine(kpts0, gt_matches, dist_thr):
+    """
+    Only support single batch
+    Input:
+    kpts0: N * ww * 2
+    gt_matches: M * 2
+    """
+    B = kpts0.shape[0] # B is the fine matches in a single pair
+    assert gt_matches.shape[0] == 1
+    kpts_warpped = torch.zeros_like(kpts0)
+    kpts_valid_mask = torch.zeros_like(kpts0[..., 0], dtype=torch.bool)
+    gt_matches_non_padding_mask = gt_matches.sum(-1) > 0
+
+    dist_matrix = torch.cdist(kpts0, gt_matches[..., :2]) # B * N * M
+    if dist_thr is not None:
+        mask = dist_matrix < dist_thr[:, None, None]
+    else:
+        mask = torch.ones_like(dist_matrix, dtype=torch.bool)
+    # Mutual-Nearest check:
+    mask = mask \
+        * (dist_matrix == dist_matrix.min(dim=2, keepdim=True)[0]) \
+        * (dist_matrix == dist_matrix.min(dim=1, keepdim=True)[0])
+
+    mask_v, all_j_ids = mask.max(dim=2)
+    b_ids, i_ids = torch.where(mask_v)
+    j_ids = all_j_ids[b_ids, i_ids]
+
+    j_uq_indices = get_unique_indices(torch.stack([b_ids, j_ids], dim=-1))
+    b_ids, i_ids, j_ids = map(lambda x: x[j_uq_indices], [b_ids, i_ids, j_ids])
+
+    i_uq_indices = get_unique_indices(torch.stack([b_ids, i_ids], dim=-1))
+    b_ids, i_ids, j_ids = map(lambda x: x[i_uq_indices], [b_ids, i_ids, j_ids])
+
+    kpts_valid_mask[b_ids, i_ids] = gt_matches_non_padding_mask[0, j_ids]
+    kpts_warpped[b_ids, i_ids] = gt_matches[..., 2:][0, j_ids]
+
+    return kpts_valid_mask, kpts_warpped
+
+@torch.no_grad()
+def warp_kpts_by_sparse_gt_matches_fast(kpts0, gt_matches, scale0, current_h, current_w):
+    B, n_gt_pts = gt_matches.shape[0], gt_matches.shape[1]
+    kpts_warpped = torch.zeros_like(kpts0)
+    kpts_valid_mask = torch.zeros_like(kpts0[..., 0], dtype=torch.bool)
+    gt_matches_non_padding_mask = gt_matches.sum(-1) > 0
+
+    all_j_idxs = torch.arange(gt_matches.shape[-2], device=gt_matches.device, dtype=torch.long)[None].expand(B, n_gt_pts)
+    all_b_idxs = torch.arange(B, device=gt_matches.device, dtype=torch.long)[:, None].expand(B, n_gt_pts)
+    gt_matches_rescale = gt_matches[..., :2] / scale0 # From original img scale to resized scale
+    in_boundary_mask = (gt_matches_rescale[..., 0] <= current_w-1) & (gt_matches_rescale[..., 0] >= 0) & (gt_matches_rescale[..., 1] <= current_h -1) & (gt_matches_rescale[..., 1] >= 0)
+
+    gt_matches_rescale = gt_matches_rescale.round().to(torch.long)
+    all_i_idxs = gt_matches_rescale[..., 1] * current_w + gt_matches_rescale[..., 0] # idx = y * w + x
+
+    # Filter:
+    b_ids, i_ids, j_ids = map(lambda x: x[gt_matches_non_padding_mask & in_boundary_mask], [all_b_idxs, all_i_idxs, all_j_idxs])
+
+    j_uq_indices = get_unique_indices(torch.stack([b_ids, j_ids], dim=-1))
+    b_ids, i_ids, j_ids = map(lambda x: x[j_uq_indices], [b_ids, i_ids, j_ids])
+
+    i_uq_indices = get_unique_indices(torch.stack([b_ids, i_ids], dim=-1))
+    b_ids, i_ids, j_ids = map(lambda x: x[i_uq_indices], [b_ids, i_ids, j_ids])
+
+    kpts_valid_mask[b_ids, i_ids] = gt_matches_non_padding_mask[b_ids, j_ids]
+    kpts_warpped[b_ids, i_ids] = gt_matches[..., 2:][b_ids, j_ids]
+
+    return kpts_valid_mask, kpts_warpped
+
+
+@torch.no_grad()
+def homo_warp_kpts(kpts0, norm_pixel_mat, homo_sample_normed, original_size0=None, original_size1=None):
+    """
+    original_size1: N * 2, (h, w)
+    """
+    normed_kpts0_h = norm_pixel_mat @ torch.cat([kpts0, torch.ones_like(kpts0[:, :, [0]])], dim=-1).transpose(2, 1)  # (N * 3 * L)
+    kpts_warpped_h = (torch.linalg.inv(norm_pixel_mat) @ homo_sample_normed @ normed_kpts0_h).transpose(2, 1)  # (N * L * 3)
+    kpts_warpped = kpts_warpped_h[..., :2] / kpts_warpped_h[..., [2]]  # N * L * 2
+    valid_mask = (kpts_warpped[..., 0] > 0) & (kpts_warpped[..., 0] < original_size1[:, [1]]) & (kpts_warpped[..., 1] > 0) \
+                & (kpts_warpped[..., 1] < original_size1[:, [0]])  # N * L
+    if original_size0 is not None:
+        valid_mask *= (kpts0[..., 0] > 0) & (kpts0[..., 0] < original_size0[:, [1]]) & (kpts0[..., 1] > 0) \
+                    & (kpts0[..., 1] < original_size0[:, [0]]) # N * L
+            
+    return valid_mask, kpts_warpped
+
+@torch.no_grad()
+# if using mask in homo warp(for coarse supervision)
+def homo_warp_kpts_with_mask(kpts0, scale, depth_mask, norm_pixel_mat, homo_sample_normed, original_size0=None, original_size1=None):
+    """
+    original_size1: N * 2, (h, w)
+    """
+    normed_kpts0_h = norm_pixel_mat @ torch.cat([kpts0, torch.ones_like(kpts0[:, :, [0]])], dim=-1).transpose(2, 1)  # (N * 3 * L)
+    kpts_warpped_h = (torch.linalg.inv(norm_pixel_mat) @ homo_sample_normed @ normed_kpts0_h).transpose(2, 1)  # (N * L * 3)
+    kpts_warpped = kpts_warpped_h[..., :2] / kpts_warpped_h[..., [2]]  # N * L * 2
+    # get coarse-level depth_mask
+    depth_mask_coarse = depth_mask[:, :, ::scale, ::scale]
+    depth_mask_coarse = depth_mask_coarse.reshape(depth_mask.shape[0], -1)
+    
+    valid_mask = (kpts_warpped[..., 0] > 0) & (kpts_warpped[..., 0] < original_size1[:, [1]]) & (kpts_warpped[..., 1] > 0) \
+                & (kpts_warpped[..., 1] < original_size1[:, [0]]) & (depth_mask_coarse != 0)  # N * L
+    if original_size0 is not None:
+        valid_mask *= (kpts0[..., 0] > 0) & (kpts0[..., 0] < original_size0[:, [1]]) & (kpts0[..., 1] > 0) \
+                    & (kpts0[..., 1] < original_size0[:, [0]]) & (depth_mask_coarse != 0) # N * L
+            
+    return valid_mask, kpts_warpped
+
+@torch.no_grad()
+# if using mask in homo warp(for fine supervision)
+def homo_warp_kpts_with_mask_f(kpts0, depth_mask, norm_pixel_mat, homo_sample_normed, original_size0=None, original_size1=None):
+    """
+    original_size1: N * 2, (h, w)
+    """
+    normed_kpts0_h = norm_pixel_mat @ torch.cat([kpts0, torch.ones_like(kpts0[:, :, [0]])], dim=-1).transpose(2, 1)  # (N * 3 * L)
+    kpts_warpped_h = (torch.linalg.inv(norm_pixel_mat) @ homo_sample_normed @ normed_kpts0_h).transpose(2, 1)  # (N * L * 3)
+    kpts_warpped = kpts_warpped_h[..., :2] / kpts_warpped_h[..., [2]]  # N * L * 2
+    valid_mask = (kpts_warpped[..., 0] > 0) & (kpts_warpped[..., 0] < original_size1[:, [1]]) & (kpts_warpped[..., 1] > 0) \
+                & (kpts_warpped[..., 1] < original_size1[:, [0]]) & (depth_mask != 0)  # N * L
+    if original_size0 is not None:
+        valid_mask *= (kpts0[..., 0] > 0) & (kpts0[..., 0] < original_size0[:, [1]]) & (kpts0[..., 1] > 0) \
+                    & (kpts0[..., 1] < original_size0[:, [0]]) & (depth_mask != 0) # N * L
+            
+    return valid_mask, kpts_warpped
+
+@torch.no_grad()
+def homo_warp_kpts_glue(kpts0, homo, original_size0=None, original_size1=None):
+    """
+    original_size1: N * 2, (h, w)
+    """
+    kpts_warpped = warp_points_torch(kpts0, homo, inverse=False)
+    valid_mask = (kpts_warpped[..., 0] > 0) & (kpts_warpped[..., 0] < original_size1[:, [1]]) & (kpts_warpped[..., 1] > 0) \
+                & (kpts_warpped[..., 1] < original_size1[:, [0]]) # N * L
+    if original_size0 is not None:
+        valid_mask *= (kpts0[..., 0] > 0) & (kpts0[..., 0] < original_size0[:, [1]]) & (kpts0[..., 1] > 0) \
+                    & (kpts0[..., 1] < original_size0[:, [0]]) # N * L
+    return valid_mask, kpts_warpped
+        
+@torch.no_grad()
+# if using mask in homo warp(for coarse supervision)
+def homo_warp_kpts_glue_with_mask(kpts0, scale, depth_mask, homo, original_size0=None, original_size1=None):
+    """
+    original_size1: N * 2, (h, w)
+    """
+    kpts_warpped = warp_points_torch(kpts0, homo, inverse=False)
+    # get coarse-level depth_mask
+    depth_mask_coarse = depth_mask[:, :, ::scale, ::scale]
+    depth_mask_coarse = depth_mask_coarse.reshape(depth_mask.shape[0], -1)
+
+    valid_mask = (kpts_warpped[..., 0] > 0) & (kpts_warpped[..., 0] < original_size1[:, [1]]) & (kpts_warpped[..., 1] > 0) \
+                & (kpts_warpped[..., 1] < original_size1[:, [0]]) & (depth_mask_coarse != 0)  # N * L
+    if original_size0 is not None:
+        valid_mask *= (kpts0[..., 0] > 0) & (kpts0[..., 0] < original_size0[:, [1]]) & (kpts0[..., 1] > 0) \
+                    & (kpts0[..., 1] < original_size0[:, [0]]) & (depth_mask_coarse != 0)  # N * L
+    return valid_mask, kpts_warpped
+
+@torch.no_grad()
+# if using mask in homo warp(for fine supervision)
+def homo_warp_kpts_glue_with_mask_f(kpts0, depth_mask, homo, original_size0=None, original_size1=None):
+    """
+    original_size1: N * 2, (h, w)
+    """
+    kpts_warpped = warp_points_torch(kpts0, homo, inverse=False)
+    valid_mask = (kpts_warpped[..., 0] > 0) & (kpts_warpped[..., 0] < original_size1[:, [1]]) & (kpts_warpped[..., 1] > 0) \
+                & (kpts_warpped[..., 1] < original_size1[:, [0]]) & (depth_mask != 0)  # N * L
+    if original_size0 is not None:
+        valid_mask *= (kpts0[..., 0] > 0) & (kpts0[..., 0] < original_size0[:, [1]]) & (kpts0[..., 1] > 0) \
+                    & (kpts0[..., 1] < original_size0[:, [0]]) & (depth_mask != 0)  # N * L
+    return valid_mask, kpts_warpped

vismatch/third_party/MatchAnything/imcui/third_party/MatchAnything/src/loftr/utils/position_encoding.py

@@ -0,0 +1,131 @@
+import math
+import torch
+from torch import nn
+
+
+class PositionEncodingSine(nn.Module):
+    """
+    This is a sinusoidal position encoding that generalized to 2-dimensional images
+    """
+
+    def __init__(self, d_model, max_shape=(256, 256), temp_bug_fix=True, npe=False):
+        """
+        Args:
+            max_shape (tuple): for 1/8 featmap, the max length of 256 corresponds to 2048 pixels
+            temp_bug_fix (bool): As noted in this [issue](https://github.com/zju3dv/LoFTR/issues/41),
+                the original implementation of LoFTR includes a bug in the pos-enc impl, which has little impact
+                on the final performance. For now, we keep both impls for backward compatability.
+                We will remove the buggy impl after re-training all variants of our released models.
+        """
+        super().__init__()
+
+        pe = torch.zeros((d_model, *max_shape))
+        y_position = torch.ones(max_shape).cumsum(0).float().unsqueeze(0)
+        x_position = torch.ones(max_shape).cumsum(1).float().unsqueeze(0)
+
+        assert npe is not None
+        if npe is not None:
+            if isinstance(npe, bool):
+                train_res_H, train_res_W, test_res_H, test_res_W = 832, 832, 832, 832
+                print('loftr no npe!!!!', npe)
+            else:
+                print('absnpe!!!!', npe)
+                train_res_H, train_res_W, test_res_H, test_res_W = npe[0], npe[1], npe[2], npe[3] # train_res_H, train_res_W, test_res_H, test_res_W
+                y_position, x_position = y_position * train_res_H / test_res_H, x_position * train_res_W / test_res_W
+
+        if temp_bug_fix:
+            div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / (d_model//2)))
+        else:  # a buggy implementation (for backward compatability only)
+            div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / d_model//2))
+        div_term = div_term[:, None, None]  # [C//4, 1, 1]
+        pe[0::4, :, :] = torch.sin(x_position * div_term)
+        pe[1::4, :, :] = torch.cos(x_position * div_term)
+        pe[2::4, :, :] = torch.sin(y_position * div_term)
+        pe[3::4, :, :] = torch.cos(y_position * div_term)
+
+        self.register_buffer('pe', pe.unsqueeze(0), persistent=False)  # [1, C, H, W]
+
+    def forward(self, x):
+        """
+        Args:
+            x: [N, C, H, W]
+        """
+        return x + self.pe[:, :, :x.size(2), :x.size(3)]
+    
+class RoPEPositionEncodingSine(nn.Module):
+    """
+    This is a sinusoidal position encoding that generalized to 2-dimensional images
+    """
+
+    def __init__(self, d_model, max_shape=(256, 256), npe=None, ropefp16=True):
+        """
+        Args:
+            max_shape (tuple): for 1/8 featmap, the max length of 256 corresponds to 2048 pixels
+            temp_bug_fix (bool): As noted in this [issue](https://github.com/zju3dv/LoFTR/issues/41),
+                the original implementation of LoFTR includes a bug in the pos-enc impl, which has little impact
+                on the final performance. For now, we keep both impls for backward compatability.
+                We will remove the buggy impl after re-training all variants of our released models.
+        """
+        super().__init__()
+
+        # pe = torch.zeros((d_model, *max_shape))
+        # y_position = torch.ones(max_shape).cumsum(0).float().unsqueeze(-1)
+        # x_position = torch.ones(max_shape).cumsum(1).float().unsqueeze(-1)
+        i_position = torch.ones(max_shape).cumsum(0).float().unsqueeze(-1) # [H, 1]
+        j_position = torch.ones(max_shape).cumsum(1).float().unsqueeze(-1) # [W, 1]
+        
+        assert npe is not None
+        if npe is not None:
+            train_res_H, train_res_W, test_res_H, test_res_W = npe[0], npe[1], npe[2], npe[3] # train_res_H, train_res_W, test_res_H, test_res_W
+            i_position, j_position = i_position * train_res_H / test_res_H, j_position * train_res_W / test_res_W
+        
+        div_term = torch.exp(torch.arange(0, d_model//4, 1).float() * (-math.log(10000.0) / (d_model//4)))
+        div_term = div_term[None, None, :]  # [1, 1, C//4]
+        # pe[0::4, :, :] = torch.sin(x_position * div_term)
+        # pe[1::4, :, :] = torch.cos(x_position * div_term)
+        # pe[2::4, :, :] = torch.sin(y_position * div_term)
+        # pe[3::4, :, :] = torch.cos(y_position * div_term)
+        sin = torch.zeros(*max_shape, d_model//2, dtype=torch.float16 if ropefp16 else torch.float32)
+        cos = torch.zeros(*max_shape, d_model//2, dtype=torch.float16 if ropefp16 else torch.float32)
+        sin[:, :, 0::2] = torch.sin(i_position * div_term).half() if ropefp16 else torch.sin(i_position * div_term)
+        sin[:, :, 1::2] = torch.sin(j_position * div_term).half() if ropefp16 else torch.sin(j_position * div_term)
+        cos[:, :, 0::2] = torch.cos(i_position * div_term).half() if ropefp16 else torch.cos(i_position * div_term)
+        cos[:, :, 1::2] = torch.cos(j_position * div_term).half() if ropefp16 else torch.cos(j_position * div_term)
+
+        sin = sin.repeat_interleave(2, dim=-1)
+        cos = cos.repeat_interleave(2, dim=-1)
+        # self.register_buffer('pe', pe.unsqueeze(0), persistent=False)  # [1, H, W, C]
+        self.register_buffer('sin', sin.unsqueeze(0), persistent=False)  # [1, H, W, C//2]
+        self.register_buffer('cos', cos.unsqueeze(0), persistent=False)  # [1, H, W, C//2]
+        
+        i_position4 = i_position.reshape(64,4,64,4,1)[...,0,:]
+        i_position4 = i_position4.mean(-3)
+        j_position4 = j_position.reshape(64,4,64,4,1)[:,0,...]
+        j_position4 = j_position4.mean(-2)
+        sin4 = torch.zeros(max_shape[0]//4, max_shape[1]//4, d_model//2, dtype=torch.float16 if ropefp16 else torch.float32)
+        cos4 = torch.zeros(max_shape[0]//4, max_shape[1]//4, d_model//2, dtype=torch.float16 if ropefp16 else torch.float32)
+        sin4[:, :, 0::2] = torch.sin(i_position4 * div_term).half() if ropefp16 else torch.sin(i_position4 * div_term)
+        sin4[:, :, 1::2] = torch.sin(j_position4 * div_term).half() if ropefp16 else torch.sin(j_position4 * div_term)
+        cos4[:, :, 0::2] = torch.cos(i_position4 * div_term).half() if ropefp16 else torch.cos(i_position4 * div_term)
+        cos4[:, :, 1::2] = torch.cos(j_position4 * div_term).half() if ropefp16 else torch.cos(j_position4 * div_term)
+        sin4 = sin4.repeat_interleave(2, dim=-1)
+        cos4 = cos4.repeat_interleave(2, dim=-1)
+        self.register_buffer('sin4', sin4.unsqueeze(0), persistent=False)  # [1, H, W, C//2]
+        self.register_buffer('cos4', cos4.unsqueeze(0), persistent=False)  # [1, H, W, C//2]
+
+        
+
+    def forward(self, x, ratio=1):
+        """
+        Args:
+            x: [N, H, W, C]
+        """
+        if ratio == 4:
+            return (x * self.cos4[:, :x.size(1), :x.size(2), :]) + (self.rotate_half(x) * self.sin4[:, :x.size(1), :x.size(2), :])
+        else:
+            return (x * self.cos[:, :x.size(1), :x.size(2), :]) + (self.rotate_half(x) * self.sin[:, :x.size(1), :x.size(2), :])
+    
+    def rotate_half(self, x):
+        x = x.unflatten(-1, (-1, 2))
+        x1, x2 = x.unbind(dim=-1)
+        return torch.stack((-x2, x1), dim=-1).flatten(start_dim=-2)