Spaces:
Sleeping
Sleeping
| from ..custom_types import * | |
| def get_gm_support(gm, x): | |
| dim = x.shape[-1] | |
| mu, p, phi, eigen = gm | |
| sigma_det = eigen.prod(-1) | |
| eigen_inv = 1 / eigen | |
| sigma_inverse = torch.matmul(p.transpose(3, 4), p * eigen_inv[:, :, :, :, None]).squeeze(1) | |
| phi = torch.softmax(phi, dim=2) | |
| const_1 = phi / torch.sqrt((2 * np.pi) ** dim * sigma_det) | |
| distance = x[:, :, None, :] - mu | |
| mahalanobis_distance = - .5 * torch.einsum('bngd,bgdc,bngc->bng', distance, sigma_inverse, distance) | |
| const_2, _ = mahalanobis_distance.max(dim=2) # for numeric stability | |
| mahalanobis_distance -= const_2[:, :, None] | |
| support = const_1 * torch.exp(mahalanobis_distance) | |
| return support, const_2 | |
| def gm_log_likelihood_loss(gms: TS, x: T, get_supports: bool = False, | |
| mask: Optional[T] = None, reduction: str = "mean") -> Union[T, Tuple[T, TS]]: | |
| batch_size, num_points, dim = x.shape | |
| support, const = get_gm_support(gms, x) | |
| probs = torch.log(support.sum(dim=2)) + const | |
| if mask is not None: | |
| probs = probs.masked_select(mask=mask.flatten()) | |
| if reduction == 'none': | |
| likelihood = probs.sum(-1) | |
| loss = - likelihood / num_points | |
| else: | |
| likelihood = probs.sum() | |
| loss = - likelihood / (probs.shape[0] * probs.shape[1]) | |
| if get_supports: | |
| return loss, support | |
| return loss | |
| def split_mesh_by_gmm(mesh: T_Mesh, gmm) -> Dict[int, T]: | |
| faces_split = {} | |
| vs, faces = mesh | |
| vs_mid_faces = vs[faces].mean(1) | |
| _, supports = gm_log_likelihood_loss(gmm, vs_mid_faces.unsqueeze(0), get_supports=True) | |
| supports = supports[0] | |
| label = supports.argmax(1) | |
| for i in range(gmm[1].shape[2]): | |
| select = label.eq(i) | |
| if select.any(): | |
| faces_split[i] = faces[select] | |
| else: | |
| faces_split[i] = None | |
| return faces_split | |
| def flatten_gmm(gmm: TS) -> T: | |
| b, gp, g, _ = gmm[0].shape | |
| mu, p, phi, eigen = [item.view(b, gp * g, *item.shape[3:]) for item in gmm] | |
| p = p.reshape(*p.shape[:2], -1) | |
| z_gmm = torch.cat((mu, p, phi.unsqueeze(-1), eigen), dim=2) | |
| return z_gmm | |