import torch from vidar.utils.read import read_pickle from vidar.geometry.camera import Camera, Camera_from_list from vidar.geometry.pose import Pose from vidar.utils.write import write_image from vidar.utils.viz import viz_normals, viz_depth from externals.vggt.training.loss import point_map_to_normal, normal_loss, check_and_fix_inf_nan from vidar.utils.depth import scale_and_shift_pred def calc_normals(points, mask, eps=1e-12): normals, masks = point_map_to_normal(points.permute(0, 2, 3, 1), mask.squeeze(1), eps=eps) normals = normals.permute(1, 4, 2, 3, 0) masks = masks.permute(1, 2, 3, 0) return normals, masks def calc_normals_loss(gt_normals, pred_normals, gt_valids, pred_valids, cos_eps=1e-8, conf=None, gamma=1.0, alpha=0.2): # # Convert point maps to surface normals using cross products # pred_normals, pred_valids = point_map_to_normal(prediction, mask, eps=cos_eps) # gt_normals, gt_valids = point_map_to_normal(target, mask, eps=cos_eps) gt_normals = gt_normals.permute(4, 0, 2, 3, 1) gt_valids = gt_valids.permute(3, 0, 1, 2) pred_normals = pred_normals.permute(4, 0, 2, 3, 1) pred_valids = pred_valids.permute(3, 0, 1, 2) # Only consider regions where both predicted and GT normals are valid all_valid = pred_valids & gt_valids # shape: (4, B, H, W) # Early return if not enough valid points divisor = torch.sum(all_valid) if divisor < 10: return 0 # Extract valid normals pred_normals = pred_normals[all_valid].clone() gt_normals = gt_normals[all_valid].clone() # Compute cosine similarity between corresponding normals dot = torch.sum(pred_normals * gt_normals, dim=-1) # Clamp dot product to [-1, 1] for numerical stability dot = torch.clamp(dot, -1 + cos_eps, 1 - cos_eps) # Compute loss as 1 - cos(theta), instead of arccos(dot) for numerical stability loss = 1 - dot # Return mean loss if we have enough valid points if loss.numel() < 10: return 0 else: loss = check_and_fix_inf_nan(loss, "normal_loss") if conf is not None: # Apply confidence weighting conf = conf.unsqueeze(0).squeeze(2).repeat(4, 1, 1, 1) conf = conf[all_valid].clone() loss = gamma * loss * conf - alpha * torch.log(conf) return loss.mean() else: return loss.mean() def calc_regression_loss(gt, pred, mask, conf, gamma=1.0, alpha=0.2): weight_absrel = 1.0 weight_scaleshift = 0.0 loss_regr_absrel = (gt - pred).abs().sum(1, keepdim=True) # scale, shift, scaled_pred = scale_and_shift_pred(pred, gt, mask) # loss_regr_scaleshift = (gt - scaled_pred).abs().sum(1, keepdim=True) # loss_regr = weight_absrel * loss_regr_absrel + \ # weight_scaleshift * loss_regr_scaleshift loss_regr = weight_absrel * loss_regr_absrel loss_conf = gamma * loss_regr * conf - alpha * torch.log(conf) return { 'loss_regr_absrel': loss_regr_absrel[mask].mean(), 'loss_regr': loss_regr[mask].mean(), 'loss_conf': loss_conf[mask].mean(), } def get_data(): data = read_pickle('data_batch5.pkl')['anydata'] output = read_pickle('output5.pkl') gt_rgb = torch.stack([val for val in data['rgb'].values()], 1) gt_depth = torch.stack([val for val in data['depth'].values()], 1) gt_cams = Camera_from_list([val for val in data['cams'].values()]) b, n, _, h, w = gt_rgb.shape gt_depth = gt_depth.view(b * n, -1, h, w) gt_mask_depth = gt_depth > 0 ############# # pred_points_wld = gt_cams.reconstruct_depth_map(gt_depth, to_world=True) # pred_conf_points_wld = torch.randn_like(gt_mask_depth.float()) + 10 pred_points_wld = output['world_points'].view(b * n, h, w, -1).permute(0, 3, 1, 2) pred_conf_points_wld = output['world_points_conf'].view(b * n, h, w).unsqueeze(1) ############# return data, pred_points_wld, pred_conf_points_wld def calc_loss(data, pred_points_wld, pred_conf_points_wld): gt_rgb = torch.stack([val for val in data['rgb'].values()], 1) gt_depth = torch.stack([val for val in data['depth'].values()], 1) gt_cams = Camera_from_list([val for val in data['cams'].values()]) b, n, _, h, w = gt_rgb.shape gt_rgb = gt_rgb.view(b * n, -1, h, w) gt_depth = gt_depth.view(b * n, -1, h, w) gt_mask_depth = gt_depth > 0 gt_points_cam = gt_cams.reconstruct_depth_map(gt_depth, to_world=False) gt_points_wld = gt_cams.reconstruct_depth_map(gt_depth, to_world=True) gt_origin_wld = gt_cams.get_origin() pred_points_cam = gt_cams.from_world(pred_points_wld) gt_normals_cam, gt_masks_normals_cam = calc_normals(gt_points_cam, gt_mask_depth) gt_normals_wld, gt_masks_normals_wld = calc_normals(gt_points_wld, gt_mask_depth) pred_normals_wld, pred_masks_normals_wld = calc_normals(pred_points_wld, gt_mask_depth) pred_normals_cam, pred_masks_normals_cam = calc_normals(pred_points_cam, gt_mask_depth) gt_depth_points_cam = ((gt_points_wld - gt_origin_wld) ** 2).sum(1, keepdim=True).sqrt() pred_depth_points_cam = ((pred_points_wld - gt_origin_wld) ** 2).sum(1, keepdim=True).sqrt() loss_normals_cam = calc_normals_loss( gt_normals_cam, pred_normals_cam, gt_masks_normals_cam, pred_masks_normals_cam, conf=pred_conf_points_wld, ) loss_normals_wld = calc_normals_loss( gt_normals_wld, pred_normals_wld, gt_masks_normals_wld, pred_masks_normals_wld, conf=pred_conf_points_wld, ) loss_normals = loss_normals_cam + loss_normals_wld loss_regr_points_wld = calc_regression_loss( gt_points_wld, pred_points_wld, mask=gt_mask_depth, conf=pred_conf_points_wld, ) loss_regr_depth_points_cam = calc_regression_loss( gt_depth_points_cam, pred_depth_points_cam, mask=gt_mask_depth, conf=pred_conf_points_wld, ) loss = \ 1.000 * (loss_regr_points_wld['loss_regr'] + loss_regr_points_wld['loss_conf'] ) + \ 0.001 * (loss_regr_depth_points_cam['loss_regr'] + loss_regr_depth_points_cam['loss_conf'] ) + \ 0.050 * loss_normals loss_dict = { 'loss_regr_points_wld': loss_regr_points_wld['loss_regr'], 'loss_regr_absrel_points_wld': loss_regr_points_wld['loss_regr_absrel'], 'loss_conf_points_wld': loss_regr_points_wld['loss_conf'], 'loss_regr_depth_points_cam': loss_regr_depth_points_cam['loss_regr'], 'loss_regr_absrel_depth_points_cam': loss_regr_depth_points_cam['loss_regr_absrel'], 'loss_conf_depth_points_cam': loss_regr_depth_points_cam['loss_conf'], 'loss_normals': loss_normals, 'objective': loss, 'loss': loss, } ########################## # from camviz import Draw # draw = Draw((1600, 900)) # draw.add3Dworld('wld', # pose=(-0.01255, -0.13684, -0.11027, 0.91111, 0.40866, 0.05339, 0.00472), # ) # draw.addBufferf('gt_pts_wld', gt_points_wld[:6].permute(0, 2, 3, 1).reshape(-1, 3)) # draw.addBufferf('gt_pts_cam', gt_points_cam[:6].permute(0, 2, 3, 1).reshape(-1, 3)) # draw.addBufferf('pred_pts', pred_points_wld[:6].permute(0, 2, 3, 1).reshape(-1, 3)) # draw.addBufferf('clr', gt_rgb[:6].permute(0, 2, 3, 1).reshape(-1, 3)) # from camviz.objects.camera import Camera as CameraCV # cv_cams =[CameraCV.from_vidar(gt_cams[i], scale=0.1 / 200) for i in range(6)] # show = 0 # while draw.input(): # if draw.RETURN: # show = (show + 1) % 2 # draw.halt(100) # draw.clear() # if show == 0: # draw['wld'].size(2).color('whi').points('gt_pts_wld', 'clr') # # draw['wld'].size(2).color('whi').points('gt_pts_cam', 'clr') # if show == 1: # draw['wld'].size(2).color('whi').points('pred_pts', 'clr') # for i in range(6): # draw['wld'].object(cv_cams[i]) # draw.update(30) ########################## # for i in range(4): # for j in range(6): # write_image(f'viz/gt_normals_wld_{j}_{i}.png', viz_normals(gt_normals_wld[j, ..., i])) # write_image(f'viz/gt_normals_cam_{j}_{i}.png', viz_normals(gt_normals_cam[j, ..., i])) # write_image(f'viz/pred_normals_wld_{j}_{i}.png', viz_normals(pred_normals_wld[j, ..., i])) # write_image(f'viz/pred_normals_cam_{j}_{i}.png', viz_normals(pred_normals_cam[j, ..., i])) # write_image(f'viz/gt_depth_wld_{j}.png', viz_depth(gt_depth_points_cam[j])) # write_image(f'viz/pred_depth_wld_{j}.png', viz_depth(pred_depth_points_cam[j])) # write_image(f'viz/gt_depth_{j}.png', viz_depth(gt_depth[j])) ########################## return loss_dict if __name__ == '__main__': data, pred_points_wld, pred_conf_points_wld = get_data() calc_loss(data, pred_points_wld, pred_conf_points_wld) # import torch # from vidar.utils.read import read_pickle # from vidar.geometry.camera import Camera # from vidar.geometry.pose import Pose # from camviz.objects.camera import Camera as CameraCV # from vidar.utils.write import write_image # from vidar.utils.viz import viz_normals, viz_depth # from externals.vggt.training.loss import point_map_to_normal, normal_loss, check_and_fix_inf_nan # from vidar.utils.depth import scale_and_shift_pred # def calc_normals(points, mask): # normals, masks = point_map_to_normal(points.permute(0, 2, 3, 1), mask.squeeze(1)) # normals = normals.permute(1, 4, 2, 3, 0) # masks = masks.permute(1, 2, 3, 0) # return normals, masks # def calc_normals_loss(gt_normals, pred_normals, gt_valids, pred_valids, cos_eps=1e-8, conf=None): # # # Convert point maps to surface normals using cross products # # pred_normals, pred_valids = point_map_to_normal(prediction, mask, eps=cos_eps) # # gt_normals, gt_valids = point_map_to_normal(target, mask, eps=cos_eps) # gt_normals = gt_normals.permute(4, 0, 2, 3, 1) # gt_valids = gt_valids.permute(3, 0, 1, 2) # pred_normals = pred_normals.permute(4, 0, 2, 3, 1) # pred_valids = pred_valids.permute(3, 0, 1, 2) # # Only consider regions where both predicted and GT normals are valid # all_valid = pred_valids & gt_valids # shape: (4, B, H, W) # # Early return if not enough valid points # divisor = torch.sum(all_valid) # if divisor < 10: # return 0 # # Extract valid normals # pred_normals = pred_normals[all_valid].clone() # gt_normals = gt_normals[all_valid].clone() # # Compute cosine similarity between corresponding normals # dot = torch.sum(pred_normals * gt_normals, dim=-1) # # Clamp dot product to [-1, 1] for numerical stability # dot = torch.clamp(dot, -1 + cos_eps, 1 - cos_eps) # # Compute loss as 1 - cos(theta), instead of arccos(dot) for numerical stability # loss = 1 - dot # # Return mean loss if we have enough valid points # if loss.numel() < 10: # return 0 # else: # loss = check_and_fix_inf_nan(loss, "normal_loss") # if conf is not None: # # Apply confidence weighting # conf = conf[None, ...].expand(4, -1, -1, -1) # conf = conf[all_valid].clone() # loss = gamma * loss * conf - alpha * torch.log(conf) # return loss.mean() # else: # return loss.mean() # def calc_regression_loss(gt, pred, mask, conf, gamma=1.0, alpha=0.2): # weight_absrel = 1.0 # weight_scaleshift = 0.1 # loss_regr_absrel = (gt - pred).abs().sum(1, keepdim=True) # scale, shift, scaled_pred = scale_and_shift_pred(pred, gt, mask) # loss_regr_scaleshift = (gt - scaled_pred).abs().sum(1, keepdim=True) # loss_regr = weight_absrel * loss_regr_absrel + \ # weight_scaleshift *loss_regr_scaleshift # loss_regr = loss_regr[mask].mean() # loss_conf = gamma * loss_regr * conf - alpha * torch.log(conf) # loss_conf = loss_conf[mask].mean() # return loss_regr, loss_conf # data = read_pickle('data_batch.pkl')['anydata'] # gt_rgb = torch.stack([val[0] for val in data['rgb'].values()], 0) # gt_depth = torch.stack([val[0] for val in data['depth'].values()], 0) # gt_intrinsics = torch.stack([val[0] for val in data['intrinsics'].values()], 0) # gt_mask_depth = gt_depth > 0 # gt_pose = Pose.from_dict(data['pose'], to_global=True, zero_origin=True, broken=True) # gt_pose = torch.stack([val.T[0] for val in gt_pose.values()], 0) # gt_cams = Camera(K=gt_intrinsics, Twc=gt_pose, hw=gt_depth) # gt_points_cam = gt_cams.reconstruct_depth_map(gt_depth, to_world=False) # gt_points_wld = gt_cams.reconstruct_depth_map(gt_depth, to_world=True) # gt_origin_wld = gt_cams.get_origin() # ############# # pred_points_wld = gt_cams.reconstruct_depth_map(gt_depth, to_world=True) # pred_conf_points_wld = torch.randn_like(gt_mask_depth.float()) + 10 # ############# # pred_points_cam = gt_cams.from_world(pred_points_wld) # gt_normals_cam, gt_masks_normals_cam = calc_normals(gt_points_cam, gt_mask_depth) # gt_normals_wld, gt_masks_normals_wld = calc_normals(gt_points_wld, gt_mask_depth) # pred_normals_wld, pred_masks_normals_wld = calc_normals(pred_points_wld, gt_mask_depth) # pred_normals_cam, pred_masks_normals_cam = calc_normals(pred_points_cam, gt_mask_depth) # gt_depth_points_wld = ((gt_points_wld - gt_origin_wld) ** 2).sum(1, keepdim=True).sqrt() # pred_depth_points_wld = ((pred_points_wld - gt_origin_wld) ** 2).sum(1, keepdim=True).sqrt() # loss_normals_cam = calc_normals_loss( # gt_normals_cam, pred_normals_cam, # gt_masks_normals_cam, pred_masks_normals_cam, # ) # loss_normals_wld = calc_normals_loss( # gt_normals_wld, pred_normals_wld, # gt_masks_normals_wld, pred_masks_normals_wld, # ) # loss_regr, loss_conf = calc_regression_loss( # gt_depth_points_wld, pred_depth_points_wld, # gt_mask_depth, pred_conf_points_wld, # ) # print(loss_regr) # print(loss_conf) # # print('aaaaaaaaaaa', gt_points_wld.shape, gt_points_wld.device, pred_points_wld.device, gt_mask_depth.device) # # from moge_loss import affine_invariant_global_loss, affine_invariant_local_loss # # aaa = affine_invariant_global_loss( # # gt_points_wld.permute(0, 2, 3, 1), # # pred_points_wld.permute(0, 2, 3, 1), # # gt_mask_depth.permute(0, 2, 3, 1).squeeze(-1), # # ) # # print(aaa) # # bbb = affine_invariant_local_loss( # # gt_points_wld.permute(0, 2, 3, 1), # # pred_points_wld.permute(0, 2, 3, 1), # # gt_mask_depth.permute(0, 2, 3, 1).squeeze(-1), # # ) # # print(bbb) # for i in range(4): # for j in range(6): # write_image(f'viz/gt_normals_wld_{j}_{i}.png', viz_normals(gt_normals_wld[j, ..., i])) # write_image(f'viz/gt_normals_cam_{j}_{i}.png', viz_normals(gt_normals_cam[j, ..., i])) # write_image(f'viz/pred_normals_wld_{j}_{i}.png', viz_normals(pred_normals_wld[j, ..., i])) # write_image(f'viz/pred_normals_cam_{j}_{i}.png', viz_normals(pred_normals_cam[j, ..., i])) # write_image(f'viz/gt_depth_wld_{j}.png', viz_depth(gt_depth_points_wld[j])) # write_image(f'viz/pred_depth_wld_{j}.png', viz_depth(pred_depth_points_wld[j])) # # write_image(f'normals_wld_cam_{j}_{i}.png', viz_normals(normals_wld_cam[j, ..., i])) # # write_image(f'masks_wld_{j}_{i}.png', gt_masks_normals_wld[j, ..., i].unsqueeze(0)) # # write_image(f'masks_cam_{j}_{i}.png', gt_masks_normals_cam[j, ..., i].unsqueeze(0)) # # write_image(f'masks_wld_cam_{j}_{i}.png', masks_normals_wld_cam[j, ..., i].unsqueeze(0)) # # from camviz import Draw # # draw = Draw((1600, 900)) # # draw.add3Dworld('wld') # # draw.addBufferf('pts', points_wld.permute(0, 2, 3, 1).reshape(-1, 3)) # # draw.addBufferf('clr', rgb.permute(0, 2, 3, 1).reshape(-1, 3)) # # cv_cams =[CameraCV.from_vidar(cams[i], scale=0.1) for i in range(6)] # # while draw.input(): # # draw.clear() # # draw['wld'].size(2).color('whi').points('pts', 'clr') # # for i in range(6): # # draw['wld'].object(cv_cams[i]) # # draw.update(30) # # # Compute loss (per instance) # # loss_list, weight_list = [], [] # # for i in range(current_batch_size): # # gt_metric_scale = None # # loss_dict, weight_dict, misc_dict = {}, {}, {} # # misc_dict['monitoring'] = monitoring(pred_points[i]) # # for k, v in config['loss'][label_type[i]].items(): # # weight_dict[k] = v['weight'] # # if v['function'] == 'affine_invariant_global_loss': # # loss_dict[k], misc_dict[k], gt_metric_scale = affine_invariant_global_loss(pred_points[i], gt_points[i], gt_mask[i], **v['params']) # # elif v['function'] == 'affine_invariant_local_loss': # # loss_dict[k], misc_dict[k] = affine_invariant_local_loss(pred_points[i], gt_points[i], gt_mask[i], gt_focal[i], gt_metric_scale, **v['params']) # # elif v['function'] == 'normal_loss': # # loss_dict[k], misc_dict[k] = normal_loss(pred_points[i], gt_points[i], gt_mask[i]) # # elif v['function'] == 'edge_loss': # # loss_dict[k], misc_dict[k] = edge_loss(pred_points[i], gt_points[i], gt_mask[i]) # # elif v['function'] == 'mask_bce_loss': # # loss_dict[k], misc_dict[k] = mask_bce_loss(pred_mask[i], gt_mask_fin[i], gt_mask_inf[i]) # # elif v['function'] == 'mask_l2_loss': # # loss_dict[k], misc_dict[k] = mask_l2_loss(pred_mask[i], gt_mask_fin[i], gt_mask_inf[i]) # # else: # # raise ValueError(f'Undefined loss function: {v["function"]}') # # weight_dict = {'.'.join(k): v for k, v in flatten_nested_dict(weight_dict).items()} # # loss_dict = {'.'.join(k): v for k, v in flatten_nested_dict(loss_dict).items()} # # loss_ = sum([weight_dict[k] * loss_dict[k] for k in loss_dict], start=torch.tensor(0.0, device=device)) # # loss_list.append(loss_)