Spaces:
Running
Running
File size: 2,672 Bytes
56cd6b7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# utility functions for global alignment
# --------------------------------------------------------
import torch
import torch.nn as nn
import numpy as np
def edge_str(i, j):
return f'{i}_{j}'
def i_j_ij(ij):
return edge_str(*ij), ij
def edge_conf(conf_i, conf_j, edge):
return float(conf_i[edge].mean() * conf_j[edge].mean())
# edge对应的两张图片经dust3r输出的置信度，分别对两张图片所有像素点的置信度取平均值再相乘，作为当前edge的置信度
def compute_edge_scores(edges, conf_i, conf_j):# edge对应的两张图片经dust3r会输出两个置信度矩阵，分别对两张图片所有像素点的置信度取平均值再相乘，作为当前edge的置信度
return {(i, j): edge_conf(conf_i, conf_j, e) for e, (i, j) in edges}
def NoGradParamDict(x):
assert isinstance(x, dict)
return nn.ParameterDict(x).requires_grad_(False)
def get_imshapes(edges, pred_i, pred_j):
n_imgs = max(max(e) for e in edges) + 1
imshapes = [None] * n_imgs
for e, (i, j) in enumerate(edges):
shape_i = tuple(pred_i[e].shape[0:2])
shape_j = tuple(pred_j[e].shape[0:2])
if imshapes[i]:
assert imshapes[i] == shape_i, f'incorrect shape for image {i}'
if imshapes[j]:
assert imshapes[j] == shape_j, f'incorrect shape for image {j}'
imshapes[i] = shape_i
imshapes[j] = shape_j
return imshapes
def get_conf_trf(mode):
if mode == 'log':
def conf_trf(x): return x.log()
elif mode == 'sqrt':
def conf_trf(x): return x.sqrt()
elif mode == 'm1':
def conf_trf(x): return x-1
elif mode in ('id', 'none'):
def conf_trf(x): return x
else:
raise ValueError(f'bad mode for {mode=}')
return conf_trf
def l2_dist(a, b, weight):
return ((a - b).square().sum(dim=-1) * weight)
def l1_dist(a, b, weight):
return ((a - b).norm(dim=-1) * weight) # torch.norm()是求范式的损失，默认是第二范式
ALL_DISTS = dict(l1=l1_dist, l2=l2_dist)
def signed_log1p(x):
sign = torch.sign(x)
return sign * torch.log1p(torch.abs(x))
def signed_expm1(x):
sign = torch.sign(x)
return sign * torch.expm1(torch.abs(x))
def cosine_schedule(t, lr_start, lr_end):
assert 0 <= t <= 1
return lr_end + (lr_start - lr_end) * (1+np.cos(t * np.pi))/2
def linear_schedule(t, lr_start, lr_end):
assert 0 <= t <= 1
return lr_start + (lr_end - lr_start) * t
|