|
|
| from __future__ import annotations |
|
|
| from typing import Callable |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from unisharp.utils.pixel_convention import scale_intrinsics_align_corners_false |
|
|
|
|
| def resize_k3_align_corners_false(k: torch.Tensor, *, sx: float, sy: float) -> torch.Tensor: |
| return scale_intrinsics_align_corners_false(k, sx=float(sx), sy=float(sy)) |
|
|
|
|
| def resize_rgb_u8_chw_high_quality(image: torch.Tensor, *, size: tuple[int, int]) -> torch.Tensor: |
| if not torch.is_tensor(image) or image.ndim != 3: |
| raise ValueError(f"Expected CHW tensor, got {tuple(image.shape) if torch.is_tensor(image) else type(image)}") |
| dst_h, dst_w = int(size[0]), int(size[1]) |
| if tuple(image.shape[-2:]) == (dst_h, dst_w): |
| return image.contiguous() |
| resized = F.interpolate( |
| image.unsqueeze(0).to(torch.float32), |
| size=(dst_h, dst_w), |
| mode="bicubic", |
| align_corners=False, |
| antialias=True, |
| ) |
| return resized[0].round().clamp(0, 255).to(torch.uint8).contiguous() |
|
|
|
|
| def project_overlap_ratio( |
| src_w2c: torch.Tensor, |
| tgt_w2c: torch.Tensor, |
| src_k: torch.Tensor, |
| tgt_k: torch.Tensor, |
| h: int, |
| w: int, |
| src_hw: tuple[int, int] | None = None, |
| tgt_hw: tuple[int, int] | None = None, |
| sample_h: int = 32, |
| sample_w: int = 56, |
| proxy_depth: float = 1.0, |
| ) -> float: |
| device = src_w2c.device |
| src_h, src_w = tuple(int(v) for v in (src_hw or (h, w))) |
| tgt_h, tgt_w = tuple(int(v) for v in (tgt_hw or (h, w))) |
| ys = torch.linspace(0, src_h - 1, steps=sample_h, device=device) |
| xs = torch.linspace(0, src_w - 1, steps=sample_w, device=device) |
| vv, uu = torch.meshgrid(ys, xs, indexing="ij") |
| u = uu.reshape(-1) |
| v = vv.reshape(-1) |
|
|
| fx, fy = src_k[0, 0], src_k[1, 1] |
| cx, cy = src_k[0, 2], src_k[1, 2] |
| x = (u - cx) / fx |
| y = (v - cy) / fy |
| z = torch.ones_like(x) |
| rays = torch.stack([x, y, z], dim=-1) |
| rays = rays / torch.norm(rays, dim=-1, keepdim=True).clamp(min=1e-6) |
| pts_src = rays * float(proxy_depth) |
|
|
| src_c2w = torch.linalg.inv(src_w2c) |
| pts_src_h = torch.cat([pts_src, torch.ones_like(pts_src[:, :1])], dim=-1) |
| pts_w = (src_c2w @ pts_src_h.T).T |
| pts_tgt = (tgt_w2c @ pts_w.T).T |
| xt, yt, zt = pts_tgt[:, 0], pts_tgt[:, 1], pts_tgt[:, 2].clamp(min=1e-6) |
| ut = tgt_k[0, 0] * (xt / zt) + tgt_k[0, 2] |
| vt = tgt_k[1, 1] * (yt / zt) + tgt_k[1, 2] |
| inside = (zt > 0.0) & (ut >= 0.0) & (ut <= float(tgt_w - 1)) & (vt >= 0.0) & (vt <= float(tgt_h - 1)) |
| return float(inside.float().mean().item()) |
|
|
|
|
| def select_targets_for_source( |
| *, |
| src_idx: int, |
| candidate_indices: list[int], |
| centers: torch.Tensor, |
| min_index_gap: int, |
| max_index_gap: int, |
| pair_max_translation_m: float, |
| pair_min_overlap: float, |
| overlap_score_fn: Callable[[int, int], float], |
| ) -> list[int]: |
| src_c = centers[int(src_idx)] |
| tgt_cands: list[int] = [] |
| for j in candidate_indices: |
| j = int(j) |
| if j == int(src_idx): |
| continue |
| gap = abs(int(j) - int(src_idx)) |
| if gap < int(min_index_gap) or gap > int(max_index_gap): |
| continue |
| trans = float(torch.norm(centers[j] - src_c, p=2).item()) |
| if trans > float(pair_max_translation_m): |
| continue |
| if float(overlap_score_fn(int(src_idx), j)) >= float(pair_min_overlap): |
| tgt_cands.append(j) |
| return tgt_cands |
|
|
