Spaces:
Running
Running
File size: 8,834 Bytes
8bd45de |
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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 |
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# utilities needed for the inference
# --------------------------------------------------------
import torch
import numpy as np
from .utils.misc import invalid_to_zeros
from .utils.geometry import geotrf, inv
def loss_of_one_batch(loss_func, batch, model, criterion, device,
use_amp=False, ret=None,
assist_model=None, train=False, epoch=0,
args=None):
if loss_func == "i2p":
return loss_of_one_batch_multiview(batch, model, criterion,
device, use_amp, ret,
args.ref_id)
elif loss_func == "i2p_corr_score":
return loss_of_one_batch_multiview_corr_score(batch, model, criterion,
device, use_amp, ret,
args.ref_id)
elif loss_func == "l2w":
return loss_of_one_batch_l2w(
batch, model, criterion,
device, use_amp, ret,
ref_ids=args.ref_ids, coord_frame_id=0,
exclude_ident=True, to_zero=True
)
else:
raise NotImplementedError
def loss_of_one_batch_multiview(batch, model, criterion, device,
use_amp=False, ret=None, ref_id=-1):
""" Function to compute the reconstruction loss of the Image-to-Points model
"""
views = batch
for view in views:
for name in 'img pts3d valid_mask camera_pose'.split(): # pseudo_focal
if name not in view:
continue
view[name] = view[name].to(device, non_blocking=True)
if ref_id == -1:
ref_id = (len(views)-1)//2
with torch.cuda.amp.autocast(enabled=bool(use_amp)):
preds = model(views, ref_id=ref_id)
assert len(preds) == len(views)
with torch.cuda.amp.autocast(enabled=False):
if criterion is None:
loss = None
else:
loss = criterion(views, preds, ref_id=ref_id)
result = dict(views=views, preds=preds, loss=loss)
for i in range(len(preds)):
result[f'pred{i+1}'] = preds[i]
result[f'view{i+1}'] = views[i]
return result[ret] if ret else result
def loss_of_one_batch_multiview_corr_score(batch, model, criterion, device,
use_amp=False, ret=None, ref_id=-1):
views = batch
for view in views:
for name in 'img pts3d valid_mask camera_pose'.split(): # pseudo_focal
if name not in view:
continue
view[name] = view[name].to(device, non_blocking=True)
if ref_id == -1:
ref_id = (len(views)-1)//2
all_loss = [0, {}]
with torch.cuda.amp.autocast(enabled=bool(use_amp)):
preds = model(views, ref_id=ref_id, return_corr_score=True)
assert len(preds) == len(views)
for i,pred in enumerate(preds):
if i == ref_id:
continue
patch_pseudo_conf = pred['pseudo_conf'] # (B,S)
true_conf = (pred['conf']-1.).mean(dim=(1,2)) # (B,) mean(exp(x))
pseudo_conf = torch.exp(patch_pseudo_conf).mean(dim=1) # (B,) mean(exp(batch(x)))
pseudo_conf = pseudo_conf / (1+pseudo_conf)
true_conf = true_conf / (1+true_conf)
dis = torch.abs(pseudo_conf-true_conf)
loss = dis.mean()
# if loss.isinf():
# print(((patch_pseudo_conf-patch_true_conf)**2).max())
all_loss[0] += loss
all_loss[1][f'pseudo_conf_loss_{i}'] = loss
result = dict(views=views, preds=preds, loss=all_loss)
for i in range(len(preds)):
result[f'pred{i+1}'] = preds[i]
result[f'view{i+1}'] = views[i]
return result[ret] if ret else result
def get_multiview_scale(pts:list, valid:list, norm_mode='avg_dis'):
# adpat from DUSt3R
for i in range(len(pts)):
assert pts[i].ndim >= 3 and pts[i].shape[-1] == 3
assert len(pts) == len(valid)
norm_mode, dis_mode = norm_mode.split('_')
if norm_mode == 'avg':
# gather all points together (joint normalization)
all_pts = []
all_nnz = 0
for i in range(len(pts)):
nan_pts, nnz = invalid_to_zeros(pts[i], valid[i], ndim=3)
# print(nnz,nan_pts.shape) #(B,) (B,H*W,3)
all_pts.append(nan_pts)
all_nnz += nnz
all_pts = torch.cat(all_pts, dim=1)
# compute distance to origin
all_dis = all_pts.norm(dim=-1)
if dis_mode == 'dis':
pass # do nothing
elif dis_mode == 'log1p':
all_dis = torch.log1p(all_dis)
else:
raise ValueError(f'bad {dis_mode=}')
norm_factor = all_dis.sum(dim=1) / (all_nnz + 1e-8)
else:
raise ValueError(f'bad {norm_mode=}')
norm_factor = norm_factor.clip(min=1e-8)
while norm_factor.ndim < pts[0].ndim:
norm_factor.unsqueeze_(-1)
# print('norm factor:', norm_factor)
return norm_factor
def loss_of_one_batch_l2w(batch, model, criterion, device,
use_amp=False, ret=None,
ref_ids=-1, coord_frame_id=0,
exclude_ident=True, to_zero=True):
""" Function to compute the reconstruction loss of the Local-to-World model
ref_ids: list of indices of the suppporting frames(excluding the coord_frame)
coord_frame_id: all the pointmaps input and output will be in the coord_frame_id's camera coordinate
exclude_ident: whether to exclude the coord_frame to simulate real-life inference scenarios
to_zero: whether to set the invalid points to zero
"""
views = batch
for view in views:
for name in 'img pts3d pts3d_cam valid_mask camera_pose'.split(): # pseudo_focal
if name not in view:
continue
view[name] = view[name].to(device, non_blocking=True)
if coord_frame_id == -1:
# ramdomly select a camera as the target camera
coord_frame_id = np.random.randint(0, len(views))
# print(coord_frame_id)
c2w = views[coord_frame_id]['camera_pose']
w2c = inv(c2w)
# exclude the frame that has the identity pose
if exclude_ident:
views.pop(coord_frame_id)
if ref_ids == -1:
ref_ids = [i for i in range(len(views)-1)] # all views except the last one
elif ref_ids == -2:
#select half of the views randomly
ref_ids = np.random.choice(len(views), len(views)//2, replace=False).tolist()
else:
assert isinstance(ref_ids, list)
for id in ref_ids:
views[id]['pts3d_world'] = geotrf(w2c, views[id]['pts3d']) #转移到目标坐标系
norm_factor_world = get_multiview_scale([views[id]['pts3d_world'] for id in ref_ids],
[views[id]['valid_mask'] for id in ref_ids],
norm_mode='avg_dis')
for id,view in enumerate(views):
if id in ref_ids:
view['pts3d_world'] = view['pts3d_world'].permute(0,3,1,2) / norm_factor_world
else:
norm_factor_src = get_multiview_scale([view['pts3d_cam']],
[view['valid_mask']],
norm_mode='avg_dis')
view['pts3d_cam'] = view['pts3d_cam'].permute(0,3,1,2) / norm_factor_src
if to_zero:
for id,view in enumerate(views):
valid_mask = view['valid_mask'].unsqueeze(1).float() # B,1,H,W
if id in ref_ids:
# print(view['pts3d_world'].shape, valid_mask.shape, (-valid_mask+1).sum())
view['pts3d_world'] = view['pts3d_world'] * valid_mask
else:
view['pts3d_cam'] = view['pts3d_cam'] * valid_mask
with torch.cuda.amp.autocast(enabled=bool(use_amp)):
preds = model(views, ref_ids=ref_ids)
assert len(preds) == len(views)
with torch.cuda.amp.autocast(enabled=False):
if criterion is None:
loss = None
else:
loss = criterion(views, preds, ref_id=ref_ids, ref_camera=w2c, norm_scale=norm_factor_world)
result = dict(views=views, preds=preds, loss=loss)
for i in range(len(preds)):
result[f'pred{i+1}'] = preds[i]
result[f'view{i+1}'] = views[i]
return result[ret] if ret else result
|