|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import torch
|
|
|
| def mse(img1, img2):
|
| return (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True)
|
|
|
| def psnr(img1, img2):
|
| mse = (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True)
|
| return 20 * torch.log10(1.0 / torch.sqrt(mse))
|
|
|
| def get_robust_pca(features: torch.Tensor, m: float = 2, remove_first_component=False):
|
|
|
|
|
| assert len(features.shape) == 2, "features should be (N, C)"
|
| reduction_mat = torch.pca_lowrank(features, q=3, niter=20)[2]
|
| colors = features @ reduction_mat
|
| if remove_first_component:
|
| colors_min = colors.min(dim=0).values
|
| colors_max = colors.max(dim=0).values
|
| tmp_colors = (colors - colors_min) / (colors_max - colors_min)
|
| fg_mask = tmp_colors[..., 0] < 0.2
|
| reduction_mat = torch.pca_lowrank(features[fg_mask], q=3, niter=20)[2]
|
| colors = features @ reduction_mat
|
| else:
|
| fg_mask = torch.ones_like(colors[:, 0]).bool()
|
| d = torch.abs(colors[fg_mask] - torch.median(colors[fg_mask], dim=0).values)
|
| mdev = torch.median(d, dim=0).values
|
| s = d / mdev
|
| rins = colors[fg_mask][s[:, 0] < m, 0]
|
| gins = colors[fg_mask][s[:, 1] < m, 1]
|
| bins = colors[fg_mask][s[:, 2] < m, 2]
|
|
|
| rgb_min = torch.tensor([rins.min(), gins.min(), bins.min()])
|
| rgb_max = torch.tensor([rins.max(), gins.max(), bins.max()])
|
| return reduction_mat, rgb_min.to(reduction_mat), rgb_max.to(reduction_mat)
|
|
|
