Dororo99/Ours_S3GS_Waymo

1#
2# Copyright (C) 2023, Inria
3# GRAPHDECO research group, https://team.inria.fr/graphdeco
4# All rights reserved.
5#
6# This software is free for non-commercial, research and evaluation use 
7# under the terms of the LICENSE.md file.
8#
9# For inquiries contact  george.drettakis@inria.fr
10#
11
12import torch
13
14def mse(img1, img2):
15    return (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True)
16
17def psnr(img1, img2):
18    mse = (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True)
19    return 20 * torch.log10(1.0 / torch.sqrt(mse))
20
21def get_robust_pca(features: torch.Tensor, m: float = 2, remove_first_component=False):
22    # features: (N, C)
23    # m: a hyperparam controlling how many std dev outside for outliers
24    assert len(features.shape) == 2, "features should be (N, C)"
25    reduction_mat = torch.pca_lowrank(features, q=3, niter=20)[2]
26    colors = features @ reduction_mat
27    if remove_first_component:
28        colors_min = colors.min(dim=0).values
29        colors_max = colors.max(dim=0).values
30        tmp_colors = (colors - colors_min) / (colors_max - colors_min)
31        fg_mask = tmp_colors[..., 0] < 0.2
32        reduction_mat = torch.pca_lowrank(features[fg_mask], q=3, niter=20)[2]
33        colors = features @ reduction_mat
34    else:
35        fg_mask = torch.ones_like(colors[:, 0]).bool()
36    d = torch.abs(colors[fg_mask] - torch.median(colors[fg_mask], dim=0).values)
37    mdev = torch.median(d, dim=0).values
38    s = d / mdev
39    rins = colors[fg_mask][s[:, 0] < m, 0]
40    gins = colors[fg_mask][s[:, 1] < m, 1]
41    bins = colors[fg_mask][s[:, 2] < m, 2]
42
43    rgb_min = torch.tensor([rins.min(), gins.min(), bins.min()])
44    rgb_max = torch.tensor([rins.max(), gins.max(), bins.max()])
45    return reduction_mat, rgb_min.to(reduction_mat), rgb_max.to(reduction_mat)
46