| 189 | |
| 190 | |
| 191 | def get_robust_pca(features: torch.Tensor, m: float = 2, remove_first_component=False): |
| 192 | # features: (N, C) |
| 193 | # m: a hyperparam controlling how many std dev outside for outliers |
| 194 | assert len(features.shape) == 2, "features should be (N, C)" |
| 195 | reduction_mat = torch.pca_lowrank(features, q=3, niter=20)[2] |
| 196 | colors = features @ reduction_mat |
| 197 | if remove_first_component: |
| 198 | colors_min = colors.min(dim=0).values |
| 199 | colors_max = colors.max(dim=0).values |
| 200 | tmp_colors = (colors - colors_min) / (colors_max - colors_min) |
| 201 | fg_mask = tmp_colors[..., 0] < 0.2 |
| 202 | reduction_mat = torch.pca_lowrank(features[fg_mask], q=3, niter=20)[2] |
| 203 | colors = features @ reduction_mat |
| 204 | else: |
| 205 | fg_mask = torch.ones_like(colors[:, 0]).bool() |
| 206 | d = torch.abs(colors[fg_mask] - torch.median(colors[fg_mask], dim=0).values) |
| 207 | mdev = torch.median(d, dim=0).values |
| 208 | s = d / mdev |
| 209 | try: |
| 210 | rins = colors[fg_mask][s[:, 0] < m, 0] |
| 211 | gins = colors[fg_mask][s[:, 1] < m, 1] |
| 212 | bins = colors[fg_mask][s[:, 2] < m, 2] |
| 213 | rgb_min = torch.tensor([rins.min(), gins.min(), bins.min()]) |
| 214 | rgb_max = torch.tensor([rins.max(), gins.max(), bins.max()]) |
| 215 | except: |
| 216 | rins = colors |
| 217 | gins = colors |
| 218 | bins = colors |
| 219 | rgb_min = torch.tensor([rins.min(), gins.min(), bins.min()]) |
| 220 | rgb_max = torch.tensor([rins.max(), gins.max(), bins.max()]) |
| 221 | |
| 222 | return reduction_mat, rgb_min.to(reduction_mat), rgb_max.to(reduction_mat) |
| 223 | |
| 224 | |
| 225 | def get_pca_map( |