Code of Neural-Singular-Hessian: Implicit Neural Representation of Unoriented Point Clouds by Enforcing Singular Hessian. ACM Transactions on Graphics (SIGGRAPH Asia 2023).
from .overfit_network import Network
from .losses import MorseLoss
from .shape_network import ShapeNetwork

models/__init__.py
bearprin/Neural-Singular-Hessian
"""
Convolutional Feature Embedding
"""
from typing import Callable, Dict, List, Tuple, Union
import torch
from .custom import grid_sample_2d, grid_sample_3d
from torch import nn
from torch_scatter import scatter_mean, scatter_max
from .ocn.common import coordinate2index, normalize_coordinate, normalize_3d_coordinate
from .ocn.layers import ResnetBlockFC
from .ocn.unet import UNet
from .ocn.unet3d import UNet3D
import numpy as np


def get_knn(vs: torch.Tensor, k: int, batch_idx: torch.Tensor) -> torch.Tensor:
    mat_square = torch.matmul(vs, vs.transpose(2, 1))
    diag = torch.diagonal(mat_square, dim1=1, dim2=2)
    diag = diag.unsqueeze(2).expand(mat_square.shape)
    dist_mat = (diag + diag.transpose(2, 1) - 2 * mat_square)
    _, index = dist_mat.topk(k + 1, dim=2, largest=False, sorted=True)
    index = index[:, :, 1:].view(-1, k) + batch_idx[:, None] * vs.shape[1]
    return index.flatten()


def extract_angles(vs: torch.Tensor, distance_k: torch.Tensor, vs_k: torch.Tensor) -> Union[
    Tuple[torch.Tensor, ...], List[torch.Tensor]]:
    proj = torch.einsum('nd,nkd->nk', vs, vs_k)
    cos_angles = torch.clamp(proj / distance_k, -1., 1.)
    proj = vs_k - vs[:, None, :] * proj[:, :, None]
    # moving same axis points
    ma = torch.abs(proj).sum(2) == 0
    num_points_to_replace = ma.sum().item()
    if num_points_to_replace:
        proj[ma] = torch.rand(num_points_to_replace,
                              vs.shape[1], device=ma.device)
    proj = proj / torch.norm(proj, p=2, dim=2)[:, :, None]
    angles = torch.acos(cos_angles)
    return angles, proj


def min_angles(dirs: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
    ref = dirs[:, 0]
    all_cos = torch.einsum('nd,nkd->nk', ref, dirs)
    all_sin = torch.cross(ref.unsqueeze(
        1).expand(-1, dirs.shape[1], -1), dirs, dim=2)
    all_sin = torch.einsum('nd,nkd->nk', up, all_sin)
    all_angles = torch.atan2(all_sin, all_cos)
    all_angles[:, 0] = 0
    all_angles[all_angles < 0] = all_angles[all_angles < 0] + 2 * np.pi
    all_angles, inds = all_angles.sort(dim=1)
    inds = torch.argsort(inds, dim=1)
    all_angles_0 = 2 * np.pi - all_angles[:, -1]
    all_angles[:, 1:] = all_angles[:, 1:] - all_angles[:, :-1]
    all_angles[:, 0] = all_angles_0
    all_angles = torch.gather(all_angles, 1, inds)
    return all_angles


def extract_rotation_invariant_features(k: int) -> Tuple[Callable[[Union[torch.Tensor, np.array]], torch.Tensor], int]:
    batch_idx = None
    num_features = k * 3 + 1

    def get_input(xyz: torch.Tensor) -> Union[Tuple[torch.Tensor, ...], List[torch.Tensor]]:
        nonlocal batch_idx
        if batch_idx is None or len(batch_idx) != xyz.shape[0] * xyz.shape[1]:
            batch_idx, _ = torch.meshgrid(
                [torch.arange(xyz.shape[0]), torch.arange(xyz.shape[1])])
            batch_idx = batch_idx.flatten().to(xyz.device)
        return xyz.view(-1, 3), batch_idx

    def extract(base_vs: Union[torch.Tensor, np.array]):
        nonlocal num_features
        with torch.no_grad():
            if type(base_vs) is np.array:
                base_vs = torch.Tensor(base_vs)
            batch_size, num_pts = base_vs.shape[:2]
            vs, batch_idx = get_input(base_vs)
            knn = get_knn(base_vs, k, batch_idx)
            vs_k = vs[knn].view(-1, k, vs.shape[1])
            distance = torch.norm(vs, p=2, dim=1)
            vs_unit = vs / distance[:, None]
            distance_k = distance[knn].view(-1, k)
            angles, proj_unit = extract_angles(vs_unit, distance_k, vs_k)
            proj_min_angle = min_angles(proj_unit, vs_unit)
            fe = torch.cat([distance.unsqueeze(1), distance_k,
                            angles, proj_min_angle], dim=1)
        return fe.view(batch_size, num_pts, num_features)

    return extract, num_features


class ConvolutionalFeature(nn.Module):
    ''' PointNet-based encoder network with ResNet blocks for each point.
        Number of input points are fixed.

    Attributes:
        c_dim (int): dimension of latent code c
        dim (int): input points dimension
        hidden_dim (int): hidden dimension of the network
        scatter_type (str): feature aggregation when doing local pooling
        unet (bool): weather to use U-Net
        unet_kwargs (str): U-Net parameters
        unet3d (bool): weather to use 3D U-Net
        unet3d_kwargs (str): 3D U-Net parameters
        reso_plane (int): defined resolution for plane feature
        reso_grid (int): defined resolution for grid feature
        plane_type (str): feature type, 'xz' - 1-plane, ['xz', 'xy', 'yz'] - 3-plane, ['grid'] - 3D grid volume
        n_blocks (int): number of blocks ResNetBlockFC layers
    '''

    def __init__(self):
        super().__init__()

        # defaults from shapenet 3plane
        self.c_dim: int = 32
        self.dim: int = 3
        self.hidden_dim: int = 32
        self.scatter_type: str = 'max'
        self.unet: bool = False
        self.unet_kwargs: dict = dict(
            depth=4, merge_mode='concat', start_filts=32)
        self.unet3d: bool = True
        self.unet3d_kwargs: dict = {
            "num_levels": 4,
            "f_maps": 32,
            "in_channels": 32,
            "out_channels": 32
        }
        self.reso_plane: int = 64
        self.reso_grid: int = 64
        self.subpixel_upsampling: int = 1  # ratio for subpixel upsampling
        self.plane_type: List[str] = ['grid']
        self.n_blocks: int = 5
        self.sample_mode: str = "bilinear"
        self.input_normals: bool = False
        self.bbox_size: float = 2.0
        self.clusternet_feature: bool = False

        self._initialize()

    def _initialize(self):
        # shortcuts
        hidden_dim = self.hidden_dim
        dim = self.dim
        n_blocks = self.n_blocks
        c_dim = self.c_dim
        unet_kwargs = self.unet_kwargs
        unet = self.unet
        unet3d = self.unet3d

        if self.clusternet_feature:
            self.extractor, dim = extract_rotation_invariant_features(8)
        self.fc_pos = nn.Linear(dim, 2 * hidden_dim)
        self.blocks = nn.ModuleList([
            ResnetBlockFC(2 * hidden_dim, hidden_dim) for i in range(n_blocks)
        ])
        self.fc_c = nn.Linear(hidden_dim, c_dim)

        self.actvn = nn.ReLU()
        self.hidden_dim = hidden_dim

        if unet:
            self.unet = UNet(c_dim * (self.subpixel_upsampling ** 2),
                             in_channels=c_dim, **unet_kwargs)
        else:
            self.unet = None

        if unet3d:
            unet3d_kwargs = self.unet3d_kwargs.copy()
            unet3d_kwargs["out_channels"] *= (self.subpixel_upsampling ** 3)
            self.unet3d = UNet3D(**self.unet3d_kwargs)
        else:
            self.unet3d = None

        if self.subpixel_upsampling > 1:
            self.ps = torch.nn.PixelShuffle(self.subpixel_upsampling)

        if self.scatter_type == 'max':
            self.scatter = scatter_max
        elif self.scatter_type == 'mean':
            self.scatter = scatter_mean
        else:
            raise ValueError('incorrect scatter type')

        self.last_epoch = -1

    def generate_plane_features(self, p, c, plane='xz'):
        # acquire indices of features in plane
        # normalize to the range of (0, 1)
        xy = normalize_coordinate(
            p.clone(), plane=plane, bbox_size=self.bbox_size)
        index = coordinate2index(xy, self.reso_plane)

        # scatter plane features from points
        fea_plane = c.new_zeros(p.size(0), self.c_dim, self.reso_plane ** 2)
        c = c.permute(0, 2, 1)  # B x 512 x torch.Tensor
        fea_plane = scatter_mean(c, index, out=fea_plane)  # B x 512 x reso^2
        # sparce matrix (B x 512 x reso x reso)
        fea_plane = fea_plane.reshape(
            p.size(0), self.c_dim, self.reso_plane, self.reso_plane)

        # process the plane features with UNet
        if self.unet is not None:
            fea_plane = self.unet(fea_plane)

        if self.subpixel_upsampling > 1:
            # pixelshuffle upsampling
            fea_plane = self.ps(fea_plane)

        return fea_plane

    def generate_grid_features(self, p, c):
        p_nor = normalize_3d_coordinate(
            p.clone(), bbox_size=self.bbox_size)
        index = coordinate2index(p_nor, self.reso_grid, coord_type='3d')
        # scatter grid features from points
        fea_grid = c.new_zeros(p.size(0), self.c_dim, self.reso_grid ** 3)
        c = c.permute(0, 2, 1)
        fea_grid = scatter_mean(c, index, out=fea_grid)  # B x C x reso^3
        # sparce matrix (B x 512 x reso x reso)
        fea_grid = fea_grid.reshape(
            p.size(0), self.c_dim, self.reso_grid, self.reso_grid, self.reso_grid)

        if self.unet3d is not None:
            fea_grid = self.unet3d(fea_grid)

        if self.subpixel_upsampling > 1:
            fea_grid = depthToVolume(fea_grid, self.subpixel_upsampling)

        return fea_grid

    def pool_local(self, xy, index, c):
        bs, fea_dim = c.size(0), c.size(2)
        keys = xy.keys()

        c_out = 0
        for key in keys:
            # scatter plane features from points
            if key == 'grid':
                fea = self.scatter(c.permute(0, 2, 1),
                                   index[key], dim_size=self.reso_grid ** 3)
            else:
                fea = self.scatter(c.permute(0, 2, 1),
                                   index[key], dim_size=self.reso_plane ** 2)
            if self.scatter == scatter_max:
                fea = fea[0]
            # gather feature back to points
            fea = fea.gather(dim=2, index=index[key].expand(-1, fea_dim, -1))
            c_out += fea
        return c_out.permute(0, 2, 1)

    def query_feature(self, fea: Dict[str, torch.Tensor], query_coords: torch.Tensor) -> torch.Tensor:
        # self.bbox_size = getattr(self.runner.data, 'bbox_size', self.bbox_size)

        # originally inside the decoder
        if self.c_dim != 0:
            plane_type = list(fea.keys())
            # merge different planes with sum
            c = 0
            if 'grid' in plane_type:
                c += self.sample_grid_feature(query_coords, fea['grid'])
            if 'xz' in plane_type:
                c += self.sample_plane_feature(query_coords,
                                               fea['xz'], plane='xz')
            if 'xy' in plane_type:
                c += self.sample_plane_feature(query_coords,
                                               fea['xy'], plane='xy')
            if 'yz' in plane_type:
                c += self.sample_plane_feature(query_coords,
                                               fea['yz'], plane='yz')
            c = c.transpose(1, 2)

        return c

    # def query_global_feature(self, fea: Dict[str, torch.Tensor]) -> torch.Tensor:
    #     """
    #     Global max and avg pooling
    #     Returns:
    #         (B, P, 2*self.c_dim)
    #     """
    #     if self.c_dim != 0:
    #         plane_type = list(fea.keys())
    #         # merge different planes with sum
    #         c = 0
    #         if 'grid' in plane_type:
    #             batch_size = fea['grid'].shape[0]
    #             c_max = fea['grid'].view(
    #                 batch_size, self.c_dim, -1).max(dim=-1)[0]
    #             c_avg = fea['grid'].view(
    #                 batch_size, self.c_dim, -1).mean(dim=-1)
    #             c += torch.cat([c_max, c_avg], dim=-1)
    #         if 'xz' in plane_type:
    #             batch_size = fea['xz'].shape[0]
    #             c_max = fea['xz'].view(
    #                 batch_size, self.c_dim, -1).max(dim=-1)[0]
    #             c_avg = fea['xz'].view(batch_size, self.c_dim, -1).mean(dim=-1)
    #             c += torch.cat([c_max, c_avg], dim=-1)
    #         if 'xy' in plane_type:
    #             batch_size = fea['xy'].shape[0]
    #             c_max = fea['xy'].view(
    #                 batch_size, self.c_dim, -1).max(dim=-1)[0]
    #             c_avg = fea['xy'].view(batch_size, self.c_dim, -1).mean(dim=-1)
    #             c += torch.cat([c_max, c_avg], dim=-1)
    #         if 'yz' in plane_type:
    #             batch_size = fea['yz'].shape[0]
    #             c_max = fea['yz'].view(
    #                 batch_size, self.c_dim, -1).max(dim=-1)[0]
    #             c_avg = fea['yz'].view(batch_size, self.c_dim, -1).mean(dim=-1)
    #             c += torch.cat([c_max, c_avg], dim=-1)
    #
    #     return c
    def query_global_feature(self, fea: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Global max pooling
        Returns:
            (B, P, self.c_dim)
        """
        if self.c_dim != 0:
            plane_type = list(fea.keys())
            # merge different planes with sum
            c_max = None
            c_avg = None
            if 'grid' in plane_type:
                batch_size = fea['grid'].shape[0]
                c_max = fea['grid'].view(
                    batch_size, self.c_dim, -1).max(dim=-1)[0]
                c_avg = fea['grid'].view(
                    batch_size, self.c_dim, -1).mean(dim=-1)
                c = c_max
                # c += torch.cat([c_max, c_avg], dim=-1)
            if 'xz' in plane_type:
                batch_size = fea['xz'].shape[0]
                c_max = fea['xz'].view(
                    batch_size, self.c_dim, -1).max(dim=-1)[0]
                c_avg = fea['xz'].view(batch_size, self.c_dim, -1).mean(dim=-1)
                # c += torch.cat([c_max, c_avg], dim=-1)
            if 'xy' in plane_type:
                batch_size = fea['xy'].shape[0]
                c_max = fea['xy'].view(
                    batch_size, self.c_dim, -1).max(dim=-1)[0]
                c_avg = fea['xy'].view(batch_size, self.c_dim, -1).mean(dim=-1)
                # c += torch.cat([c_max, c_avg], dim=-1)
            if 'yz' in plane_type:
                batch_size = fea['yz'].shape[0]
                c_max = fea['yz'].view(
                    batch_size, self.c_dim, -1).max(dim=-1)[0]
                c_avg = fea['yz'].view(batch_size, self.c_dim, -1).mean(dim=-1)
                # c += torch.cat([c_max, c_avg], dim=-1)
        c = c_max
        return c

    def sample_plane_feature(self, p, c, plane='xz'):
        # normalize to the range of (0, 1)
        xy = normalize_coordinate(
            p.clone(), plane=plane, bbox_size=self.bbox_size)
        xy = xy[:, :, None].float()
        vgrid = 2.0 * xy - 1.0  # normalize to (-1, 1)
        if c.shape[0] == 1 and vgrid.shape[0] > 0:
            c = c.expand(vgrid.shape[0], -1, -1, -1)
        c = grid_sample_2d(c, vgrid).squeeze(-1)
        return c

    def sample_grid_feature(self, p, c):
        # normalize to the range of (0, 1)
        p_nor = normalize_3d_coordinate(
            p.clone(), bbox_size=self.bbox_size)
        p_nor = p_nor[:, :, None, None].float()
        vgrid = 2.0 * p_nor - 1.0  # normalize to (-1, 1)
        # acutally trilinear interpolation if mode = 'bilinear'
        if c.shape[0] == 1 and vgrid.shape[0] > 0:
            c = c.expand(vgrid.shape[0], -1, -1, -1, -1)
        c = grid_sample_3d(c, vgrid).squeeze(-1).squeeze(-1)
        return c

    def encode(self, input: torch.Tensor) -> Dict[str, torch.Tensor]:
        input_coords = input[..., :3]
        if self.input_normals:
            input_feats = input[..., 3:(3 + self.dim)]
        elif self.clusternet_feature:
            input_feats = self.extractor(input_coords)
        else:
            input_feats = input_coords.clone()

        # acquire the index for each point
        coord = {}
        index = {}
        if 'xz' in self.plane_type:
            coord['xz'] = normalize_coordinate(input_coords.clone(
            ), plane='xz', bbox_size=self.bbox_size)
            index['xz'] = coordinate2index(coord['xz'], self.reso_plane)
        if 'xy' in self.plane_type:
            coord['xy'] = normalize_coordinate(input_coords.clone(
            ), plane='xy', bbox_size=self.bbox_size)
            index['xy'] = coordinate2index(coord['xy'], self.reso_plane)
        if 'yz' in self.plane_type:
            coord['yz'] = normalize_coordinate(input_coords.clone(
            ), plane='yz', bbox_size=self.bbox_size)
            index['yz'] = coordinate2index(coord['yz'], self.reso_plane)
        if 'grid' in self.plane_type:
            coord['grid'] = normalize_3d_coordinate(
                input_coords.clone(), bbox_size=self.bbox_size)
            index['grid'] = coordinate2index(
                coord['grid'], self.reso_grid, coord_type='3d')

        net = self.fc_pos(input_feats)

        net = self.blocks[0](net)
        for block in self.blocks[1:]:
            pooled = self.pool_local(coord, index, net)
            net = torch.cat([net, pooled], dim=2)
            net = block(net)

        c = self.fc_c(net)

        fea = {}
        if 'grid' in self.plane_type:
            fea['grid'] = self.generate_grid_features(input_coords, c)
        if 'xz' in self.plane_type:
            fea['xz'] = self.generate_plane_features(
                input_coords, c, plane='xz')
        if 'xy' in self.plane_type:
            fea['xy'] = self.generate_plane_features(
                input_coords, c, plane='xy')
        if 'yz' in self.plane_type:
            fea['yz'] = self.generate_plane_features(
                input_coords, c, plane='yz')

        return fea

    def query_local_coordinates(self, input_coords):
        if "grid " in self.plane_type and len(self.plane_type) > 1:
            # make sure plane resolution is the same as grid resolution
            assert (self.reso_grid ==
                    self.reso_plane), "Local coordinates mapping only supports equal grid_reso and plane_reso"

        # assuming input in between -0.5 to 0.5
        p_nor = normalize_3d_coordinate(
            input_coords.clone(), bbox_size=self.bbox_size)
        vgrid = 2.0 * p_nor - 1.0  # normalize to (-1, 1)

        if 'grid' in self.plane_type:
            cell_size = 2.0 / self.reso_plane
        else:
            cell_size = 2.0 / self.reso_grid

        return torch.remainder(vgrid, cell_size) / cell_size

    def forward(self, args) -> torch.Tensor:
        """
        Args:
            input_coords: (N, P, 3) input point positions (can use relative ones as well)
            query_coords (N, Q, 3) query point position
        Returns:
            query_fea: (N, P, c_dim)
        """
        input_coords = args['pointclouds']
        query_coords = args['coords']

        assert (query_coords.size(-1) == 3)

        fea = self.encode(input_coords)
        query_feat = self.query_feature(fea, query_coords)

        return query_feat

    def save(self, path):
        torch.save(self, path)


def depthToVolume(x, upsample_ratio):
    N, C, D, H, W = x.shape
    x = x.view(N, upsample_ratio, upsample_ratio, upsample_ratio, C //
               (upsample_ratio ** 3), D, H, W)  # (N, bs, bs, bs, C//bs^3, D, H, W)
    # (N, C//bs^2, D, bs H, bs, W, bs)
    x = x.permute(0, 4, 5, 1, 6, 2, 7, 3).contiguous()
    x = x.view(N, C // (upsample_ratio ** 3), D * upsample_ratio, H *
               upsample_ratio, W * upsample_ratio)  # (N, C//bs^2, H * bs, W * bs)
    return x


def volumeToDepth(x, upsample_ratio):
    N, C, D, H, W = x.shape
    x = x.view(N, C, D // upsample_ratio, upsample_ratio, H // upsample_ratio, upsample_ratio,
               W // upsample_ratio, upsample_ratio)  # (N, C, D//bs, bs, H//bs, bs, W//bs, bs)
    # (N, bs, bs, bs, C, D//bs, H//bs, W//bs)
    x = x.permute(0, 3, 5, 7, 1, 2, 4, 6).contiguous()
    x = x.view(N, C * (upsample_ratio ** 3), H // upsample_ratio,
               W // upsample_ratio)  # (N, C*bs^2, H//bs, W//bs)
    return x

models/convolutionalfeature.py
bearprin/Neural-Singular-Hessian
import torch
import torch.autograd as autograd
import torch.nn.functional as F

def grid_sample_2d(input, grid):
    # grid (B, Ho, Wo, 2)
    # input (B, C, H, W)
    # output (B, C, Ho, Wo)
    B, Ho, Wo, _ = grid.shape
    P = Ho * Wo
    _, C, H, W = input.shape

    # grid to index
    reso = torch.tensor([2.0 / (W-1), 2.0 / (H-1)]).view(1, 1, 1, 2).to(device=grid.device)
    # B, Ho, Wo, 2
    grid = (grid+1.0)/reso
    w1h1 = torch.floor(grid).long()
    w2h2 = w1h1 + 1
    assert(w1h1.min() >= 0 and w2h2[..., 0].max() <= W and w2h2[..., 1].max() <= H)

    input = F.pad(input, (0, 1, 0, 1), mode='replicate')

    # B, 4*P: Q11, Q12, Q21, Q22
    index = torch.stack([w1h1[...,1]*(W+1)+w1h1[...,0], (w1h1[...,1]+1)*(W+1)+w1h1[...,0],
                            w1h1[...,1]*(W+1)+w1h1[...,0]+1, (w1h1[...,1]+1)*(W+1)+w1h1[...,0]+1], dim=1).view(B, -1)
    # 4 of (B, C, P) (f11, f12, f21, f22)
    input_4samples = torch.gather(input.view(B, C, -1), -1, index.unsqueeze(1).expand(-1, C, -1)).split(P, dim=-1)

    # B, Ho, Wo, 2 -> B, P
    diff_x2x, diff_y2y = torch.unbind((w2h2 - grid).view(B, -1, 2), dim=-1)
    diff_xx1, diff_yy1 = torch.unbind((grid - w1h1).view(B, -1, 2), dim=-1)

    f11, f12, f21, f22 = input_4samples
    # B,1,P,2 @ B,C,P,2,2
    result = torch.stack([diff_x2x, diff_xx1],dim=-1).reshape(B,1,P,1,2).expand(-1,C,-1,-1,-1).reshape(-1,1,2) @ torch.stack([f11, f12, f21, f22],dim=-1).reshape(B,C,P,2,2).reshape(-1,2,2) @ torch.stack([diff_y2y, diff_yy1],dim=-1).reshape(B,1,P,2,1).expand(-1,C,-1,-1,-1).reshape(-1,2,1)
    result = result.view(B,C,Ho,Wo)

    return result

def grid_sample_3d(input, grid):
    """
    grid (B, Do, Ho, Wo, 3)
    input (B, C, D, H, W)
    output (B, C, Do, Ho, Wo)
    """
    B, Do, Ho, Wo, _ = grid.shape
    P = Ho * Wo * Do
    _, C, D, H, W = input.shape

    # ref = F.grid_sample(input, grid, padding_mode="border", align_corners=True)
    # grid to index B,1,1,1,3
    reso = torch.tensor([2.0/ (D-1), 2.0 / (W-1), 2.0 / (H-1)]).view(1, 1, 1, 1, 3).to(device=grid.device)
    # B, Ho, Wo, 2
    grid = (grid+1.0)/reso
    x000 = torch.floor(grid).long()

    input = F.pad(input, (0, 1, 0, 1, 0, 1), mode='replicate')

    # B, 8*P: Q000, Q001, Q010, Q011, Q100, Q101, Q110, Q111
    index = torch.stack([(x000[...,1]+x000[...,2]*(H+1))*(W+1)+x000[...,0],
                         (x000[...,1]+(x000[...,2]+1)*(H+1))*(W+1)+x000[...,0],
                         (x000[...,1]+1+x000[...,2]*(H+1))*(W+1)+x000[...,0],
                         (x000[...,1]+1+(x000[...,2]+1)*(H+1))*(W+1)+x000[...,0],
                         (x000[...,1]+x000[...,2]*(H+1))*(W+1)+x000[...,0]+1,
                         (x000[...,1]+(x000[...,2]+1)*(H+1))*(W+1)+x000[...,0]+1,
                         (x000[...,1]+1+x000[...,2]*(H+1))*(W+1)+x000[...,0]+1,
                         (x000[...,1]+1+(x000[...,2]+1)*(H+1))*(W+1)+x000[...,0]+1,
                         ], dim=1).view(B, -1)
    # 2 of (B, C, 4P) corresponding f000, f001, f010, f011, f100, f101, f110, f111
    f0xx, f1xx = torch.gather(input.view(B, C, -1), -1, index.unsqueeze(1).expand(-1, C, -1)).split(P*4, dim=-1)

    # B, Ho, Wo, 3 -> B, P
    xd, yd, zd = torch.unbind((grid - x000).view(B, -1, 3), dim=-1)

    # B, C, 4P
    fxx = f0xx * (1-xd).repeat(1, 4).unsqueeze(1) + f1xx * xd.repeat(1,4).unsqueeze(1)

    # f00, f01: B, C, 2P
    f0x, f1x = fxx.split(2*P, dim=-1)
    fx = f0x * (1-yd).repeat(1, 2).unsqueeze(1) + f1x*yd.repeat(1,2).unsqueeze(1)

    f0, f1 = fx.split(P, dim=-1)
    result = f0*(1-zd).unsqueeze(1) + f1*zd.unsqueeze(1)

    result = result.view(B,C,Do,Ho,Wo)
    # assert(torch.allclose(ref, result))
    return result


if __name__ == "__main__":
    from torch.autograd.gradcheck import gradgradcheck, gradcheck
    torch.manual_seed(0)

    # X = torch.linspace(0, 5, steps=6, dtype=torch.double).view(1, 1, 1, 6).expand(1, 1, 6, -1).contiguous()
    X = torch.rand((1, 1, 6, 6), dtype=torch.double)
    X.requires_grad_()

    grid = torch.rand((1, 3, 3, 2), dtype=torch.double) - 0.5
    grid *= 2.0

    grid.requires_grad_()

    S = grid_sample_2d(X, grid).sum()
    S.backward()
    print(grid.grad)
    grid.grad.zero_()

    S = grid_sample_2d(X, grid).sum()
    grid_grad = autograd.grad(S, grid, create_graph=True)[0]
    S_grad = grid_grad.sum()

    S_grad.backward()
    print(grid.grad)

    # 3D
    X = torch.rand((1, 1, 6, 6, 6), dtype=torch.double)
    X.requires_grad_()
    grid = torch.rand((1, 2, 2, 2, 3), dtype=torch.double) - 0.5
    grid *= 2.0

    grid.requires_grad_()

    S = grid_sample_3d(X, grid).sum()
    S.backward()
    print(grid.grad)
    grid.grad.zero_()

    S = grid_sample_3d(X, grid).sum()
    grid_grad = autograd.grad(S, grid, create_graph=True)[0]
    S_grad = grid_grad.sum()

    S_grad.backward()
    print(grid.grad)

models/custom.py
bearprin/Neural-Singular-Hessian
import numpy as np

import torch
import torch.nn as nn


def kaiming_leaky_init(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        torch.nn.init.kaiming_normal_(m.weight, a=0.2, mode='fan_in', nonlinearity='leaky_relu')


def frequency_init(freq):
    def init(m):
        with torch.no_grad():
            # if hasattr(m, 'weight'):
            if isinstance(m, nn.Linear):
                num_input = m.weight.size(-1)
                m.weight.uniform_(-np.sqrt(6 / num_input) / freq, np.sqrt(6 / num_input) / freq)

    return init


def first_layer_film_sine_init(m):
    with torch.no_grad():
        # if hasattr(m, 'weight'):
        if isinstance(m, nn.Linear):
            num_input = m.weight.size(-1)
            m.weight.uniform_(-1 / num_input, 1 / num_input)


class MappingNet(nn.Module):
    def __init__(self, dim: int, out_dim_per_layer: int, n_out_layers: int = 1, hidden_size: int = 256):
        super().__init__()
        self.dim: int = dim
        self.out_dim_per_layer: int = out_dim_per_layer
        self.n_out_layers: int = n_out_layers
        self.hidden_size: int = hidden_size

        last_length = self.out_dim_per_layer * 2 * self.n_out_layers
        self.net = nn.Sequential(nn.Linear(self.dim, self.hidden_size), nn.LeakyReLU(0.2, inplace=True),
                                 nn.Linear(self.hidden_size, self.hidden_size), nn.LeakyReLU(0.2, inplace=True),
                                 nn.Linear(self.hidden_size, self.hidden_size), nn.LeakyReLU(0.2, inplace=True),
                                 nn.Linear(self.hidden_size, last_length)
                                 )

        self.net.apply(kaiming_leaky_init)
        with torch.no_grad():
            self.net[-1].weight *= 0.25

    def forward(self, x):
        """
        Args:
            x (B, ..., dim)
        Returns:
            frequencies: (B, ..., n_out_layers, out_dim_per_layer)
            biases (B, ..., n_out_layers, out_dim_per_layer)
        """
        output = self.net(x).view(x.shape[:-1] + (self.n_out_layers, self.out_dim_per_layer * 2))
        freqs, biases = output.split(self.out_dim_per_layer, dim=-1)

        return freqs, biases


class FiLMLayer(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.layer = nn.Linear(input_dim, hidden_dim)

    def forward(self, x, freq, phase_shift):
        """
        sin(freq * lin(x) + phase_shift)
        Args:
            x (B, ..., input_dim)
            freq: (B,...,hidden_dim)
            phase_shift: (B,...,hidden_dim)
        Returns
            output (B, ..., hidden_dim)
        """
        x = self.layer(x)
        assert (freq.shape == phase_shift.shape)
        if freq.ndim == (x.ndim - 1):
            freq = freq.unsqueeze(-2)
            phase_shift = phase_shift.unsqueeze(-2)

        return torch.sin(freq * x + phase_shift)


class FilmSiren(nn.Module):

    def __init__(self, hidden_size, n_layers, c_dim=32):
        super().__init__()
        self.hidden_size: int = hidden_size
        self.mapping_hidden_size: int = hidden_size
        self.n_layers: int = n_layers
        self.dim: int = 3
        self.c_dim: int = c_dim
        self.out_dim: int = 1
        self.share_frequencies: bool = False
        self.base_omega: int = 30

        self._initialize()

    def _initialize(self):
        if self.c_dim > 0:
            self.frequency_net = MappingNet(self.c_dim, self.hidden_size,
                                            n_out_layers=(1 if self.share_frequencies else self.n_layers),
                                            hidden_size=self.mapping_hidden_size)

        sirens = []
        for i in range(self.n_layers):
            if i == 0:
                lin = FiLMLayer(self.dim, self.hidden_size)
            else:
                lin = FiLMLayer(self.hidden_size, self.hidden_size)
            sirens.append(lin)

        self.sirens = nn.ModuleList(sirens)

        self.final_linear = nn.Linear(self.hidden_size, self.out_dim)

        self.sirens.apply(frequency_init(self.base_omega))
        self.final_linear.apply(frequency_init(self.base_omega))
        self.sirens[0].apply(first_layer_film_sine_init)
        self.epoch = 0

    def forward(self, query_points, query_feat=None):
        coords = query_points
        c = query_feat
        # mappingnet (B,...,n_layers,hidden_size)
        frequencies, biases = self.frequency_net(c)
        frequencies = frequencies * self.base_omega / 2 + self.base_omega

        if frequencies.shape[-2] != self.n_layers:
            shp = frequencies.shape
            new_shp = shp[:-2] + (self.n_layers,) + shp[-1:]
            frequencies = frequencies.expand(new_shp)
            biases = biases.expand(new_shp)

        x = coords
        _it = 0
        for layer, gamma, beta in zip(self.sirens, frequencies.unbind(dim=-2), biases.unbind(dim=-2)):
            x = layer(x, gamma, beta)
            _it += 1

        x = self.final_linear(x)

        return x  # sdf

models/filmsiren.py
bearprin/Neural-Singular-Hessian
import torch
import torch.nn as nn
import utils.utils as utils


def eikonal_loss(nonmnfld_grad, mnfld_grad, eikonal_type='abs'):
    # Compute the eikonal loss that penalises when ||grad(f)|| != 1 for points on and off the manifold
    # shape is (bs, num_points, dim=3) for both grads
    # Eikonal
    if nonmnfld_grad is not None and mnfld_grad is not None:
        all_grads = torch.cat([nonmnfld_grad, mnfld_grad], dim=-2)
    elif nonmnfld_grad is not None:
        all_grads = nonmnfld_grad
    elif mnfld_grad is not None:
        all_grads = mnfld_grad

    if eikonal_type == 'abs':
        eikonal_term = ((all_grads.norm(2, dim=2) - 1).abs()).mean()
    else:
        eikonal_term = ((all_grads.norm(2, dim=2) - 1).square()).mean()

    return eikonal_term


def relax_eikonal_loss(nonmnfld_grad, mnfld_grad, min=.8, max=0.1, eikonal_type='abs', udf=False):
    # Compute the eikonal loss that penalises when ||grad(f)|| != 1 for points on and off the manifold
    # shape is (bs, num_points, dim=3) for both grads
    # Eikonal
    if nonmnfld_grad is not None and mnfld_grad is not None:
        all_grads = torch.cat([nonmnfld_grad, mnfld_grad], dim=-2)
    elif nonmnfld_grad is not None:
        all_grads = nonmnfld_grad
    elif mnfld_grad is not None:
        all_grads = mnfld_grad

    grad_norm = all_grads.norm(2, dim=-1) + 1e-12
    if udf:
        pass
    else:
        term = torch.relu(-(grad_norm - min))
    if eikonal_type == 'abs':
        eikonal_term = term.abs().mean()
    else:
        eikonal_term = term.square().mean()
    return eikonal_term


def latent_rg_loss(latent_reg, device):
    # compute the VAE latent representation regularization loss
    if latent_reg is not None:
        reg_loss = latent_reg.mean()
    else:
        reg_loss = torch.tensor([0.0], device=device)

    return reg_loss


class MorseLoss(nn.Module):
    def __init__(self, weights=None, loss_type='siren_wo_n_w_morse', div_decay='none',
                 div_type='l1', bidirectional_morse=True, udf=False):
        super().__init__()
        if weights is None:
            weights = [3e3, 1e2, 1e2, 5e1, 1e2, 1e1]
        self.weights = weights  # sdf, intern, normal, eikonal, div
        self.loss_type = loss_type
        self.div_decay = div_decay
        self.div_type = div_type
        self.use_morse = True if 'morse' in self.loss_type else False
        self.bidirectional_morse = bidirectional_morse
        self.udf = udf

    def forward(self, output_pred, mnfld_points, nonmnfld_points, mnfld_n_gt=None, near_points=None):
        dims = mnfld_points.shape[-1]
        device = mnfld_points.device

        #########################################
        # Compute required terms
        #########################################

        non_manifold_pred = output_pred["nonmanifold_pnts_pred"]
        manifold_pred = output_pred["manifold_pnts_pred"]
        latent_reg = output_pred["latent_reg"]

        div_loss = torch.tensor([0.0], device=mnfld_points.device)
        morse_loss = torch.tensor([0.0], device=mnfld_points.device)
        curv_term = torch.tensor([0.0], device=mnfld_points.device)
        latent_reg_term = torch.tensor([0.0], device=mnfld_points.device)
        normal_term = torch.tensor([0.0], device=mnfld_points.device)

        # compute gradients for div (divergence), curl and curv (curvature)
        if manifold_pred is not None:
            mnfld_grad = utils.gradient(mnfld_points, manifold_pred)
        else:
            mnfld_grad = None

        nonmnfld_grad = utils.gradient(nonmnfld_points, non_manifold_pred)

        morse_nonmnfld_points = None
        morse_nonmnfld_grad = None
        if self.use_morse and near_points is not None:
            morse_nonmnfld_points = near_points
            morse_nonmnfld_grad = utils.gradient(near_points, output_pred['near_points_pred'])
        elif self.use_morse and near_points is None:
            morse_nonmnfld_points = nonmnfld_points
            morse_nonmnfld_grad = nonmnfld_grad

        if self.use_morse:
            nonmnfld_dx = utils.gradient(morse_nonmnfld_points, morse_nonmnfld_grad[:, :, 0])
            nonmnfld_dy = utils.gradient(morse_nonmnfld_points, morse_nonmnfld_grad[:, :, 1])

            mnfld_dx = utils.gradient(mnfld_points, mnfld_grad[:, :, 0])
            mnfld_dy = utils.gradient(mnfld_points, mnfld_grad[:, :, 1])
            if dims == 3:
                nonmnfld_dz = utils.gradient(morse_nonmnfld_points, morse_nonmnfld_grad[:, :, 2])
                nonmnfld_hessian_term = torch.stack((nonmnfld_dx, nonmnfld_dy, nonmnfld_dz), dim=-1)

                mnfld_dz = utils.gradient(mnfld_points, mnfld_grad[:, :, 2])
                mnfld_hessian_term = torch.stack((mnfld_dx, mnfld_dy, mnfld_dz), dim=-1)
            else:
                nonmnfld_hessian_term = torch.stack((nonmnfld_dx, nonmnfld_dy), dim=-1)
                mnfld_hessian_term = torch.stack((mnfld_dx, mnfld_dy), dim=-1)

            nonmnfld_det = torch.det(nonmnfld_hessian_term)
            mnfld_det = torch.det(mnfld_hessian_term)

            morse_mnfld = torch.tensor([0.0], device=mnfld_points.device)
            morse_nonmnfld = torch.tensor([0.0], device=mnfld_points.device)
            if self.div_type == 'l2':
                morse_nonmnfld = nonmnfld_det.square().mean()
                if self.bidirectional_morse:
                    morse_mnfld = mnfld_det.square().mean()
            elif self.div_type == 'l1':
                morse_nonmnfld = nonmnfld_det.abs().mean()
                # morse_nonmnfld = morse_nonmnfld_grad.norm(dim=-1).square().mean()
                # morse_nonmnfld = nonmnfld_hessian_term.norm(dim=[-1, -2]).square().mean()
                # nonmnfld_divergence = nonmnfld_dx[:, :, 0] + nonmnfld_dy[:, :, 1] + nonmnfld_dz[:, :, 2]
                # morse_nonmnfld = torch.clamp(torch.abs(nonmnfld_divergence), 0.1, 50).mean()
                if self.bidirectional_morse:
                    morse_mnfld = mnfld_det.abs().mean()

            morse_loss = 0.5 * (morse_nonmnfld + morse_mnfld)

        # latent regulariation for multiple shape learning
        latent_reg_term = latent_rg_loss(latent_reg, device)

        # normal term
        if mnfld_n_gt is not None:
            if 'igr' in self.loss_type:
                normal_term = ((mnfld_grad - mnfld_n_gt).abs()).norm(2, dim=1).mean()
            else:
                normal_term = (
                        1 - torch.abs(torch.nn.functional.cosine_similarity(mnfld_grad, mnfld_n_gt, dim=-1))).mean()

        # signed distance function term
        sdf_term = torch.abs(manifold_pred).mean()
        # sdf_term = (torch.abs(manifold_pred) * torch.exp(manifold_pred.abs())).mean()

        # eikonal term
        # eikonal_term = eikonal_loss(nonmnfld_grad, mnfld_grad=mnfld_grad, eikonal_type='abs')
        # Sometimes > relax may leading to bad results, use another type relax
        # eikonal_term = relax_eikonal_loss(None, mnfld_grad=mnfld_grad, udf=self.udf)
        eikonal_term = eikonal_loss(None, mnfld_grad=mnfld_grad, eikonal_type='abs')

        # inter term
        inter_term = torch.exp(-1e2 * torch.abs(non_manifold_pred)).mean()
        #########################################
        # Losses
        #########################################

        # losses used in the paper
        if self.loss_type == 'siren':  # SIREN loss
            loss = self.weights[0] * sdf_term + self.weights[1] * inter_term + \
                   self.weights[2] * normal_term + self.weights[3] * eikonal_term
        elif self.loss_type == 'siren_w_morse':
            loss = self.weights[0] * sdf_term + self.weights[1] * inter_term + \
                   self.weights[2] * normal_term + self.weights[3] * eikonal_term + \
                   self.weights[4] * morse_loss
        elif self.loss_type == 'siren_wo_n':  # SIREN loss without normal constraint
            self.weights[2] = 0
            loss = self.weights[0] * sdf_term + self.weights[1] * inter_term + self.weights[3] * eikonal_term
        elif self.loss_type == 'siren_wo_n_w_morse':
            self.weights[2] = 0
            loss = self.weights[0] * sdf_term + self.weights[1] * inter_term + self.weights[3] * eikonal_term + \
                   self.weights[5] * morse_loss
        elif self.loss_type == 'siren_wo_n_wo_e_wo_morse':
            loss = self.weights[0] * sdf_term + self.weights[1] * inter_term
        elif self.loss_type == 'igr':  # IGR loss
            self.weights[1] = 0
            loss = self.weights[0] * sdf_term + self.weights[2] * normal_term + self.weights[3] * eikonal_term
        elif self.loss_type == 'igr_wo_n':  # IGR without normals loss
            self.weights[1] = 0
            self.weights[2] = 0
            loss = self.weights[0] * sdf_term + self.weights[3] * eikonal_term
        elif self.loss_type == 'igr_wo_n_w_morse':
            self.weights[1] = 0
            self.weights[2] = 0
            loss = self.weights[0] * sdf_term + self.weights[3] * eikonal_term + self.weights[5] * morse_loss
        elif self.loss_type == 'siren_w_div':  # SIREN loss with divergence term
            loss = self.weights[0] * sdf_term + self.weights[1] * inter_term + \
                   self.weights[2] * normal_term + self.weights[3] * eikonal_term + \
                   self.weights[4] * div_loss
        elif self.loss_type == 'siren_wo_e_w_morse':
            self.weights[3] = 0
            self.weights[4] = 0
            loss = self.weights[0] * sdf_term + self.weights[1] * inter_term + \
                   self.weights[2] * normal_term + self.weights[5] * morse_loss
        elif self.loss_type == 'siren_wo_e_wo_n_w_morse':
            self.weights[2] = 0
            self.weights[3] = 0
            self.weights[4] = 0
            loss = self.weights[0] * sdf_term + self.weights[1] * inter_term + self.weights[5] * morse_loss
        elif self.loss_type == 'siren_wo_n_w_div':  # SIREN loss without normals and with divergence constraint
            loss = self.weights[0] * sdf_term + self.weights[1] * inter_term + self.weights[3] * eikonal_term + \
                   self.weights[4] * div_loss
        else:
            print(self.loss_type)
            raise Warning("unrecognized loss type")

        # If multiple surface reconstruction, then latent and latent_reg are defined so reg_term need to be used
        if latent_reg is not None:
            loss += self.weights[6] * latent_reg_term

        return {"loss": loss, 'sdf_term': sdf_term, 'inter_term': inter_term, 'latent_reg_term': latent_reg_term,
                'eikonal_term': eikonal_term, 'normals_loss': normal_term, 'div_loss': div_loss,
                'curv_loss': curv_term.mean(), 'morse_term': morse_loss}, mnfld_grad

    def update_morse_weight(self, current_iteration, n_iterations, params=None):
        # `params`` should be (start_weight, *optional middle, end_weight) where optional middle is of the form [percent, value]*
        # Thus (1e2, 0.5, 1e2 0.7 0.0, 0.0) means that the weight at [0, 0.5, 0.75, 1] of the training process, the weight should
        #   be [1e2,1e2,0.0,0.0]. Between these points, the weights change as per the div_decay parameter, e.g. linearly, quintic, step etc.
        #   Thus the weight stays at 1e2 from 0-0.5, decay from 1e2 to 0.0 from 0.5-0.75, and then stays at 0.0 from 0.75-1.

        if not hasattr(self, 'decay_params_list'):
            assert len(params) >= 2, params
            assert len(params[1:-1]) % 2 == 0
            self.decay_params_list = list(zip([params[0], *params[1:-1][1::2], params[-1]], [0, *params[1:-1][::2], 1]))

        curr = current_iteration / n_iterations
        we, e = min([tup for tup in self.decay_params_list if tup[1] >= curr], key=lambda tup: tup[1])
        w0, s = max([tup for tup in self.decay_params_list if tup[1] <= curr], key=lambda tup: tup[1])

        # Divergence term anealing functions
        if self.div_decay == 'linear':  # linearly decrease weight from iter s to iter e
            if current_iteration < s * n_iterations:
                self.weights[5] = w0
            elif current_iteration >= s * n_iterations and current_iteration < e * n_iterations:
                self.weights[5] = w0 + (we - w0) * (current_iteration / n_iterations - s) / (e - s)
            else:
                self.weights[5] = we
        elif self.div_decay == 'quintic':  # linearly decrease weight from iter s to iter e
            if current_iteration < s * n_iterations:
                self.weights[5] = w0
            elif current_iteration >= s * n_iterations and current_iteration < e * n_iterations:
                self.weights[5] = w0 + (we - w0) * (1 - (1 - (current_iteration / n_iterations - s) / (e - s)) ** 5)
            else:
                self.weights[5] = we
        elif self.div_decay == 'step':  # change weight at s
            if current_iteration < s * n_iterations:
                self.weights[5] = w0
            else:
                self.weights[5] = we
        elif self.div_decay == 'none':
            pass
        else:
            raise Warning("unsupported div decay value")

models/losses.py
bearprin/Neural-Singular-Hessian
import numpy as np


def normalize_coord(p, vol_range, plane='xz'):
    ''' Normalize coordinate to [0, 1] for sliding-window experiments

    Args:
        p (tensor): point
        vol_range (numpy array): volume boundary
        plane (str): feature type, ['xz', 'xy', 'yz'] - canonical planes; ['grid'] - grid volume
    '''
    p[:, 0] = (p[:, 0] - vol_range[0][0]) / (vol_range[1][0] - vol_range[0][0])
    p[:, 1] = (p[:, 1] - vol_range[0][1]) / (vol_range[1][1] - vol_range[0][1])
    p[:, 2] = (p[:, 2] - vol_range[0][2]) / (vol_range[1][2] - vol_range[0][2])

    if plane == 'xz':
        x = p[:, [0, 2]]
    elif plane =='xy':
        x = p[:, [0, 1]]
    elif plane =='yz':
        x = p[:, [1, 2]]
    else:
        x = p
    return x


def normalize_coordinate(p, plane='xz',bbox_size=1.0):
    ''' Normalize coordinate to [0, 1] for unit cube experiments

    Args:
        p (tensor): point
        padding (float): conventional padding paramter of ONet for unit cube, so [-0.5, 0.5] -> [-0.45, 0.45]
        plane (str): plane feature type, ['xz', 'xy', 'yz']
    '''
    if plane == 'xz':
        xy = p[:, :, [0, 2]]
    elif plane =='xy':
        xy = p[:, :, [0, 1]]
    else:
        xy = p[:, :, [1, 2]]

    xy_new = xy / bbox_size # (-0.5, 0.5)
    xy_new = xy_new + 0.5 # range (0, 1)

    # f there are outliers out of the range
    if xy_new.max() >= 1:
        xy_new[xy_new >= 1] = 1 - 10e-6
    if xy_new.min() < 0:
        xy_new[xy_new < 0] = 0.0
    return xy_new

def normalize_3d_coordinate(p, bbox_size=1.0):
    ''' Normalize coordinate to [0, 1] for unit cube experiments.
        Corresponds to our 3D model

    Args:
        p (tensor): point
        padding (float): conventional padding paramter of ONet for unit cube, so [-0.5, 0.5] -> [-0.55, 0.55]
    '''
    p_nor = p / bbox_size # (-0.5, 0.5)
    p_nor = p_nor + 0.5 # range (0, 1)
    # f there are outliers out of the range
    if p_nor.max() >= 1:
        p_nor[p_nor >= 1] = 1 - 10e-4
    if p_nor.min() < 0:
        p_nor[p_nor < 0] = 0.0
    return p_nor

def coordinate2index(x, reso, coord_type='2d'):
    ''' Normalize coordinate to [0, 1] for unit cube experiments.
        Corresponds to our 3D model

    Args:
        x (tensor): coordinate
        reso (int): defined resolution
        coord_type (str): coordinate type
    '''
    x = (x * reso).long()
    if coord_type == '2d': # plane
        index = x[:, :, 0] + reso * x[:, :, 1]
    elif coord_type == '3d': # grid
        index = x[:, :, 0] + reso * (x[:, :, 1] + reso * x[:, :, 2])
    index = index[:, None, :]
    return index

def coord2index(p, vol_range, reso=None, plane='xz'):
    ''' Normalize coordinate to [0, 1] for sliding-window experiments.
        Corresponds to our 3D model

    Args:
        p (tensor): points
        vol_range (numpy array): volume boundary
        reso (int): defined resolution
        plane (str): feature type, ['xz', 'xy', 'yz'] - canonical planes; ['grid'] - grid volume
    '''
    # normalize to [0, 1]
    x = normalize_coord(p, vol_range, plane=plane)

    if isinstance(x, np.ndarray):
        x = np.floor(x * reso).astype(int)
    else: #* pytorch tensor
        x = (x * reso).long()

    if x.shape[1] == 2:
        index = x[:, 0] + reso * x[:, 1]
        index[index > reso**2] = reso**2
    elif x.shape[1] == 3:
        index = x[:, 0] + reso * (x[:, 1] + reso * x[:, 2])
        index[index > reso**3] = reso**3

    return index[None]

models/ocn/common.py
bearprin/Neural-Singular-Hessian
import numpy as np
import torch
import torch.nn as nn
from torch import distributions as dist


class AbsLayer(nn.Module):
    def __init__(self):
        super(AbsLayer, self).__init__()

    def forward(self, x):
        return torch.abs(x)


def exists(val):
    return val is not None


def cast_tuple(val, repeat=1):
    return val if isinstance(val, tuple) else ((val,) * repeat)


class Decoder(nn.Module):
    def __init__(self, udf=False):
        super(Decoder, self).__init__()
        self.nl = nn.Identity() if not udf else AbsLayer()

    def forward(self, *args, **kwargs):
        res = self.fc_block(*args, **kwargs)
        res = self.nl(res)
        return res


class Modulator(nn.Module):
    def __init__(self, dim_in, dim_hidden, num_layers):
        super().__init__()
        self.layers = nn.ModuleList([])

        for ind in range(num_layers):
            is_first = ind == 0
            dim = dim_in if is_first else (dim_hidden + dim_in)

            self.layers.append(nn.Sequential(
                nn.Linear(dim, dim_hidden),
                nn.ReLU()
            ))

    def forward(self, z):
        x = z
        hiddens = []

        for layer in self.layers:
            x = layer(x)
            hiddens.append(x)
            x = torch.cat((x, z), dim=-1)

        return tuple(hiddens)


def maxpool(x, dim=-1, keepdim=False):
    out, _ = x.max(dim=dim, keepdim=keepdim)
    return out


class SimplePointnet(nn.Module):
    ''' PointNet-based encoder network. Based on: https://github.com/autonomousvision/occupancy_networks


    Args:
        c_dim (int): dimension of latent code c
        dim (int): input points dimension
        hidden_dim (int): hidden dimension of the network
    '''

    def __init__(self, c_dim=128, dim=3, hidden_dim=128):
        super().__init__()
        self.c_dim = c_dim

        self.fc_pos = nn.Linear(dim, 2 * hidden_dim)
        self.fc_0 = nn.Linear(2 * hidden_dim, hidden_dim)
        self.fc_1 = nn.Linear(2 * hidden_dim, hidden_dim)
        self.fc_2 = nn.Linear(2 * hidden_dim, hidden_dim)
        self.fc_3 = nn.Linear(2 * hidden_dim, hidden_dim)
        self.fc_mean = nn.Linear(hidden_dim, c_dim)
        self.fc_std = nn.Linear(hidden_dim, c_dim)
        torch.nn.init.constant_(self.fc_mean.weight, 0)
        torch.nn.init.constant_(self.fc_mean.bias, 0)

        torch.nn.init.constant_(self.fc_std.weight, 0)
        torch.nn.init.constant_(self.fc_std.bias, -10)

        self.actvn = nn.ReLU()
        self.pool = maxpool

    def forward(self, p):
        net = self.fc_pos(p)
        net = self.fc_0(self.actvn(net))
        pooled = self.pool(net, dim=1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=2)

        net = self.fc_1(self.actvn(net))
        pooled = self.pool(net, dim=1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=2)

        net = self.fc_2(self.actvn(net))
        pooled = self.pool(net, dim=1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=2)

        net = self.fc_3(self.actvn(net))

        net = self.pool(net, dim=1)

        c_mean = self.fc_mean(self.actvn(net))
        c_std = self.fc_std(self.actvn(net))

        return c_mean, c_std


class Network(nn.Module):
    def __init__(self, latent_size=0, in_dim=3, decoder_hidden_dim=256, nl='sine', decoder_n_hidden_layers=4,
                 init_type='siren', sphere_init_params=[1.6, 1.0], udf=False, vae=False):
        super().__init__()
        self.latent_size = latent_size
        self.vae = vae

        self.encoder = SimplePointnet(c_dim=latent_size, dim=3) if latent_size > 0 and vae == True else None
        # self.encoder = ResnetPointnet(c_dim=latent_size, dim=3) if latent_size > 0 and vae == True else None
        self.modulator = Modulator(
            dim_in=latent_size,
            dim_hidden=decoder_hidden_dim,
            num_layers=decoder_n_hidden_layers + 1  # +1 for input layer
        ) if latent_size > 0 else None

        self.init_type = init_type
        self.decoder = Decoder(udf=udf)
        self.decoder.fc_block = FCBlock(in_dim, 1, num_hidden_layers=decoder_n_hidden_layers,
                                        hidden_features=decoder_hidden_dim,
                                        outermost_linear=True, nonlinearity=nl, init_type=init_type,
                                        sphere_init_params=sphere_init_params)  # SIREN decoder

    def forward(self, non_mnfld_pnts, mnfld_pnts=None, near_points=None, only_nonmnfld=False, latent=None):
        batch_size = non_mnfld_pnts.shape[0]
        if self.latent_size > 0 and self.encoder is not None:
            # encoder
            q_latent_mean, q_latent_std = self.encoder(mnfld_pnts)
            q_z = dist.Normal(q_latent_mean, torch.exp(q_latent_std))
            latent = q_z.rsample()
            latent_reg = 1.0e-3 * (q_latent_mean.abs().mean(dim=-1) + (q_latent_std + 1).abs().mean(dim=-1))
            # modulate
            modulate = exists(self.modulator)
            mods = self.modulator(latent) if modulate else None
        elif self.latent_size > 0:
            modulate = exists(self.modulator)
            mods = self.modulator(latent) if modulate else None
            latent_reg = 1e-3 * latent.norm(-1).mean()
        else:
            latent = None
            latent_reg = None
            mods = None
        if mnfld_pnts is not None and not only_nonmnfld:
            manifold_pnts_pred = self.decoder(mnfld_pnts, mods)
        else:
            manifold_pnts_pred = None
        nonmanifold_pnts_pred = self.decoder(non_mnfld_pnts, mods)

        near_points_pred = None
        if near_points is not None:
            near_points_pred = self.decoder(near_points, mods)

        return {"manifold_pnts_pred": manifold_pnts_pred,
                "nonmanifold_pnts_pred": nonmanifold_pnts_pred,
                'near_points_pred': near_points_pred,
                "latent_reg": latent_reg,
                }

    def get_latent_mods(self, mnfld_pnts=None, latent=None, rand_predict=True):
        mods = None
        if self.vae:
            if rand_predict:
                q_latent_mean, q_latent_std = self.encoder(mnfld_pnts)
                q_z = dist.Normal(q_latent_mean, torch.exp(q_latent_std))
                latent = q_z.rsample()
            else:
                latent, _ = self.encoder(mnfld_pnts)
        modulate = exists(self.modulator)
        mods = self.modulator(latent) if modulate else None
        return mods


class FCBlock(nn.Module):
    '''A fully connected neural network that also allows swapping out the weights when used with a hypernetwork.
    Can be used just as a normal neural network though, as well.
    '''

    def __init__(self, in_features, out_features, num_hidden_layers, hidden_features,
                 outermost_linear=False, nonlinearity='sine', init_type='siren',
                 sphere_init_params=[1.6, 1.0]):
        super().__init__()
        print("decoder initialising with {} and {}".format(nonlinearity, init_type))

        self.first_layer_init = None
        self.sphere_init_params = sphere_init_params
        self.init_type = init_type

        nl_dict = {'sine': Sine(), 'relu': nn.ReLU(inplace=True), 'softplus': nn.Softplus(beta=100),
                   'tanh': nn.Tanh(), 'sigmoid': nn.Sigmoid()}
        nl = nl_dict[nonlinearity]

        self.net = []
        self.net.append(nn.Sequential(nn.Linear(in_features, hidden_features), nl))

        for i in range(num_hidden_layers):
            self.net.append(nn.Sequential(nn.Linear(hidden_features, hidden_features), nl))

        if outermost_linear:
            self.net.append(nn.Sequential(nn.Linear(hidden_features, out_features)))
        else:
            self.net.append(nn.Sequential(nn.Linear(hidden_features, out_features), nl))

        self.net = nn.Sequential(*self.net)

        if init_type == 'siren':
            self.net.apply(sine_init)
            self.net[0].apply(first_layer_sine_init)

        elif init_type == 'geometric_sine':
            self.net.apply(geom_sine_init)
            self.net[0].apply(first_layer_geom_sine_init)
            self.net[-2].apply(second_last_layer_geom_sine_init)
            self.net[-1].apply(last_layer_geom_sine_init)

        elif init_type == 'mfgi':
            self.net.apply(geom_sine_init)
            self.net[0].apply(first_layer_mfgi_init)
            self.net[1].apply(second_layer_mfgi_init)
            self.net[-2].apply(second_last_layer_geom_sine_init)
            self.net[-1].apply(last_layer_geom_sine_init)

        elif init_type == 'geometric_relu':
            self.net.apply(geom_relu_init)
            self.net[-1].apply(geom_relu_last_layers_init)

    def forward(self, coords, mods=None):
        mods = cast_tuple(mods, len(self.net))
        x = coords

        for layer, mod in zip(self.net, mods):
            x = layer(x)
            if exists(mod):
                if mod.shape[1] != 1:
                    mod = mod[:, None, :]
                x = x * mod
        if mods[0] is not None:
            x = self.net[-1](x)  # last layer

        if self.init_type == 'mfgi' or self.init_type == 'geometric_sine':
            radius, scaling = self.sphere_init_params
            output = torch.sign(x) * torch.sqrt(x.abs() + 1e-8)
            output -= radius  # 1.6
            output *= scaling  # 1.0

        return x


class Sine(nn.Module):
    def forward(self, input):
        # See SIREN paper sec. 3.2, final paragraph, and supplement Sec. 1.5 for discussion of factor 30
        return torch.sin(30 * input)


################################# SIREN's initialization ###################################
def sine_init(m):
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_input = m.weight.size(-1)
            # See SIREN paper supplement Sec. 1.5 for discussion of factor 30
            m.weight.uniform_(-np.sqrt(6 / num_input) / 30, np.sqrt(6 / num_input) / 30)


def first_layer_sine_init(m):
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_input = m.weight.size(-1)
            # See SIREN paper sec. 3.2, final paragraph, and supplement Sec. 1.5 for discussion of factor 30
            m.weight.uniform_(-1 / num_input, 1 / num_input)


################################# sine geometric initialization ###################################

def geom_sine_init(m):
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_output = m.weight.size(0)
            m.weight.uniform_(-np.sqrt(3 / num_output), np.sqrt(3 / num_output))
            m.bias.uniform_(-1 / (num_output * 1000), 1 / (num_output * 1000))
            m.weight.data /= 30
            m.bias.data /= 30


def first_layer_geom_sine_init(m):
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_output = m.weight.size(0)
            m.weight.uniform_(-np.sqrt(3 / num_output), np.sqrt(3 / num_output))
            m.bias.uniform_(-1 / (num_output * 1000), 1 / (num_output * 1000))
            m.weight.data /= 30
            m.bias.data /= 30


def second_last_layer_geom_sine_init(m):
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_output = m.weight.size(0)
            assert m.weight.shape == (num_output, num_output)
            m.weight.data = 0.5 * np.pi * torch.eye(num_output) + 0.001 * torch.randn(num_output, num_output)
            m.bias.data = 0.5 * np.pi * torch.ones(num_output, ) + 0.001 * torch.randn(num_output)
            m.weight.data /= 30
            m.bias.data /= 30


def last_layer_geom_sine_init(m):
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_input = m.weight.size(-1)
            assert m.weight.shape == (1, num_input)
            assert m.bias.shape == (1,)
            # m.weight.data = -1 * torch.ones(1, num_input) + 0.001 * torch.randn(num_input)
            m.weight.data = -1 * torch.ones(1, num_input) + 0.00001 * torch.randn(num_input)
            m.bias.data = torch.zeros(1) + num_input


################################# multi frequency geometric initialization ###################################
periods = [1, 30]  # Number of periods of sine the values of each section of the output vector should hit
# periods = [1, 60] # Number of periods of sine the values of each section of the output vector should hit
portion_per_period = np.array([0.25, 0.75])  # Portion of values per section/period


def first_layer_mfgi_init(m):
    global periods
    global portion_per_period
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_input = m.weight.size(-1)
            num_output = m.weight.size(0)
            num_per_period = (portion_per_period * num_output).astype(int)  # Number of values per section/period
            assert len(periods) == len(num_per_period)
            assert sum(num_per_period) == num_output
            weights = []
            for i in range(0, len(periods)):
                period = periods[i]
                num = num_per_period[i]
                scale = 30 / period
                weights.append(torch.zeros(num, num_input).uniform_(-np.sqrt(3 / num_input) / scale,
                                                                    np.sqrt(3 / num_input) / scale))
            W0_new = torch.cat(weights, axis=0)
            m.weight.data = W0_new


def second_layer_mfgi_init(m):
    global portion_per_period
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_input = m.weight.size(-1)
            assert m.weight.shape == (num_input, num_input)
            num_per_period = (portion_per_period * num_input).astype(int)  # Number of values per section/period
            k = num_per_period[0]  # the portion that only hits the first period
            # W1_new = torch.zeros(num_input, num_input).uniform_(-np.sqrt(3 / num_input), np.sqrt(3 / num_input) / 30) * 0.00001
            W1_new = torch.zeros(num_input, num_input).uniform_(-np.sqrt(3 / num_input),
                                                                np.sqrt(3 / num_input) / 30) * 0.0005
            W1_new_1 = torch.zeros(k, k).uniform_(-np.sqrt(3 / num_input) / 30, np.sqrt(3 / num_input) / 30)
            W1_new[:k, :k] = W1_new_1
            m.weight.data = W1_new


################################# geometric initialization used in SAL and IGR ###################################
def geom_relu_init(m):
    with torch.no_grad():
        if hasattr(m, 'weight'):
            out_dims = m.out_features

            m.weight.normal_(mean=0.0, std=np.sqrt(2) / np.sqrt(out_dims))
            m.bias.data = torch.zeros_like(m.bias.data)


def geom_relu_last_layers_init(m):
    radius_init = 1
    with torch.no_grad():
        if hasattr(m, 'weight'):
            num_input = m.weight.size(-1)
            m.weight.normal_(mean=np.sqrt(np.pi) / np.sqrt(num_input), std=0.00001)
            m.bias.data = torch.Tensor([-radius_init])

models/overfit_network.py
bearprin/Neural-Singular-Hessian
import torch.nn as nn

from .filmsiren import FilmSiren
from .convolutionalfeature import ConvolutionalFeature


class ShapeNetwork(nn.Module):
    def __init__(self, decoder_hidden_dim=256, decoder_n_hidden_layers=5):
        super().__init__()
        self.encoder = ConvolutionalFeature()
        self.decoder = FilmSiren(hidden_size=decoder_hidden_dim, n_layers=decoder_n_hidden_layers)

    def forward(self, non_mnfld_pnts=None, mnfld_pnts=None, near_points=None):
        return self.forward_dense(non_mnfld_pnts, mnfld_pnts, near_points)

    def forward_dense(self, non_mnfld_pnts=None, mnfld_pnts=None, near_points=None):
        latent_reg = None
        # mnfld
        global_feat = self.encoder.encode(mnfld_pnts)
        mnfld_pnts_feat = self.encoder.query_feature(global_feat, mnfld_pnts)
        manifold_pnts_pred = self.decoder(mnfld_pnts, mnfld_pnts_feat)
        # nonfld
        nonmanifold_pnts_pred = None
        if non_mnfld_pnts is not None:
            non_mnfld_pnts_feat = self.encoder.query_feature(global_feat, non_mnfld_pnts)
            nonmanifold_pnts_pred = self.decoder(non_mnfld_pnts, non_mnfld_pnts_feat)

        near_points_pred = None
        if near_points is not None:
            near_points_feat = self.encoder.query_feature(global_feat, near_points)
            near_points_pred = self.decoder(near_points, near_points_feat)

        return {"manifold_pnts_pred": manifold_pnts_pred,
                "nonmanifold_pnts_pred": nonmanifold_pnts_pred,
                'near_points_pred': near_points_pred,
                "latent_reg": latent_reg,
                }

models/shape_network.py
bearprin/Neural-Singular-Hessian
import torch
import torch.utils.data as data
import numpy as np
import os


class DFaustDataSet(data.Dataset):

    def __init__(self, dataset_path, split_path, points_batch=10000, d_in=3, with_normals=False, grid_range=1.1,
                 val=False, part_i=None):
        # dataset_path, gt_path, scan_path
        self.grid_range = grid_range

        print('getting npy files')
        dataset_path = os.path.expanduser(dataset_path)
        self.dataset_path = dataset_path
        assert os.path.exists(dataset_path)
        split_path = os.path.expanduser(split_path)
        assert os.path.exists(split_path), split_path
        with open(split_path, "r") as f:
            split = f.readlines()
        # self.npyfiles_mnfld = get_instance_filenames(dataset_path, split)  # list of filenames
        self.npyfiles_mnfld = open(os.path.join(dataset_path, split_path), 'r+').readlines()
        print(len(self.npyfiles_mnfld), 'npy files found')
        # splitting the dataset into 15 parts for parallel training
        # idx = int(np.ceil(len(self.npyfiles_mnfld) / 15))
        # if part_i is not None:
        #     self.npyfiles_mnfld = self.npyfiles_mnfld[part_i * idx:(part_i + 1) * idx]
        self.n_points = points_batch
        self.with_normals = with_normals
        self.d_in = d_in

    def load_points(self, index):
        return np.load(os.path.join(self.dataset_path, self.npyfiles_mnfld[index].strip())) # (250000, 6) which has xyz, normal xyz
        # return np.load(self.npyfiles_mnfld[index])  # (N, 6) which has xyz, normal xyz

    def get_info(self, index):
        shape_name, pose, tag = self.npyfiles_mnfld[index].split('/')[-3:]
        return shape_name, pose, tag[:tag.find('.npy')]

    def __getitem__(self, index):
        # index = torch.tensor([0]).long()
        point_set_mnlfld = torch.from_numpy(self.load_points(index)).float()  # (250000, 6) which has xyz, normal xyz

        random_idx = torch.randperm(point_set_mnlfld.shape[0])[:self.n_points]
        point_set_mnlfld = torch.index_select(point_set_mnlfld, 0, random_idx)  # (pnts, 6)

        mnfld_points = point_set_mnlfld[:, :self.d_in]

        if self.with_normals:
            normals = point_set_mnlfld[:, -self.d_in:]  # todo adjust to case when we get no sigmas
        else:
            normals = torch.empty(0)

        # Q_far
        nonmnfld_points = np.random.uniform(-self.grid_range, self.grid_range,
                                            size=(self.n_points, 3)).astype(np.float32)  # (n_points, 3)
        nonmnfld_points = torch.from_numpy(nonmnfld_points).float()

        # Q_near
        dist = torch.cdist(mnfld_points, mnfld_points)
        sigmas = torch.topk(dist, k=51, dim=1, largest=False)[0][:, -1:]  # (n_points, 1)
        near_points = (mnfld_points + sigmas * torch.randn(mnfld_points.shape[0],
                                                           mnfld_points.shape[1]))
        return {'mnfld_points': mnfld_points, 'mnfld_n': normals, 'nonmnfld_points': nonmnfld_points,
                'near_points': near_points, 'indices': index, 'name': self.npyfiles_mnfld[index]}

    def __len__(self):
        return len(self.npyfiles_mnfld)


def get_instance_filenames(base_dir, split):
    npyfiles = []
    for line in split:
        line = line.strip()
        npyfiles.append(os.path.join(base_dir, line))
    return npyfiles

shapespace/dfaust_dataset.py
bearprin/Neural-Singular-Hessian
import os
import sys
import time

import numpy as np
import tqdm

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
from torch.utils.data import DataLoader
import pytorch_lightning as pl
from pytorch_lightning.callbacks import LearningRateMonitor
from pytorch_lightning.callbacks.progress import TQDMProgressBar

from shapespace.dfaust_dataset import DFaustDataSet
from models import ShapeNetwork, MorseLoss
import utils.utils as utils
import shapespace.shapespace_dfaust_args as shapespace_dfaust_args

# get training parameters
args = shapespace_dfaust_args.get_args()
logdir = os.path.join(args.logdir)
os.makedirs(logdir, exist_ok=True)
log_file, log_writer_train, log_writer_test, model_outdir = utils.setup_logdir(logdir, args)

project_dir = os.path.join(logdir, "trained_models")

os.system('cp %s %s' % (__file__, logdir))  # backup the current training file
os.system('cp %s %s' % ('../models/new_network.py', logdir))  # backup the models files
os.system('cp %s %s' % ('../models/convolutionalfeature.py', logdir))  # backup the models files
os.system('cp %s %s' % ('../models/losses.py', logdir))  # backup the models files


class DFaustDataModule(pl.LightningDataModule):
    def __init__(self, args):
        super().__init__()
        self.args = args

    def setup(self, stage=None):
        self.train_set = DFaustDataSet(self.args.dataset_path, self.args.test_split_path,
                                       # gt_path=self.args.gt_path,
                                       # scan_path=self.args.scan_path, \
                                       with_normals=False, part_i=args.test_part_i)

    def train_dataloader(self):
        return DataLoader(self.train_set, batch_size=self.args.batch_size, num_workers=32, persistent_workers=True,
                          shuffle=False,
                          pin_memory=True)


class BaseTrainer(pl.LightningModule):
    def __init__(self, args):
        super(BaseTrainer, self).__init__()
        self.args = args
        self.learning_rate = args.test_lr
        self.net = ShapeNetwork(decoder_hidden_dim=args.decoder_hidden_dim,
                                decoder_n_hidden_layers=args.decoder_n_hidden_layers)
        self.criterion = MorseLoss(weights=args.test_loss_weights, loss_type=args.test_loss_type,
                                   div_decay=args.morse_decay,
                                   div_type=args.morse_type, bidirectional_morse=args.bidirectional_morse)

        self.args.morse_near = False
        self.automatic_optimization = False

    def training_step(self, batch, batch_idx):
        self.net.train()
        mnfld_points, _, nonmnfld_points, near_points, indices, name = batch['mnfld_points'], batch[
            'mnfld_n'], batch['nonmnfld_points'], batch['near_points'], batch['indices'], batch['name']
        # load pretrained model
        print('load pretrained model')
        self.load_state_dict(
            torch.load(os.path.join(project_dir, 'last.ckpt'), map_location=mnfld_points.device)['state_dict'])
        optimizer = torch.optim.Adam(self.net.parameters(), lr=self.learning_rate, amsgrad=True)
        optimizer.load_state_dict(torch.load(os.path.join(project_dir, 'last.ckpt'))['optimizer_states'][0])
        # save the pre_optim state
        self.generate_mesh(mnfld_points[:1], name, dir_name='pre_optim')
        mnfld_points.requires_grad_()
        nonmnfld_points.requires_grad_()
        near_points.requires_grad_()
        for i in tqdm.tqdm(range(args.test_iter)):
            # for i in tqdm.tqdm(range(50)):
            self.net.zero_grad()
            # decay lr
            if i == args.test_iter // 2:
                for g in optimizer.param_groups:
                    g['lr'] = self.learning_rate / 10
            output_pred = self.net(nonmnfld_points, mnfld_points,
                                   near_points=near_points if self.args.morse_near else None)

            loss_dict, _ = self.criterion(output_pred, mnfld_points, nonmnfld_points,
                                          near_points=near_points if args.morse_near else None)
            loss_dict["loss"].backward()
            if args.grad_clip_norm > 0:
                torch.nn.utils.clip_grad_norm_(self.net.parameters(), args.grad_clip_norm)
            optimizer.step()
            # self.optimizers().step()
            # if self.local_rank == 0 and self.global_rank == 0 and i % 50 == 0:
            #     weights = self.criterion.weights
            #     utils.log_string("Weights: {}".format(weights), log_file)
            #     utils.log_string('Epoch: {}, Loss: {:.5f} = L_Mnfld: {:.5f} + '
            #                      'L_NonMnfld: {:.5f} + L_Nrml: {:.10f} + L_Eknl: {:.5f} + L_Div: {:.5f} + L_Morse: {:.5f} + L_Latent: {:.10f}'.format(
            #         self.current_epoch, loss_dict["loss"].item(), weights[0] * loss_dict["sdf_term"].item(),
            #                                                       weights[1] * loss_dict["inter_term"].item(),
            #                                                       weights[2] * loss_dict["normals_loss"].item(),
            #                                                       weights[3] * loss_dict["eikonal_term"].item(),
            #                                                       weights[4] * loss_dict["div_loss"].item(),
            #                                                       weights[5] * loss_dict['morse_term'].item(),
            #                                                       weights[6] * loss_dict['latent_reg_term'].item()),
            #         log_file)
            #     utils.log_string('Unweighted L_s : L_Mnfld: {:.5f},  '
            #                      'L_NonMnfld: {:.5f},  L_Nrml: {:.10f},  L_Eknl: {:.5f}, L_Morse: {:.5f}, L_Latent: {:.10f}'.format(
            #         loss_dict["sdf_term"].item(), loss_dict["inter_term"].item(),
            #         loss_dict["normals_loss"].item(), loss_dict["eikonal_term"].item(),
            #         loss_dict['morse_term'].item(), loss_dict['latent_reg_term'].item()),
            #         log_file)
            #     self.generate_mesh(mnfld_points[:1], name, dir_name='iter', i=i)
        return {'mnfld': mnfld_points[:1], 'name': name}

    def generate_mesh(self, mnfld, name, dir_name='', i=None):
        cat = '/'.join(name[0].split('/')[-3:-1])
        name = os.path.splitext(name[0].split('/')[-1])[0]
        self.net.eval()
        if self.global_rank == 0 and self.local_rank == 0:
            with torch.no_grad():
                t0 = time.time()
                out_dir = "{}/vis_results_test_{}/{}".format(args.logdir, dir_name, cat)
                out_dir_pts = "{}/vis_results_test_pts_{}/{}".format(args.logdir, dir_name, cat)
                os.makedirs(out_dir, exist_ok=True)
                os.makedirs(out_dir_pts, exist_ok=True)
                global_feat = self.net.encoder.encode(mnfld)
                try:
                    pred_mesh = utils.implicit2mesh(decoder=self.net, mods=None, feat=global_feat,
                                                    grid_res=512,
                                                    get_mesh=True, device=next(self.net.parameters()).device)
                    if i is None:
                        pred_mesh.export(os.path.join(out_dir, "{}.ply".format(name)))
                    else:
                        pred_mesh.export(os.path.join(out_dir, "{}_iter_{}.ply".format(name, i)))
                    np.savetxt('{}/{}.xyz'.format(out_dir_pts, name), mnfld[0].cpu().numpy())
                except Exception as e:
                    print('Can not plot')
                    print(e)
                print('Plot took {:.3f}s'.format(time.time() - t0))

    def training_step_end(self, step_output):
        mnfld = step_output['mnfld']
        name = step_output['name']
        self.generate_mesh(mnfld, name, dir_name='res')

    def configure_optimizers(self):
        # Setup Adam optimizers
        optimizer = torch.optim.Adam(self.net.parameters(), lr=self.learning_rate, amsgrad=True)
        return optimizer


lr_monitor = LearningRateMonitor(logging_interval='step')
pl.seed_everything(args.seed, workers=True)
trainer = pl.Trainer(
    max_epochs=180,
    # auto_select_gpus=True,
    accelerator='gpu',
    devices=1,
    callbacks=[TQDMProgressBar(refresh_rate=10), lr_monitor],
    benchmark=True,
    deterministic=False,
    logger=True,
)
base_trainer = BaseTrainer(args)
dm = DFaustDataModule(args)
trainer.fit(base_trainer, dm, ckpt_path=os.path.join(project_dir, 'last.ckpt') if os.path.exists(
    os.path.join(project_dir, 'last.ckpt')) else None)

shapespace/pl_conv_finetune_shapespace.py
bearprin/Neural-Singular-Hessian
import os
import sys
import time

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
from torch.utils.data import DataLoader
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
from pytorch_lightning.callbacks.progress import TQDMProgressBar
from pytorch_lightning.strategies import DDPStrategy

from shapespace.dfaust_dataset import DFaustDataSet
from models import ShapeNetwork, MorseLoss
import utils.utils as utils
import shapespace.shapespace_dfaust_args as shapespace_dfaust_args

# get training parameters
args = shapespace_dfaust_args.get_args()
logdir = os.path.join(args.logdir)
os.makedirs(logdir, exist_ok=True)
log_file, log_writer_train, log_writer_test, model_outdir = utils.setup_logdir(logdir, args)

os.makedirs(os.path.join(logdir, "trained_models"), exist_ok=True)
project_dir = os.path.join(logdir, "trained_models")

os.system('cp %s %s' % (__file__, logdir))  # backup the current training file
os.system('cp %s %s' % ('../models/new_network.py', logdir))  # backup the models files
os.system('cp %s %s' % ('../models/convolutionalfeature.py', logdir))  # backup the models files
os.system('cp %s %s' % ('../models/losses.py', logdir))  # backup the models files


class DFaustDataModule(pl.LightningDataModule):
    def __init__(self, args):
        super().__init__()
        self.args = args

    def setup(self, stage=None):
        self.train_set = DFaustDataSet(self.args.dataset_path, self.args.train_split_path,
                                       with_normals=args.with_normal)

    def train_dataloader(self):
        return DataLoader(self.train_set, batch_size=self.args.batch_size, num_workers=32, persistent_workers=True,
                          pin_memory=True)


class BaseTrainer(pl.LightningModule):
    def __init__(self, args):
        super(BaseTrainer, self).__init__()
        self.args = args
        self.learning_rate = args.lr
        self.net = ShapeNetwork(decoder_hidden_dim=args.decoder_hidden_dim,
                              decoder_n_hidden_layers=args.decoder_n_hidden_layers)

        self.criterion = MorseLoss(weights=args.loss_weights, loss_type=args.loss_type, div_decay=args.morse_decay,
                                   div_type=args.morse_type, bidirectional_morse=args.bidirectional_morse)

    def training_step(self, batch, batch_idx):
        self.net.train()
        self.net.zero_grad(set_to_none=True)
        mnfld_points, mnfld_n_gt, nonmnfld_points, near_points, indices = batch['mnfld_points'], batch[
            'mnfld_n'], batch['nonmnfld_points'], batch['near_points'], batch['indices']
        mnfld_points.requires_grad_()
        nonmnfld_points.requires_grad_()
        near_points.requires_grad_()

        output_pred = self.net(nonmnfld_points, mnfld_points, near_points=near_points if self.args.morse_near else None)

        loss_dict, _ = self.criterion(output_pred, mnfld_points, nonmnfld_points,
                                      near_points=near_points if args.morse_near else None)
        for key, value in loss_dict.items():
            self.log(key, value, on_step=True, logger=True)
        self.log('total_loss', loss_dict['loss'], on_step=True, logger=True, prog_bar=True)
        if self.local_rank == 0 and self.global_rank == 0 and batch_idx % 30 == 0:
            weights = self.criterion.weights
            utils.log_string("Weights: {}".format(weights), log_file)
            utils.log_string('Epoch: {}, Loss: {:.5f} = L_Mnfld: {:.5f} + '
                             'L_NonMnfld: {:.5f} + L_Nrml: {:.5f} + L_Eknl: {:.5f} + L_Div: {:.5f} + L_Morse: {:.5f} + L_Latent: {:.10f}'.format(
                self.current_epoch, loss_dict["loss"].item(), weights[0] * loss_dict["sdf_term"].item(),
                                                              weights[1] * loss_dict["inter_term"].item(),
                                                              weights[2] * loss_dict["normals_loss"].item(),
                                                              weights[3] * loss_dict["eikonal_term"].item(),
                                                              weights[4] * loss_dict["div_loss"].item(),
                                                              weights[5] * loss_dict['morse_term'].item(),
                                                              weights[6] * loss_dict['latent_reg_term'].item()),
                log_file)
            utils.log_string('Unweighted L_s : L_Mnfld: {:.5f},  '
                             'L_NonMnfld: {:.5f},  L_Nrml: {:.5f},  L_Eknl: {:.5f}, L_Morse: {:.5f}, L_Latent: {:.10f}'.format(
                loss_dict["sdf_term"].item(), loss_dict["inter_term"].item(),
                loss_dict["normals_loss"].item(), loss_dict["eikonal_term"].item(),
                loss_dict['morse_term'].item(), loss_dict['latent_reg_term'].item()),
                log_file)
        return {'loss': loss_dict['loss'], 'mnfld': mnfld_points[:1]}

    def training_epoch_end(self, outputs):
        mnfld = outputs[0]['mnfld']
        self.net.eval()
        if self.global_rank == 0 and self.local_rank == 0:
            with torch.no_grad():
                t0 = time.time()
                out_dir = "{}/vis_results/".format(args.logdir)
                os.makedirs(out_dir, exist_ok=True)
                global_feat = self.net.encoder.encode(mnfld)
                try:
                    pred_mesh = utils.implicit2mesh(decoder=self.net, mods=None, feat=global_feat,
                                                    grid_res=128,
                                                    get_mesh=True, device=next(self.net.parameters()).device)
                    pred_mesh.export(os.path.join(out_dir, "pred_mesh_{}.ply".format(self.current_epoch)))
                except Exception as e:
                    print('Can not plot')
                    print(e)
                print('Plot took {:.3f}s'.format(time.time() - t0))
        # update weights
        curr_epoch = self.current_epoch
        self.criterion.update_morse_weight(curr_epoch, self.args.num_epochs, self.args.decay_params)

    def configure_optimizers(self):
        # Setup Adam optimizers
        optimizer = torch.optim.Adam(self.trainer.model.parameters(), lr=self.learning_rate, amsgrad=True)
        lr_sch = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=150 * 10, T_mult=2, eta_min=1e-6)

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": lr_sch,
                "interval": "step",
                "frequency": 1,
            },
        }


check_callback = ModelCheckpoint(
    dirpath=project_dir,
    filename='model-{epoch:02d}',
    save_top_k=-1,
    save_last=True
)
lr_monitor = LearningRateMonitor(logging_interval='step')
pl.seed_everything(args.seed, workers=True)
trainer = pl.Trainer(gradient_clip_val=args.grad_clip_norm,
                     max_epochs=args.num_epochs,
                     auto_select_gpus=True,
                     accelerator='gpu',
                     strategy=DDPStrategy(find_unused_parameters=False),
                     devices=8,
                     callbacks=[check_callback, TQDMProgressBar(refresh_rate=10), lr_monitor],
                     accumulate_grad_batches=8,
                     benchmark=True,
                     deterministic=False,
                     logger=True,
                     )
base_trainer = BaseTrainer(args)
dm = DFaustDataModule(args)

trainer.fit(base_trainer, dm, ckpt_path=os.path.join(project_dir, 'last.ckpt') if os.path.exists(
    os.path.join(project_dir, 'last.ckpt')) else None)

shapespace/pl_conv_train_shapespace.py
bearprin/Neural-Singular-Hessian
import argparse
import torch
import os
import numpy as np

import argparse


def add_args(parser):
    parser.add_argument('--logdir', type=str,
                        default='./log_conv_all_half_eikonal_1e-4_amsgrad_200_epoch_cos_1e-6_1500_2/',
                        help='log directory')
    parser.add_argument('--model_name', type=str, default='model', help='trained model name')
    parser.add_argument('--seed', type=int, default=3627473, help='random seed')
    parser.add_argument('--n_points', type=int, default=10000, help='number of points in each point cloud')
    parser.add_argument('--grid_res', type=int, default=128, help='uniform grid resolution')
    parser.add_argument('--with_normal', type=bool, default=False)

    # training parameters
    parser.add_argument('--dataset_path', type=str, default='/home/bearprin/DualOctreeGNN/data/dfaust/dataset',
                        help='path to dataset folder')
    parser.add_argument('--test_part_i', type=int, default=None)
    parser.add_argument('--train_split_path', type=str,
                        default='/home/bearprin/DualOctreeGNN/data/dfaust/filelist/train.txt',
                        help='path to split file')
    parser.add_argument('--test_split_path', type=str,
                        default='/home/bearprin/DualOctreeGNN/data/dfaust/filelist/test.txt',
                        help='path to split file')
    parser.add_argument('--num_epochs', type=int, default=200)
    parser.add_argument('--lr', type=float, default=1e-4, help='initial learning rate')
    parser.add_argument('--test_lr', type=float, default=1e-5, help='initial learning rate')
    parser.add_argument('--test_iter', type=int, default=3000, help='test_optim_iter')
    parser.add_argument('--grad_clip_norm', type=float, default=10.0, help='Value to clip gradients to')
    parser.add_argument('--batch_size', type=int, default=1, help='number of samples in a minibatch')

    # Network architecture and loss
    parser.add_argument('--init_type', type=str, default='siren',
                        help='initialization type siren | geometric_sine | geometric_relu | mfgi')
    parser.add_argument('--decoder_hidden_dim', type=int, default=256, help='length of decoder hidden dim')
    parser.add_argument('--decoder_n_hidden_layers', type=int, default=5, help='number of decoder hidden layers')
    parser.add_argument('--nl', type=str, default='sine', help='type of non linearity sine | relu')
    parser.add_argument('--sphere_init_params', nargs='+', type=float, default=[1.6, 0.1],
                        help='radius and scaling')
    parser.add_argument('--loss_type', type=str, default='siren_wo_n_w_morse')
    parser.add_argument('--test_loss_type', type=str, default='siren_wo_n')
    parser.add_argument('--decay_params', nargs='+', type=float, default=[3, 0.1, 3, 0.2, 0.001, 0],
                        help='epoch number to evaluate')
    parser.add_argument('--morse_type', type=str, default='l1', help='divergence term norm l1 | l2')
    parser.add_argument('--morse_decay', type=str, default='linear',
                        help='divergence term importance decay none | step | linear')
    parser.add_argument('--loss_weights', nargs='+', type=float, default=[7e3, 6e2, 1e2, 5e1, 1e2, 3, 1],
                        help='loss terms weights sdf | inter | normal | eikonal | div | morse')
    parser.add_argument('--test_loss_weights', nargs='+', type=float, default=[7e3, 6e2, 1e2, 5e1, 1e2, 0, 0],
                        help='loss terms weights sdf | inter | normal | eikonal | div | morse')
    parser.add_argument('--bidirectional_morse', action='store_true')
    parser.add_argument('--morse_near', default=True)
    return parser


def get_args():
    parser = argparse.ArgumentParser()
    parser = add_args(parser)
    args = parser.parse_args()
    return args

shapespace/shapespace_dfaust_args.py
bearprin/Neural-Singular-Hessian
import torch.utils.data as data
import numpy as np
import scipy.spatial as spatial
import open3d as o3d
import trimesh


class ReconDataset(data.Dataset):
    # A class to generate synthetic examples of basic shapes.
    # Generates clean and noisy point clouds sampled  + samples on a grid with their distance to the surface (not used in DiGS paper)
    def __init__(self, file_path, n_points, n_samples=128, res=128, sample_type='grid', sapmling_std=0.005,
                 requires_dist=False, requires_curvatures=False, grid_range=1.1):
        self.file_path = file_path
        self.n_points = n_points
        self.n_samples = n_samples
        # self.grid_res = res
        # self.sample_type = sample_type  # grid | gaussian | combined
        # self.sampling_std = sapmling_std
        # self.requires_dist = requires_dist
        # self.nonmnfld_dist, self.nonmnfld_n, self.mnfld_curvs = None, None, None
        # self.requires_curvatures = requires_curvatures  # assumes a subdirectory names "estimated props" in dataset path
        # load data
        self.o3d_point_cloud = o3d.io.read_point_cloud(self.file_path)
        self.grid_range = grid_range

        # extract center and scale points and normals
        self.points, self.mnfld_n = self.get_mnfld_points()
        self.bbox = np.array([np.min(self.points, axis=0), np.max(self.points, axis=0)]).transpose()
        self.bbox_trimesh = trimesh.PointCloud(self.points).bounding_box.copy()

        self.point_idxs = np.arange(self.points.shape[0], dtype=np.int32)
        # record sigma
        self.sample_gaussian_noise_around_shape()

    def get_mnfld_points(self):
        # Returns points on the manifold
        points = np.asarray(self.o3d_point_cloud.points, dtype=np.float32)
        normals = np.asarray(self.o3d_point_cloud.normals, dtype=np.float32)
        if normals.shape[0] == 0:
            normals = np.zeros_like(points)
        # center and scale data/point cloud
        self.cp = points.mean(axis=0)
        points = points - self.cp[None, :]
        self.scale = np.abs(points).max()
        points = points / self.scale

        return points, normals

    def sample_gaussian_noise_around_shape(self):
        kd_tree = spatial.KDTree(self.points)
        # query each point for sigma
        dist, _ = kd_tree.query(self.points, k=51, workers=-1)
        sigmas = dist[:, -1:]
        self.sigmas = sigmas
        return

    def __getitem__(self, index):
        manifold_idxes_permutation = np.random.permutation(self.points.shape[0])
        mnfld_idx = manifold_idxes_permutation[:self.n_points]
        manifold_points = self.points[mnfld_idx]  # (n_points, 3)
        manifold_normals = self.mnfld_n[mnfld_idx]  # (n_points, 3)

        nonmnfld_points = np.random.uniform(-self.grid_range, self.grid_range,
                                            size=(self.n_points, 3)).astype(np.float32)  # (n_points, 3)

        near_points = (manifold_points + self.sigmas[mnfld_idx] * np.random.randn(manifold_points.shape[0],
                                                                                  manifold_points.shape[1])).astype(
            np.float32)

        return {'points': manifold_points, 'mnfld_n': manifold_normals, 'nonmnfld_points': nonmnfld_points,
                'near_points': near_points}

    def get_train_data(self, batch_size):
        manifold_idxes_permutation = np.random.permutation(self.points.shape[0])
        mnfld_idx = manifold_idxes_permutation[:batch_size]
        manifold_points = self.points[mnfld_idx]  # (n_points, 3)
        near_points = (manifold_points + self.sigmas[mnfld_idx] * np.random.randn(manifold_points.shape[0],
                                                                                  manifold_points.shape[1])).astype(
            np.float32)

        return manifold_points, near_points, self.points

    def gen_new_data(self, dense_pts):
        self.points = dense_pts
        kd_tree = spatial.KDTree(self.points)
        # query each point for sigma^2
        dist, _ = kd_tree.query(self.points, k=51, workers=-1)
        sigmas = dist[:, -1:]
        self.sigmas = sigmas

    def __len__(self):
        return self.n_samples

surface_reconstruction/recon_dataset.py
bearprin/Neural-Singular-Hessian
import os
import threading

if __name__ == '__main__':
    data_path = '../data/sdf/input'

    logdir = './log/sdf'

    # data process
    n_points = 15000  # points number
    n_samples = 10000  # iterations number
    batch_size = 1
    grid_res = 256  # mc resolution
    # network
    layers = 4
    decoder_hidden_dim = 256
    sphere_init_params = (1.6, 0.1)
    init_type = 'siren'  # 'siren' | 'mfgi'
    # loss
    loss_type = 'siren_wo_n_w_morse'
    loss_weights = (7e3, 6e2, 1e2, 5e1, 0, 3)
    morse_type = 'l1'
    morse_decay = 'linear'  # 'linear' | 'quintic' | 'step'
    decay_params = (3, 0.2, 3, 0.4, 0.001, 0.0001)
    # opt
    lr = 5e-5
    grad_clip = 10

    files = list()
    for f in sorted(os.listdir(data_path)):
        if os.path.splitext(f)[1] in ['.xyz', '.ply']:  # only accept .xyz and .ply files
            files.append(f)
    used = list()
    device_ID = [0]  # set used GPU
    i = 0
    while True:
        if len(files) == 0:
            break
        if len(device_ID) != 0 and (
                device_ID[i] not in used or device_ID.count(device_ID[i]) > used.count(device_ID[i])):
            f = files.pop(0)


            def sub_thread():
                id = device_ID[i]
                used.append(id)
                os.system(f'CUDA_VISIBLE_DEVICES={id} python train_surface_reconstruction.py \
                          --logdir {logdir} --data_path {os.path.join(data_path, f)} --n_samples {n_samples} --n_points {n_points} --grid_res {grid_res} \
                          --lr {lr} --grad_clip_norm {grad_clip} \
                          --init_type {init_type} --decoder_hidden_dim {decoder_hidden_dim} --decoder_n_hidden_layers {layers} \
                          --loss_type {loss_type} --loss_weights {loss_weights[0]} {loss_weights[1]} {loss_weights[2]} {loss_weights[3]} {loss_weights[4]} {loss_weights[5]} \
                          --decay_params {decay_params[0]} {decay_params[1]}  {decay_params[2]} {decay_params[3]} {decay_params[4]} {decay_params[5]} --morse_type {morse_type} --morse_decay {morse_decay} --morse_near --output_any'
                          )

                used.remove(id)


            t = threading.Thread(target=sub_thread)
            t.start()
        i = (i + 1) % len(device_ID)

surface_reconstruction/run_sdf_recon.py
bearprin/Neural-Singular-Hessian
import argparse


def add_args(parser):
    parser.add_argument('--logdir', type=str, default='./log/', help='log directory')
    parser.add_argument('--model_name', type=str, default='model', help='trained model name')
    parser.add_argument('--seed', type=int, default=3627473, help='random seed')
    parser.add_argument('--data_path', type=str, help='path to input dir')
    parser.add_argument('--n_samples', type=int, default=10000,
                        help='numbers of epochs')
    parser.add_argument('--n_points', type=int, default=15000, help='number of points in each point cloud')
    parser.add_argument('--grid_res', type=int, default=128, help='uniform grid resolution')

    # training parameters
    parser.add_argument('--num_epochs', type=int, default=1, help='always be 1')
    parser.add_argument('--lr', type=float, default=5e-5, help='initial learning rate')
    parser.add_argument('--grad_clip_norm', type=float, default=10.0, help='Value to clip gradients to')
    parser.add_argument('--batch_size', type=int, default=1, help='number of samples in a minibatch')
    parser.add_argument('--load_path', type=str, default=None)

    # Network architecture and loss
    parser.add_argument('--init_type', type=str, default='mfgi',
                        help='initialization type siren | geometric_sine | geometric_relu | mfgi')
    parser.add_argument('--decoder_hidden_dim', type=int, default=256, help='length of decoder hidden dim')
    parser.add_argument('--decoder_n_hidden_layers', type=int, default=4, help='number of decoder hidden layers')
    parser.add_argument('--latent_size', type=int, default=0)
    parser.add_argument('--nl', type=str, default='sine', help='type of non linearity sine | relu')
    parser.add_argument('--sphere_init_params', nargs='+', type=float, default=[1.6, 0.1],
                        help='radius and scaling')
    parser.add_argument('--udf', action='store_true', help='enable udf')
    parser.add_argument('--output_any', action='store_true')

    parser.add_argument('--loss_type', type=str, default='siren_wo_e_wo_n_w_morse')
    parser.add_argument('--decay_params', nargs='+', type=float, default=[5, 0.2, 5, 0.4, 0.001, 0],
                        help='epoch number to evaluate')
    parser.add_argument('--morse_type', type=str, default='l1', help='divergence term norm l1 | l2')
    parser.add_argument('--morse_decay', type=str, default='linear',
                        help='divergence term importance decay none | step | linear')
    parser.add_argument('--loss_weights', nargs='+', type=float, default=[7e3, 6e2, 6e2, 5e1, 0, 10],
                        help='loss terms weights sdf | inter | normal | eikonal | div | morse')
    parser.add_argument('--bidirectional_morse', action='store_true')
    parser.add_argument('--morse_near', action='store_true')
    return parser


def get_args():
    parser = argparse.ArgumentParser()
    parser = add_args(parser)
    args = parser.parse_args()
    return args

surface_reconstruction/surface_recon_args.py
bearprin/Neural-Singular-Hessian
import os
import sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import numpy as np
import torch
import torch.optim as optim
from torchinfo import summary

from models import Network, MorseLoss
import utils.utils as utils
import utils.visualizations as vis
import surface_recon_args
import recon_dataset as dataset

# get training parameters
args = surface_recon_args.get_args()
file_name = os.path.splitext(args.data_path.split('/')[-1])[0]
logdir = os.path.join(args.logdir, file_name)
os.makedirs(logdir, exist_ok=True)

# set up logging
log_file, log_writer_train, log_writer_test, model_outdir = utils.setup_logdir(logdir, args)
os.system('cp %s %s' % (__file__, logdir))  # backup the current training file
os.system('cp %s %s' % ('recon_dataset.py', logdir))  # backup the current training file
os.system('cp %s %s' % ('../models/overfit_network.py', logdir))  # backup the models files
os.system('cp %s %s' % ('../models/losses.py', logdir))  # backup the losses files

device = 'cpu' if not torch.cuda.is_available() else 'cuda'

# get data loaders
utils.same_seed(args.seed)
train_set = dataset.ReconDataset(args.data_path, args.n_points, args.n_samples, args.grid_res)


train_dataloader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=4,
                                               pin_memory=True)
# get model
net = Network(in_dim=3, decoder_hidden_dim=args.decoder_hidden_dim, nl=args.nl,
              decoder_n_hidden_layers=args.decoder_n_hidden_layers, init_type=args.init_type,
              sphere_init_params=args.sphere_init_params, udf=args.udf)
net.to(device)
summary(net.decoder, (1, 1024, 3))

n_parameters = utils.count_parameters(net)
utils.log_string("Number of parameters in the current model:{}".format(n_parameters), log_file)

# Setup Adam optimizers
optimizer = optim.Adam(net.parameters(), lr=args.lr, weight_decay=0.0)
n_iterations = args.n_samples * (args.num_epochs)
print('n_iterations: ', n_iterations)

net.to(device)

criterion = MorseLoss(weights=args.loss_weights, loss_type=args.loss_type, div_decay=args.morse_decay,
                      div_type=args.morse_type, bidirectional_morse=args.bidirectional_morse, udf=args.udf)

num_batches = len(train_dataloader)
refine_flag = True
min_cd = np.inf
max_f1 = -np.inf
# For each epoch
for epoch in range(args.num_epochs):
    # For each batch in the dataloader
    for batch_idx, data in enumerate(train_dataloader):
        if batch_idx != 0 and (batch_idx % 500 == 0 or batch_idx == len(train_dataloader) - 1):
            output_dir = os.path.join(logdir, 'vis')
            os.makedirs(output_dir, exist_ok=True)
            vis.plot_cuts_iso(net.decoder, save_path=os.path.join(output_dir, str(batch_idx) + '.html'))
            try:
                shapename = file_name
                output_dir = os.path.join(logdir, 'result_meshes')
                os.makedirs(output_dir, exist_ok=True)
                cp, scale, bbox = train_set.cp, train_set.scale, train_set.bbox
                mesh_dict = None
                if args.udf:
                    res_dict = utils.udf2mesh(net.decoder, None,
                                              args.grid_res,
                                              translate=-cp,
                                              scale=1 / scale,
                                              get_mesh=True, device=device, bbox=bbox)
                    mesh = res_dict['mesh']
                    mesh = utils.normalize_mesh_export(mesh)
                else:
                    mesh = utils.implicit2mesh(net.decoder, None,
                                               args.grid_res,
                                               translate=-cp,
                                               scale=1 / scale,
                                               get_mesh=True, device=device, bbox=bbox)

                pred_mesh = mesh.copy()
                output_ply_filepath = os.path.join(output_dir,
                                                   shapename + '_iter_{}.ply'.format(batch_idx))

                print('Saving to ', output_ply_filepath)
                mesh.export(output_ply_filepath)
            except Exception as e:
                print(e)
                print('Could not generate mesh\n')

        net.zero_grad()
        net.train()

        mnfld_points, mnfld_n_gt, nonmnfld_points, near_points = data['points'].to(device), data['mnfld_n'].to(device), \
            data['nonmnfld_points'].to(device), data['near_points'].to(device)

        mnfld_points.requires_grad_()
        nonmnfld_points.requires_grad_()
        near_points.requires_grad_()

        output_pred = net(nonmnfld_points, mnfld_points, near_points=near_points if args.morse_near else None)
        loss_dict, _ = criterion(output_pred, mnfld_points, nonmnfld_points, mnfld_n_gt,
                                 near_points=near_points if args.morse_near else None)
        lr = torch.tensor(optimizer.param_groups[0]['lr'])
        loss_dict["lr"] = lr
        utils.log_losses(log_writer_train, epoch, batch_idx, num_batches, loss_dict, args.batch_size)

        loss_dict["loss"].backward()

        if args.grad_clip_norm > 0:
            torch.nn.utils.clip_grad_norm_(net.parameters(), args.grad_clip_norm)

        optimizer.step()

        # Output training stats
        if batch_idx % 10 == 0:
            weights = criterion.weights
            utils.log_string("Weights: {}, lr={:.3e}".format(weights, lr), log_file)
            utils.log_string('Epoch: {} [{:4d}/{} ({:.0f}%)] Loss: {:.5f} = L_Mnfld: {:.5f} + '
                             'L_NonMnfld: {:.5f} + L_Nrml: {:.5f} + L_Eknl: {:.5f} + L_Div: {:.5f} + L_Morse: {:.5f}'.format(
                epoch, batch_idx * args.batch_size, len(train_set), 100. * batch_idx / len(train_dataloader),
                loss_dict["loss"].item(), weights[0] * loss_dict["sdf_term"].item(),
                       weights[1] * loss_dict["inter_term"].item(),
                       weights[2] * loss_dict["normals_loss"].item(), weights[3] * loss_dict["eikonal_term"].item(),
                       weights[4] * loss_dict["div_loss"].item(), weights[5] * loss_dict['morse_term'].item(),
            ),
                log_file)
            utils.log_string('Epoch: {} [{:4d}/{} ({:.0f}%)] Unweighted L_s : L_Mnfld: {:.5f},  '
                             'L_NonMnfld: {:.5f},  L_Nrml: {:.5f},  L_Eknl: {:.5f}, L_Morse: {:.5f}'.format(
                epoch, batch_idx * args.batch_size, len(train_set), 100. * batch_idx / len(train_dataloader),
                loss_dict["sdf_term"].item(), loss_dict["inter_term"].item(),
                loss_dict["normals_loss"].item(), loss_dict["eikonal_term"].item(),
                loss_dict['morse_term'].item()),
                log_file)
            utils.log_string('', log_file)

        criterion.update_morse_weight(epoch * args.n_samples + batch_idx, args.num_epochs * args.n_samples,
                                      args.decay_params)  # assumes batch size of 1

surface_reconstruction/train_surface_reconstruction.py
bearprin/Neural-Singular-Hessian
import os
import random
import io
import warnings

from tensorboardX import SummaryWriter
from PIL import Image
from scipy.spatial import KDTree as KDTree
from tqdm import tqdm
from skimage import measure
import trimesh

import numpy as np
import torch
import torch.nn.functional as F
from torch.autograd import grad
import torch.backends.cudnn as cudnn

def get_aver(distances, face):
    return (distances[face[0]] + distances[face[1]] + distances[face[2]]) / 3.0


def remove_far(gt_pts, mesh, dis_trunc=0.1, is_use_prj=False):
    # gt_pts: trimesh
    # mesh: trimesh

    gt_kd_tree = KDTree(gt_pts)
    distances, vertex_ids = gt_kd_tree.query(mesh.vertices, p=2, distance_upper_bound=dis_trunc, workers=-1)
    faces_remaining = []
    faces = mesh.faces

    if is_use_prj:
        normals = gt_pts.vertex_normals
        closest_points = gt_pts.vertices[vertex_ids]
        closest_normals = normals[vertex_ids]
        direction_from_surface = mesh.vertices - closest_points
        distances = direction_from_surface * closest_normals
        distances = np.sum(distances, axis=1)

    for i in range(faces.shape[0]):
        if get_aver(distances, faces[i]) < dis_trunc:
            faces_remaining.append(faces[i])
    print(len(faces_remaining))
    mesh_cleaned = mesh.copy()
    mesh_cleaned.faces = faces_remaining
    mesh_cleaned.remove_unreferenced_vertices()

    return mesh_cleaned


def remove_outlier(gt_pts, q_pts, dis_trunc=0.003):
    # gt_pts: trimesh
    # mesh: trimesh
    gt_kd_tree = KDTree(gt_pts)
    distances, q_ids = gt_kd_tree.query(q_pts, p=2, distance_upper_bound=dis_trunc, workers=-1)
    q_pts = q_pts[distances < dis_trunc]
    return q_pts


def same_seed(seed):
    """

    :param seed:
    :return:
    """
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    cudnn.benchmark = False
    cudnn.deterministic = True


def normalize_mesh_export(mesh, file_out=None):
    # unit to [-0.5, 0.5]
    bounds = mesh.extents
    if bounds.min() == 0.0:
        return

    # translate to origin
    translation = (mesh.bounds[0] + mesh.bounds[1]) * 0.5
    translation = trimesh.transformations.translation_matrix(direction=-translation)
    mesh.apply_transform(translation)

    # scale to unit cube
    scale = 1.0 / bounds.max()
    scale_trafo = trimesh.transformations.scale_matrix(factor=scale)
    mesh.apply_transform(scale_trafo)
    if file_out is not None:
        mesh.export(file_out)
    return mesh


def eval_reconstruct_gt(rec_mesh: trimesh.Trimesh, gt_pts, sample_num=100000):
    def distance_p2p(points_src, normals_src, points_tgt, normals_tgt):
        ''' Computes minimal distances of each point in points_src to points_tgt.

        Args:
            points_src (numpy array): source points
            normals_src (numpy array): source normals
            points_tgt (numpy array): target points
            normals_tgt (numpy array): target normals
        '''
        kdtree = KDTree(points_tgt)
        dist, idx = kdtree.query(points_src, workers=-1)

        if normals_src is not None and normals_tgt is not None:
            normals_src = \
                normals_src / np.linalg.norm(normals_src, axis=-1, keepdims=True)
            normals_tgt = \
                normals_tgt / np.linalg.norm(normals_tgt, axis=-1, keepdims=True)

            #        normals_dot_product = (normals_tgt[idx] * normals_src).sum(axis=-1)
            #        # Handle normals that point into wrong direction gracefully
            #        # (mostly due to mehtod not caring about this in generation)
            #        normals_dot_product = np.abs(normals_dot_product)

            normals_dot_product = np.abs(normals_tgt[idx] * normals_src)
            normals_dot_product = normals_dot_product.sum(axis=-1)
        else:
            normals_dot_product = np.array(
                [np.nan] * points_src.shape[0], dtype=np.float32)
        return dist, normals_dot_product

    def get_threshold_percentage(dist, thresholds):
        ''' Evaluates a point cloud.
        Args:
            dist (numpy array): calculated distance
            thresholds (numpy array): threshold values for the F-score calculation
        '''
        in_threshold = [
            (dist <= t).mean() for t in thresholds
        ]
        return in_threshold

    def eval_pointcloud(pointcloud, pointcloud_tgt,
                        normals=None, normals_tgt=None, thresholds=[0.005]):
        ''' Evaluates a point cloud.

        Args:
            pointcloud (numpy array): predicted point cloud
            pointcloud_tgt (numpy array): target point cloud
            normals (numpy array): predicted normals
            normals_tgt (numpy array): target normals
        '''
        # Return maximum losses if pointcloud is empty

        pointcloud = np.asarray(pointcloud)
        pointcloud_tgt = np.asarray(pointcloud_tgt)

        # Completeness: how far are the points of the target point cloud
        # from thre predicted point cloud
        completeness, completeness_normals = distance_p2p(
            pointcloud_tgt, normals_tgt, pointcloud, normals
        )
        recall = get_threshold_percentage(completeness, thresholds)
        completeness2 = completeness ** 2

        completeness = completeness.mean()
        completeness2 = completeness2.mean()
        completeness_normals = completeness_normals.mean()

        # Accuracy: how far are th points of the predicted pointcloud
        # from the target pointcloud
        accuracy, accuracy_normals = distance_p2p(
            pointcloud, normals, pointcloud_tgt, normals_tgt
        )
        precision = get_threshold_percentage(accuracy, thresholds)
        accuracy2 = accuracy ** 2

        accuracy = accuracy.mean()
        accuracy2 = accuracy2.mean()
        accuracy_normals = accuracy_normals.mean()
        # print(completeness,accuracy,completeness2,accuracy2)
        # Chamfer distance
        chamferL2 = 0.5 * (completeness2 + accuracy2)
        normals_correctness = (
                0.5 * completeness_normals + 0.5 * accuracy_normals
        )
        chamferL1 = 0.5 * (completeness + accuracy)

        # F-Score
        F = [
            2 * precision[i] * recall[i] / (precision[i] + recall[i])
            for i in range(len(precision))
        ]
        return normals_correctness, chamferL1, chamferL2, F[0]

    gt_pts = normalize_mesh_export(gt_pts)
    rec_mesh = normalize_mesh_export(rec_mesh)

    # sample point for rec
    pts_rec, idx = rec_mesh.sample(sample_num, return_index=True)
    normals_rec = rec_mesh.face_normals[idx]
    # sample point for gt
    pts_gt = None
    normals_gt = None
    if isinstance(gt_pts, trimesh.PointCloud):
        if gt_pts.shape[0] < sample_num:
            sample_num = gt_pts.shape[0]
        idx = np.random.choice(gt_pts.vertices.shape[0], sample_num, replace=False)
        pts_gt = gt_pts.vertices[idx]
        normals_gt = None
    elif isinstance(gt_pts, trimesh.Trimesh):
        pts_gt, idx = gt_pts.sample(sample_num, return_index=True)
        normals_gt = gt_pts.face_normals[idx]

    normals_correctness, chamferL1, chamferL2, f1_mu = eval_pointcloud(pts_rec, pts_gt, normals_rec, normals_gt)

    return normals_correctness, chamferL1, chamferL2, f1_mu


def center_and_scale(points, cp=None, scale=None):
    # center a point cloud and scale it to unite sphere.
    if cp is None:
        cp = points.mean(axis=1)
    points = points - cp[:, None, :]
    if scale is None:
        scale = np.linalg.norm(points, axis=-1).max(-1)
    points = points / scale[:, None, None]
    return points, cp, scale


def log_losses(writer, epoch, bach_idx, num_batches, loss_dict, batch_size):
    # helper function to log losses to tensorboardx writer
    fraction_done = (bach_idx + 1) / num_batches
    iteration = (epoch + fraction_done) * num_batches * batch_size
    for loss in loss_dict.keys():
        writer.add_scalar(loss, loss_dict[loss].item(), iteration)
    return iteration


def log_weight_hist(writer, epoch, bach_idx, num_batches, net_blocks, batch_size):
    # helper function to log losses to tensorboardx writer
    fraction_done = (bach_idx + 1) / num_batches
    iteration = (epoch + fraction_done) * num_batches * batch_size
    for i, block in enumerate(net_blocks):
        writer.add_histogram('layer_weights_' + str(i), block[0].weight, iteration)
        writer.add_histogram('layer_biases_' + str(i), block[0].bias
                             , iteration)
    return iteration


def log_images(writer, iteration, contour_img, curl_img, eikonal_img, div_image, z_diff_img, example_idx):
    # helper function to log images to tensorboardx writer
    writer.add_image('implicit_function/' + str(example_idx), contour_img.transpose(2, 0, 1), iteration)
    writer.add_image('curl/' + str(example_idx), curl_img.transpose(2, 0, 1), iteration)
    writer.add_image('eikonal_term/' + str(example_idx), eikonal_img.transpose(2, 0, 1), iteration)
    writer.add_image('divergence/' + str(example_idx), div_image.transpose(2, 0, 1), iteration)
    writer.add_image('z_diff/' + str(example_idx), z_diff_img.transpose(2, 0, 1), iteration)


def log_string(out_str, log_file):
    # helper function to log a string to file and print it
    log_file.write(out_str + '\n')
    log_file.flush()
    print(out_str)


def setup_logdir(logdir, args=None):
    # helper function to set up logging directory

    os.makedirs(logdir, exist_ok=True)
    log_writer_train = SummaryWriter(os.path.join(logdir, 'train'))
    log_writer_test = SummaryWriter(os.path.join(logdir, 'test'))
    log_filename = os.path.join(logdir, 'out.log')
    log_file = open(log_filename, 'w')
    model_outdir = os.path.join(logdir, 'trained_models')
    os.makedirs(model_outdir, exist_ok=True)

    if args is not None:
        params_filename = os.path.join(model_outdir, '%s_params.pth' % (args.model_name))
        torch.save(args, params_filename)  # save parameters
        log_string("input params: \n" + str(args), log_file)
    else:
        warnings.warn("Training options not provided. Not saving training options...")

    return log_file, log_writer_train, log_writer_test, model_outdir


def backup_code(logdir, dir_list=[], file_list=[]):
    # backup models code
    code_bkp_dir = os.path.join(logdir, 'code_bkp')
    os.makedirs(code_bkp_dir, exist_ok=True)
    for dir_name in dir_list:
        print("copying directory {} to {}".format(dir_name, code_bkp_dir))
        os.system('cp -r %s %s' % (dir_name, code_bkp_dir))  # backup the current model code
    for file_name in file_list:
        print("copying file {} to {}".format(file_name, code_bkp_dir))
        os.system('cp %s %s' % (file_name, code_bkp_dir))


def plotly_fig2array(fig):
    # convert Plotly fig to  an array
    fig_bytes = fig.to_image(format="png")
    buf = io.BytesIO(fig_bytes)
    img = Image.open(buf).convert('RGB')
    img = np.asarray(img)
    return img


def gradient(inputs, outputs, create_graph=True, retain_graph=True):
    d_points = torch.ones_like(outputs, requires_grad=False, device=outputs.device)
    points_grad = grad(
        outputs=outputs,
        inputs=inputs,
        grad_outputs=d_points,
        create_graph=create_graph,
        retain_graph=retain_graph,
        only_inputs=True)[0]  # [:, -3:]
    return points_grad


# def gradient(inputs, outputs):
#     outputs.sum().backward(retain_graph=True)
#     return inputs.grad
# sdf = net(q)
# sdf.sum().
# q_grad = q.grad.detach()  # (1, 500, 3)
# q_grad = F.normalize(q_grad, dim=2)


def get_cuda_ifavailable(torch_obj, device=None):
    # if cuda is available return a cuda obeject
    if (torch.cuda.is_available()):
        return torch_obj.cuda(device=device)
    else:
        return torch_obj


def compute_props(decoder, latent, z_gt, device):
    # compute derivative properties on a grid
    res = z_gt.shape[1]
    x, y, grid_points = get_2d_grid_uniform(resolution=res, range=1.2, device=device)
    grid_points.requires_grad_()
    if latent is not None:
        grid_points_latent = torch.cat([latent.expand(grid_points.shape[0], -1), grid_points], dim=1)
    else:
        grid_points_latent = grid_points
    z = decoder(grid_points_latent)
    z_np = z.detach().cpu().numpy().reshape(x.shape[0], x.shape[0])

    # plot z difference image
    z_diff = np.abs(np.abs(np.reshape(z_np, [res, res])) - np.abs(z_gt)).reshape(x.shape[0], x.shape[0])

    return x, y, z_np, z_diff


def compute_deriv_props(decoder, latent, z_gt, device):
    # compute derivative properties on a grid
    res = z_gt.shape[1]
    x, y, grid_points = get_2d_grid_uniform(resolution=res, range=1.2, device=device)
    grid_points.requires_grad_()
    if latent is not None:
        grid_points_latent = torch.cat([latent.expand(grid_points.shape[0], -1), grid_points], dim=1)
    else:
        grid_points_latent = grid_points
    z = decoder(grid_points_latent)
    z_np = z.detach().cpu().numpy().reshape(x.shape[0], x.shape[0])

    # compute derivatives
    grid_grad = gradient(grid_points, z)
    dx = gradient(grid_points, grid_grad[:, 0], create_graph=False, retain_graph=True)
    dy = gradient(grid_points, grid_grad[:, 1], create_graph=False, retain_graph=False)

    grid_curl = (dx[:, 1] - dy[:, 0]).cpu().detach().numpy().reshape(x.shape[0], x.shape[0])

    # compute eikonal term (gradient magnitude)
    eikonal_term = ((grid_grad.norm(2, dim=-1) - 1) ** 2).cpu().detach().numpy().reshape(x.shape[0], x.shape[0])

    # compute divergence
    grid_div = (dx[:, 0] + dy[:, 1]).detach().cpu().numpy().reshape(x.shape[0], x.shape[0])

    # compute det of hessian
    hessian = torch.stack([dx, dy], dim=1)
    eigs = torch.linalg.eigvalsh(hessian).detach().cpu().numpy().reshape(x.shape[0], x.shape[0], 2)
    hessian_det = torch.det(hessian).cpu().detach().numpy().reshape(x.shape[0], x.shape[0])

    # plot z difference image
    z_diff = np.abs(np.abs(np.reshape(z_np, [res, res])) - np.abs(z_gt)).reshape(x.shape[0], x.shape[0])

    return x, y, z_np, z_diff, eikonal_term, grid_div, grid_curl, hessian_det, eigs


def get_2d_grid_uniform(resolution=100, range=1.2, device=None):
    # generate points on a uniform grid within  a given range
    x = np.linspace(-range, range, resolution)
    y = x
    xx, yy = np.meshgrid(x, y)
    grid_points = get_cuda_ifavailable(torch.tensor(np.vstack([xx.ravel(), yy.ravel()]).T, dtype=torch.float),
                                       device=device)
    return x, y, grid_points


def get_3d_grid(resolution=100, bbox=1.2 * np.array([[-1, 1], [-1, 1], [-1, 1]]), device=None, eps=0.1,
                dtype=np.float16):
    # generate points on a uniform grid within  a given range
    # reimplemented from SAL : https://github.com/matanatz/SAL/blob/master/code/utils/plots.py
    # and IGR : https://github.com/amosgropp/IGR/blob/master/code/utils/plots.py

    shortest_axis = np.argmin(bbox[:, 1] - bbox[:, 0])
    if (shortest_axis == 0):
        x = np.linspace(bbox[0, 0] - eps, bbox[0, 1] + eps, resolution)
        length = np.max(x) - np.min(x)
        y = np.arange(bbox[1, 0] - eps, bbox[1, 1] + length / (x.shape[0] - 1) + eps, length / (x.shape[0] - 1))
        z = np.arange(bbox[2, 0] - eps, bbox[2, 1] + length / (x.shape[0] - 1) + eps, length / (x.shape[0] - 1))
    elif (shortest_axis == 1):
        y = np.linspace(bbox[1, 0] - eps, bbox[1, 1] + eps, resolution)
        length = np.max(y) - np.min(y)
        x = np.arange(bbox[0, 0] - eps, bbox[0, 1] + length / (y.shape[0] - 1) + eps, length / (y.shape[0] - 1))
        z = np.arange(bbox[2, 0] - eps, bbox[2, 1] + length / (y.shape[0] - 1) + eps, length / (y.shape[0] - 1))
    elif (shortest_axis == 2):
        z = np.linspace(bbox[2, 0] - eps, bbox[2, 1] + eps, resolution)
        length = np.max(z) - np.min(z)
        x = np.arange(bbox[0, 0] - eps, bbox[0, 1] + length / (z.shape[0] - 1) + eps, length / (z.shape[0] - 1))
        y = np.arange(bbox[1, 0] - eps, bbox[1, 1] + length / (z.shape[0] - 1) + eps, length / (z.shape[0] - 1))

    xx, yy, zz = np.meshgrid(x.astype(dtype), y.astype(dtype), z.astype(dtype))  #
    # grid_points = get_cuda_ifavailable(torch.tensor(np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T, dtype=torch.float16),
    #                                          device=device)
    grid_points = torch.tensor(np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T, dtype=torch.float32)
    return {"grid_points": grid_points,
            "shortest_axis_length": length,
            "xyz": [x, y, z],
            "shortest_axis_index": shortest_axis}


def scale_pc_to_unit_sphere(points, cp=None, s=None):
    if cp is None:
        cp = points.mean(axis=0)
    points = points - cp[None, :]
    if s is None:
        s = np.linalg.norm(points, axis=-1).max(-1)
    points = points / s
    return points, cp, s


def recon_metrics(pc1, pc2, one_sided=False, alphas=[0.0001, 0.0002, 0.0003, 0.0004, 0.0005],
                  percentiles=[5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 95], k=[10, 25, 50], return_all=False):
    # Compute reconstruction benchmarc evaluation metrics :
    # chamfer and hausdorff distance metrics between two point clouds pc1 and pc2 [nx3]
    # percentage of distance points metric (not used in the paper)
    # and the meal local chamfer variance
    # pc1 is the reconstruction and pc2 is the gt data

    scale = np.abs(pc2).max()

    # compute one side
    pc1_kd_tree = KDTree(pc1)
    one_distances, one_vertex_ids = pc1_kd_tree.query(pc2, n_jobs=4)
    cd12 = np.mean(one_distances)
    hd12 = np.max(one_distances)
    cdmed12 = np.median(one_distances)
    cd21 = None
    hd21 = None
    cdmed21 = None
    pods2 = None
    cdp2 = None
    chamfer_distance = cd12
    hausdorff_distance = hd12

    # compute chamfer distance percentiles cdp
    cdp1 = np.percentile(one_distances, percentiles, interpolation='lower')

    # compute PoD
    pod1 = []
    for alpha in alphas:
        pod1.append((one_distances < alpha * scale).sum() / one_distances.shape[0])

    if not one_sided:
        # compute second side
        pc2_kd_tree = KDTree(pc2)
        two_distances, two_vertex_ids = pc2_kd_tree.query(pc1, n_jobs=4)
        cd21 = np.mean(two_distances)
        hd21 = np.max(two_distances)
        cdmed21 = np.median(two_distances)
        chamfer_distance = 0.5 * (cd12 + cd21)
        hausdorff_distance = np.max((hd12, hd21))
        # compute chamfer distance percentiles cdp
        cdp2 = np.percentile(two_distances, percentiles)
        # compute PoD
        pod2 = []
        for alpha in alphas:
            pod2.append((two_distances < alpha * scale).sum() / two_distances.shape[0])

    # compute double sided pod
    pod12 = []
    for alpha in alphas:
        pod12.append(((one_distances < alpha * scale).sum() + (two_distances < alpha * scale).sum()) /
                     (one_distances.shape[0] + two_distances.shape[0]))
    cdp12 = np.percentile(np.concatenate([one_distances, two_distances]),
                          percentiles)  # compute chamfer distance percentiles cdp

    nn1_dist, local_idx2 = pc1_kd_tree.query(pc1, max(k), n_jobs=-1)
    nn1_dist_2pc2 = two_distances[local_idx2]
    malcv = [(nn1_dist_2pc2[:, :k0] / nn1_dist.mean(axis=1, keepdims=True)).var(axis=1).mean() for k0 in k]

    if return_all:
        return chamfer_distance, hausdorff_distance, (cd12, cd21, cdmed12, cdmed21, hd12, hd21), (pod1, pod2, pod12), \
            (cdp1.tolist(), cdp2.tolist(), cdp12.tolist()), malcv, (one_distances, two_distances)

    return chamfer_distance, hausdorff_distance, (cd12, cd21, cdmed12, cdmed21, hd12, hd21), (pod1, pod2, pod12), \
        (cdp1.tolist(), cdp2.tolist(), cdp12.tolist()), malcv


def load_reconstruction_data(file_path, n_points=30000, sample_type='vertices'):
    extension = file_path.split('.')[-1]
    if extension == 'xyz':
        points = np.loadtxt(file_path)
    elif extension == 'ply':
        mesh = trimesh.load_mesh(file_path)

        if hasattr(mesh, 'faces') and not sample_type == 'vertices':
            # sample points if its a triangle mesh
            points = trimesh.sample.sample_surface(mesh, n_points)[0]
        else:
            # use the vertices if its a point cloud
            points = mesh.vertices
    # Center and scale points
    # cp = points.mean(axis=0)
    # points = points - cp[None, :]
    # scale = np.abs(points).max()
    # points = points / scale
    return np.array(points).astype('float32')


def implicit2mesh(decoder, mods, grid_res, translate=[0., 0., 0.], scale=1, get_mesh=True, device=None,
                  bbox=np.array([[-1, 1], [-1, 1], [-1, 1]]), feat=None, hash_tree=None):
    # compute a mesh from the implicit representation in the decoder.
    # Uses marching cubes.
    # reimplemented from SAL get surface trace function : https://github.com/matanatz/SAL/blob/master/code/utils/plots.py
    print('in implicit2mesh')
    print(grid_res, translate, scale, bbox)
    mesh = None
    grid_dict = get_3d_grid(resolution=grid_res, bbox=bbox, device=device)
    print('Finished getting grid_dict')
    cell_width = grid_dict['xyz'][0][2] - grid_dict['xyz'][0][1]
    pnts = grid_dict["grid_points"]

    z = []
    for point in tqdm(torch.split(pnts, 100000, dim=0)):
        # point: (100000, 3)
        point = get_cuda_ifavailable(point, device=device)
        if feat is not None:
            if point.dim() == 2:
                point = point.unsqueeze(0)
            query_feat = decoder.encoder.query_feature(feat, point)
            z.append(decoder.decoder(point, query_feat).detach().squeeze(0).cpu().numpy())
        elif hash_tree is not None:
            if point.dim() == 2:
                point = point.unsqueeze(0)
            # query_feat = decoder.encoder.query_feature(feat, point)
            query_feat = decoder.encoder.query_feature(hash_tree, feat, point)
            z.append(decoder.decoder(point, query_feat).detach().squeeze(0).cpu().numpy())
        else:
            z.append(decoder(point.type(torch.float32), mods).detach().squeeze(0).cpu().numpy())
    z = np.concatenate(z, axis=0).reshape(grid_dict['xyz'][1].shape[0], grid_dict['xyz'][0].shape[0],
                                          grid_dict['xyz'][2].shape[0]).transpose([1, 0, 2]).astype(np.float64)

    print(z.min(), z.max())

    threshs = [0.00]
    mesh_list = list()
    for thresh in threshs:
        if (np.sum(z > 0.0) < np.sum(z < 0.0)):
            thresh = -thresh
        verts, faces, normals, values = measure.marching_cubes(volume=z, level=thresh,
                                                               spacing=(cell_width, cell_width, cell_width))

        verts = verts + np.array([grid_dict['xyz'][0][0], grid_dict['xyz'][1][0], grid_dict['xyz'][2][0]])
        verts = verts * (1 / scale) - translate

        if get_mesh:
            mesh = trimesh.Trimesh(verts, faces, vertex_normals=normals, vertex_colors=values, validate=True)
        mesh_list.append(mesh)
    return mesh_list[0]


# @nb.jit()
def surface_extraction_single(ndf, grad, b_max, b_min, resolution):
    '''
    From CAP-UDF (https://github.com/junshengzhou/CAP-UDF)
    '''
    v_all = []
    t_all = []
    threshold = 0.005  # accelerate extraction
    v_num = 0
    for i in range(resolution - 1):
        for j in range(resolution - 1):
            for k in range(resolution - 1):
                ndf_loc = ndf[i:i + 2]
                ndf_loc = ndf_loc[:, j:j + 2, :]
                ndf_loc = ndf_loc[:, :, k:k + 2]
                if np.min(ndf_loc) > threshold:
                    continue
                grad_loc = grad[i:i + 2]
                grad_loc = grad_loc[:, j:j + 2, :]
                grad_loc = grad_loc[:, :, k:k + 2]

                res = np.ones((2, 2, 2))
                for ii in range(2):
                    for jj in range(2):
                        for kk in range(2):
                            if np.dot(grad_loc[0][0][0], grad_loc[ii][jj][kk]) < 0:
                                res[ii][jj][kk] = -ndf_loc[ii][jj][kk]
                            else:
                                res[ii][jj][kk] = ndf_loc[ii][jj][kk]

                if res.min() < 0:
                    vertices, triangles, _, _ = measure.marching_cubes(
                        res, 0.0)
                    # print(vertices)
                    # vertices -= 1.5
                    # vertices /= 128
                    vertices[:, 0] += i  # / resolution
                    vertices[:, 1] += j  # / resolution
                    vertices[:, 2] += k  # / resolution
                    triangles += v_num
                    # vertices =
                    # vertices[:,1] /= 3  # TODO
                    v_all.append(vertices)
                    t_all.append(triangles)

                    v_num += vertices.shape[0]
                    # print(v_num)

    v_all = np.concatenate(v_all)
    t_all = np.concatenate(t_all)

    mesh = trimesh.Trimesh(v_all, t_all)
    mesh.remove_duplicate_faces()
    mesh.remove_degenerate_faces()
    mesh.fill_holes()

    return mesh.vertices, mesh.faces


def inner_iteration(ndf_loc, grad_loc, i, j, k):
    ndf_loc = np.ascontiguousarray(ndf_loc)
    grad_loc = np.ascontiguousarray(grad_loc)
    res = np.ones((2, 2, 2))
    for ii in range(2):
        for jj in range(2):
            for kk in range(2):
                if np.dot(grad_loc[0][0][0], grad_loc[ii][jj][kk]) < 0:
                    res[ii][jj][kk] = -ndf_loc[ii][jj][kk]
                else:
                    res[ii][jj][kk] = ndf_loc[ii][jj][kk]
    shape_dict = dict()
    if res.min() < 0:
        v, f, _, _ = measure.marching_cubes(res, 0.0)
        v[:, 0] = v[:, 0] + i
        v[:, 1] = v[:, 1] + j
        v[:, 2] = v[:, 2] + k
        shape_dict['v'] = v
        shape_dict['f'] = f
    else:
        shape_dict['v'] = None
        shape_dict['f'] = None

    return shape_dict


import numba

@numba.jit(nopython=True, fastmath=True)
def map(ndf, grad, resolution=128, threshold=0.01):
    call_params = list()
    for i in range(resolution - 1):
        for j in range(resolution - 1):
            for k in range(resolution - 1):
                ndf_loc = ndf[i:i + 2]
                ndf_loc = ndf_loc[:, j:j + 2, :]
                ndf_loc = ndf_loc[:, :, k:k + 2]
                if np.min(ndf_loc) > threshold:
                    continue
                grad_loc = grad[i:i + 2]
                grad_loc = grad_loc[:, j:j + 2, :]
                grad_loc = grad_loc[:, :, k:k + 2]
                call_params.append((ndf_loc, grad_loc, i, j, k))
    return call_params


def reduce(res_dict_list):
    v_all = list()
    f_all = list()
    v_num = 0
    for res in res_dict_list:
        if res['v'] is not None:
            v_all.append(res['v'])
            f_all.append(res['f'] + v_num)
            v_num += res['v'].shape[0]
    v_all = np.concatenate(v_all)
    f_all = np.concatenate(f_all)
    mesh = trimesh.Trimesh(v_all, f_all)
    return mesh


num_processes = 8

from .utils_mp import start_process_pool


def surface_extraction_mp(ndf, grad, resolution):
    '''
    Modified From CAP-UDF (https://github.com/junshengzhou/CAP-UDF)
    '''
    call_params = map(ndf, grad, resolution)
    res_dict_list = start_process_pool(inner_iteration, call_params, num_processes)
    mesh = reduce(res_dict_list)

    mesh.remove_duplicate_faces()
    mesh.remove_degenerate_faces()
    mesh.fill_holes()

    return mesh.vertices, mesh.faces


def udf2mesh(decoder, latent, grid_res, translate=[0., 0., 0.], scale=1, get_mesh=True, device=None,
             bbox=np.array([[-1, 1], [-1, 1], [-1, 1]])):
    # using mc from MeshUDF
    print('in implicit2mesh')
    print(grid_res, translate, scale, bbox)
    mesh = None

    def get_query_point(device, bd=1.05, resolution=128):
        shape = (resolution, resolution, resolution)
        vxs = torch.arange(-bd, bd, bd * 2 / resolution)
        vys = torch.arange(-bd, bd, bd * 2 / resolution)
        vzs = torch.arange(-bd, bd, bd * 2 / resolution)
        pxs = vxs.view(-1, 1, 1).expand(*shape).contiguous().view(resolution ** 3)
        pys = vys.view(1, -1, 1).expand(*shape).contiguous().view(resolution ** 3)
        pzs = vzs.view(1, 1, -1).expand(*shape).contiguous().view(resolution ** 3)
        p = torch.stack([pxs, pys, pzs], dim=1).reshape(resolution ** 3, 3).to(device)
        return p

    pnts = get_query_point(resolution=grid_res, device=device)
    z = []
    n = []
    for point in tqdm(torch.split(pnts, 100000, dim=0)):
        # point: (100000, 3)
        if latent is not None:
            point = torch.cat([point, latent.unsqueeze(0).repeat(point.shape[0], 1), ], dim=1)
        point = get_cuda_ifavailable(point, device=device)
        point.requires_grad_()
        pred = decoder(point.type(torch.float32))
        normals = gradient(point, pred).detach()
        z.append(pred.detach().cpu().float().numpy())
        n.append(normals.detach().cpu().float().numpy())
    z = np.concatenate(z, axis=0).reshape(grid_res, grid_res, grid_res).astype(np.float32)
    n = np.concatenate(n, axis=0).reshape(grid_res, grid_res, grid_res, 3).astype(np.float32)
    print(z.min(), z.max())

    pnts = pnts.detach().cpu().float().numpy().reshape(grid_res, grid_res, grid_res, 3)
    # verts, faces = surface_extraction_single(z, n, resolution=grid_res, b_min=bbox.min(1), b_max=bbox.max(1))
    verts, faces = surface_extraction_mp(z, n, resolution=grid_res)
    # verts, faces = EMC(grids_coords=pnts, grids_udf=z, grids_udf_grad=n,
    #                    voxel_size=2. / grid_res, res=grid_res)

    # # verts = verts + np.array([grid_dict['xyz'][0][0], grid_dict['xyz'][1][0], grid_dict['xyz'][2][0]])
    # # verts = verts * (1 / scale) - translate
    #
    if get_mesh:
        mesh = trimesh.Trimesh(verts, faces, validate=True)
    res_dict = dict()
    res_dict['mesh'] = mesh
    res_dict['z'] = z
    res_dict['n'] = n
    return res_dict


def convert_xyz_to_ply_with_noise(file_path, noise=None):
    # convert ply file in xyznxnynz format to ply file
    points = np.loadtxt(file_path)
    if noise is None:
        mesh = trimesh.Trimesh(points[:, :3], [], vertex_normals=points[:, 3:])
        mesh.export(file_path.split('.')[0] + '.ply', vertex_normal=True)
    else:
        for std in noise:
            bbox_scale = np.abs(points).max()
            var = std * std
            cov_mat = bbox_scale * np.array([[var, 0., 0.], [0., var, 0.], [0., 0., var]])
            noise = np.random.multivariate_normal([0., 0., 0.], cov_mat, size=points.shape[0], check_valid='warn',
                                                  tol=1e-8)
            mesh = trimesh.Trimesh(points[:, :3] + noise, [], vertex_normals=points[:, 3:])
            mesh.export(file_path.split('.')[0] + '_' + str(std) + '.ply', vertex_normal=True)


def count_parameters(model):
    # count the number of parameters in a given model
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


if __name__ == "__main__":
    file_path = '/home/sitzikbs/Datasets/Reconstruction_IKEA_sample/interior_room.xyz'
    convert_xyz_to_ply_with_noise(file_path, noise=[0.01])

utils/utils.py
bearprin/Neural-Singular-Hessian
import plotly.graph_objects as go
import plotly.figure_factory as ff
from plotly.subplots import make_subplots
import numpy as np
import plotly.offline as offline
import utils.utils as utils
import torch
from PIL import Image
import os


def plot_cuts_iso(decode_points_func,
                  box_size=(2.5, 2.5, 2.5), max_n_eval_pts=1e6,
                  resolution=256, thres=0.0, imgs_per_cut=1, save_path=None) -> go.Figure:
    """ plot levelset at a certain cross section, assume inputs are centered
    Args:
        decode_points_func: A function to extract the SDF/occupancy logits of (N, 3) points
        box_size (List[float]): bounding box dimension
        max_n_eval_pts (int): max number of points to evaluate in one inference
        resolution (int): cross section resolution xy
        thres (float): levelset value
        imgs_per_cut (int): number of images for each cut (plotted in rows)
    Returns:
        a numpy array for the image
    """
    xmax, ymax, zmax = [b / 2 for b in box_size]
    xx, yy = np.meshgrid(np.linspace(-xmax, xmax, resolution),
                         np.linspace(-ymax, ymax, resolution))
    xx = xx.ravel()
    yy = yy.ravel()

    fig = make_subplots(rows=imgs_per_cut, cols=3,
                        subplot_titles=('xz', 'xy', 'yz'),
                        shared_xaxes='all', shared_yaxes='all',
                        vertical_spacing=0.01, horizontal_spacing=0.01,
                        )

    def _plot_cut(fig, idx, pos, decode_points_func, xmax, ymax, resolution):
        """ plot one cross section pos (3, N) """
        # evaluate points in serial
        field_input = torch.tensor(pos.T, dtype=torch.float).cuda()
        feat = torch.zeros((field_input.shape[0], 1)).cuda()
        feat[:, 0] = 1
        values = decode_points_func(field_input).flatten()
        if isinstance(values, torch.Tensor):
            values = values.detach().cpu().numpy()
            # np.save('xy.npy', values)
        values = values.reshape(resolution, resolution)
        contour_dict = dict(autocontour=False,
                            colorscale='RdBu',
                            reversescale=True,
                            contours=dict(
                                start=-0.01 + thres,
                                end=0.01 + thres,
                                size=0.001,
                                showlabels=True,  # show labels on contours
                                # labelfont=dict(  # label font properties
                                #     size=12,
                                #     color='white',
                                # )
                            ), )
        r_idx = idx // 3

        fig.add_trace(
            go.Contour(x=np.linspace(-xmax, xmax, resolution),
                       y=np.linspace(-ymax, ymax, resolution),
                       z=values,
                       **contour_dict
                       ),
            col=idx % 3 + 1, row=r_idx + 1  # 1-idx
        )

        fig.update_xaxes(
            range=[-xmax, xmax],  # sets the range of xaxis
            constrain="range",  # meanwhile compresses the xaxis by decreasing its "domain"
            col=idx % 3 + 1, row=r_idx + 1)
        fig.update_yaxes(
            range=[-ymax, ymax],
            col=idx % 3 + 1, row=r_idx + 1
        )

    steps = np.stack([np.linspace(-b / 2, b / 2, imgs_per_cut + 2)[1:-1] for b in box_size], axis=-1)
    for index in range(imgs_per_cut):
        position_cut = [np.vstack([xx, np.full(xx.shape[0], steps[index, 1]), yy]),
                        np.vstack([xx, yy, np.full(xx.shape[0], steps[index, 2])]),
                        np.vstack([np.full(xx.shape[0], steps[index, 0]), xx, yy]), ]
        _plot_cut(
            fig, index * 3, position_cut[0], decode_points_func, xmax, zmax, resolution)
        _plot_cut(
            fig, index * 3 + 1, position_cut[1], decode_points_func, xmax, ymax, resolution)
        _plot_cut(
            fig, index * 3 + 2, position_cut[2], decode_points_func, ymax, zmax, resolution)

    fig.update_layout(
        title='iso-surface',
        height=1200 * imgs_per_cut,
        width=1200 * 3,
        autosize=False,
        scene=dict(aspectratio=dict(x=1, y=1))
    )

    if save_path is not None:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        fig.write_html(save_path)

    return fig


def plot_contours(x, y, z, colorscale='Geyser', fig=None):
    # plots a contour image given hight map matrix (2d array)
    layout = go.Layout(width=1200, height=1200,
                       xaxis=dict(side="bottom", range=[-1, 1], ),
                       yaxis=dict(side="left", range=[-1, 1], ),
                       scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                  yaxis=dict(range=[-1, 1], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text='distance map', y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Serif', size=24, color='red'))
                       )
    if fig is None:
        traces = []
    else:
        traces = list(fig.data)
    traces.append(go.Contour(x=x.flatten(), y=y.flatten(), z=z.flatten(),
                             colorscale=colorscale,
                             contours=dict(
                                 start=-1,
                                 end=1,
                                 size=0.025,
                             ),
                             ))  # contour trace

    fig = go.Figure(data=traces, layout=layout)
    fig.show()


def plot_scatter(x, y):
    layout = go.Layout(width=1200, height=1200, scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                                           yaxis=dict(range=[-1, 1], autorange=False),
                                                           aspectratio=dict(x=1, y=1)),
                       title=dict(text='scatter plot'))
    fig = go.Figure(data=go.Scatter(
        x=x,
        y=y,
        mode='markers',
        marker=dict(
            size=16
        )
    ), layout=layout)

    fig.show()


def plot_contour_and_scatter_sal_training(clean_points, decoder, latent=None,
                                          example_idx=0, n_gt=None, n_pred=None,
                                          show_ax=True, show_bar=True, title_text='', colorscale='RdBu',
                                          res=256):
    # plot contour and scatter plot given input points, SAL decoder and the point's latent representation

    clean_points = clean_points[example_idx].detach().cpu().numpy()

    if n_gt is not None:
        n_gt = n_gt[example_idx].detach().cpu().numpy()
    if n_pred is not None:
        n_pred = n_pred[example_idx].detach().cpu().numpy()
    x, y, grid_points = utils.get_2d_grid_uniform(resolution=res, range=1.2, device=latent.device)
    if latent is not None:
        grid_points = torch.cat([latent[example_idx].expand(grid_points.shape[0], -1), grid_points], dim=1)
    z = decoder(grid_points).detach().cpu().numpy()

    traces = []
    # plot implicit function contour
    traces.append(go.Contour(x=x, y=y, z=z.reshape(x.shape[0], x.shape[0]),
                             colorscale=colorscale,
                             # autocontour=True,
                             contours=dict(
                                 start=-1,
                                 end=1,
                                 size=0.025,
                             ), showscale=show_bar
                             ))  # contour trace
    # plot clean points and noisy points scatters
    traces.append(go.Scatter(x=clean_points[:, 0], y=clean_points[:, 1],
                             mode='markers', marker=dict(size=16, color=(0, 0, 0))))  # clean points scatter

    # plot normal vectors:
    if n_gt is not None:
        f = ff.create_quiver(clean_points[:, 0], clean_points[:, 1], n_gt[:, 0], n_gt[:, 1],
                             line=dict(color='rgb(0, 255, 0)'))
        traces.append(f.data[0])
    if n_pred is not None:
        f = ff.create_quiver(clean_points[:, 0], clean_points[:, 1], n_pred[:, 0], n_pred[:, 1],
                             line=dict(color='rgb(255, 0, 0)'))
        traces.append(f.data[0])

    layout = go.Layout(width=1200, height=1200,
                       xaxis=dict(side="bottom", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       yaxis=dict(side="left", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                  yaxis=dict(range=[-1, 1], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text=title_text, y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Sherif', size=24, color='red'))
                       )

    fig1 = go.Figure(data=traces, layout=layout)
    offline.plot(fig1, auto_open=False)
    img = utils.plotly_fig2array(fig1)
    return img


# def plot_contour_scatter_and_curl(clean_points, decoder, latent,
#                                           example_idx=0, n_gt=None, n_pred=None, z_gt=None,
#                                           show_ax=True, show_bar=True, title_text='', colorscale='Geyser',
#                                           res=128, nonmnfld_points=None):
#     # plot contour and scatter plot given input points, SAL decoder and the point's laten representation
#
#     clean_points = clean_points[example_idx].detach().cpu().numpy()
#
#     if n_gt is not None:
#         n_gt = n_gt[example_idx].detach().cpu().numpy()
#     if n_pred is not None:
#         n_pred = n_pred[example_idx].detach().cpu().numpy()
#     x, y, grid_points = utils.get_2d_grid_uniform(resolution=res, range=1.2, device=latent.device)
#     grid_points.requires_grad_()
#     grid_points_latent = torch.cat([latent[example_idx].expand(grid_points.shape[0], -1), grid_points], dim=1)
#     z = decoder(grid_points_latent)
#     z_np = z.detach().cpu().numpy()
#     traces = []
#     # plot implicit function contour
#
#     traces.append(go.Contour(x=x, y=y, z=z_np.reshape(x.shape[0], x.shape[0]),
#                                colorscale=colorscale,
#                                # autocontour=True,
#                                contours=dict(
#                                    start=-1,
#                                    end=1,
#                                    size=0.025,
#                                ), showscale=show_bar
#                                ))  # contour trace
#
#     # plot clean points and noisy points scatters
#     clean_points_scatter = go.Scatter(x=clean_points[:, 0], y=clean_points[:, 1],
#                                 mode='markers', marker=dict(size=16, color=(0, 0, 0)))
#     traces.append(clean_points_scatter) # clean points scatter
#
#     if nonmnfld_points is not None:
#         nonmnfld_points = nonmnfld_points[example_idx].detach().cpu().numpy()
#         nonmnfld_points_scatter = go.Scatter(x=nonmnfld_points[:, 0], y=nonmnfld_points[:, 1],
#                                           mode='markers', marker=dict(size=16, color='rgb(255, 255, 255)'))
#         traces.append(nonmnfld_points_scatter)  # nonmnfld points scatter
#
#
#     # plot normal vectors:
#     if n_gt is not None:
#         f = ff.create_quiver(clean_points[:, 0], clean_points[:, 1], n_gt[:, 0], n_gt[:, 1], line=dict(color='rgb(0, 255, 0)'))
#         traces.append(f.data[0])
#     if n_pred is not None:
#         f = ff.create_quiver(clean_points[:, 0], clean_points[:, 1], n_pred[:, 0], n_pred[:, 1],
#                              line=dict(color='rgb(255, 0, 0)'))
#         traces.append(f.data[0])
#
#     layout = go.Layout(width=1200, height=1200,
#                        xaxis=dict(side="bottom", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
#                        yaxis=dict(side="left", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
#                        scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
#                                   yaxis=dict(range=[-1, 1], autorange=False),
#                                   aspectratio=dict(x=1, y=1)),
#                        showlegend=False,
#                        title=dict(text=title_text, y=0.95, x=0.5, xanchor='center', yanchor='middle',
#                                   font=dict(family='Sherif', size=24, color='red'))
#                        )
#
#     fig1 = go.Figure(data=traces, layout=layout)
#     offline.plot(fig1, auto_open=False)
#     dist_img = utils.plotly_fig2array(fig1)
#
#     # plot curl
#     grid_grad = utils.gradient(grid_points, z)
#     dx = utils.gradient(grid_points, grid_grad[:, 0], create_graph=False, retain_graph=True)
#     dy = utils.gradient(grid_points, grid_grad[:, 1], create_graph=False, retain_graph=False)
#     grid_curl = (dx[:, 1] - dy[:, 0]).cpu().detach().numpy()
#
#     traces = [clean_points_scatter]
#     traces.append(go.Contour(x=x, y=y, z=grid_curl.reshape(x.shape[0], x.shape[0]),
#                                colorscale=colorscale,
#                                # autocontour=True,
#                                contours=dict(
#                                    start=-1e-5,
#                                    end=1e-5,
#                                    size=0.25*1e-5,
#                                ), showscale=show_bar
#                                ))  # contour trace
#     # f = ff.create_quiver(grid_points[:, 0].detach().cpu().numpy(), grid_points[:, 1].detach().cpu().numpy(),
#     # grid_grad[:, 0].detach().cpu().numpy(), grid_grad[:, 1].detach().cpu().numpy(),
#     #                      line=dict(color='rgb(0, 0, 0)'))
#     # traces.append(f.data[0])
#     fig2 = go.Figure(data=traces, layout=layout)
#     offline.plot(fig2, auto_open=False)
#     curl_img = utils.plotly_fig2array(fig2)
#
#     # plot eikonal
#     traces = [clean_points_scatter]
#     eikonal_term = ((grid_grad.norm(2, dim=-1) - 1) ** 2).cpu().detach().numpy()
#     traces.append(go.Contour(x=x, y=y, z=eikonal_term.reshape(x.shape[0], x.shape[0]),
#                                colorscale=colorscale,
#                                # autocontour=True,
#                                contours=dict(
#                                    start=-1,
#                                    end=1,
#                                    size=0.025,
#                                ), showscale=show_bar
#                                ))  # contour trace
#     fig3 = go.Figure(data=traces, layout=layout)
#     offline.plot(fig3, auto_open=False)
#     eikonal_img = utils.plotly_fig2array(fig3)
#
#     #plot divergence
#     grid_div = (dx[:, 0] + dy[:, 1]).detach().cpu().numpy()
#     traces = [clean_points_scatter]
#     traces.append(go.Contour(x=x, y=y, z=grid_div.reshape(x.shape[0], x.shape[0]),
#                                colorscale=colorscale,
#                                # autocontour=True,
#                                contours=dict(
#                                    start=-10,
#                                    end=10,
#                                    size=0.25,
#                                ), showscale=show_bar
#                                ))  # contour trace
#     fig4 = go.Figure(data=traces, layout=layout)
#     offline.plot(fig2, auto_open=False)
#     div_img = utils.plotly_fig2array(fig4)
#
#     #plot z difference image
#     z_diff = np.abs(np.abs(np.reshape(z_np, [res, res])) - np.abs(z_gt))
#     traces = [clean_points_scatter]
#     traces.append(go.Contour(x=x, y=y, z=z_diff.reshape(x.shape[0], x.shape[0]),
#                                colorscale=colorscale,
#                                # autocontour=True,
#                                contours=dict(
#                                    start=-0.5,
#                                    end=0.5,
#                                    size=0.025,
#                                ), showscale=show_bar
#                                ))  # contour trace
#     fig5 = go.Figure(data=traces, layout=layout)
#     offline.plot(fig5, auto_open=False)
#     z_diff_img = utils.plotly_fig2array(fig5)
#
#     return dist_img, curl_img, eikonal_img, div_img, z_diff_img, grid_div

def plot_contour_props(x_grid, y_grid, z_grid, clean_points,
                       z_diff,
                       example_idx=0, n_gt=None, n_pred=None,
                       show_ax=True, show_bar=True, title_text='', colorscale='Geyser',
                       nonmnfld_points=None, marker_size=10,
                       plot_second_derivs=False):
    # plot contour and scatter plot given input points, SAL decoder and the point's laten representation
    div_img, curl_img = 0, 0
    if clean_points is not None:
        clean_points = clean_points[example_idx].detach().cpu().numpy()
    if n_gt is not None:
        n_gt = n_gt[example_idx].detach().cpu().numpy()
    if n_pred is not None:
        n_pred = n_pred[example_idx].detach().cpu().numpy()

    traces = []
    print(z_grid.min(), z_grid.max())

    # plot implicit function contour
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid,
                             colorscale=colorscale,
                             # autocontour=True,
                             contours=dict(
                                 start=-1,
                                 end=1,
                                 size=0.025,
                             ), showscale=show_bar
                             ))  # contour trace
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid,
                             contours=dict(start=0, end=0, coloring='lines'),
                             line=dict(width=3),
                             showscale=False,
                             colorscale=[[0, 'rgb(100, 100, 100)'], [1, 'rgb(100, 100, 100)']]))  # black bold zero line

    # plot clean points and noisy points scatters
    if clean_points is not None:
        num = 200
        clean_points_scatter = go.Scatter(x=clean_points[:num, 0], y=clean_points[:num, 1],
                                          mode='markers', marker=dict(size=marker_size, color='rgb(0, 0, 0)'))
        traces.append(clean_points_scatter)  # clean points scatter

    if nonmnfld_points is not None:
        num = 200
        nonmnfld_points = nonmnfld_points[example_idx].detach().cpu().numpy()
        nonmnfld_points_scatter = go.Scatter(x=nonmnfld_points[:num, 0], y=nonmnfld_points[:num, 1],
                                             mode='markers', marker=dict(size=marker_size, color='rgb(255, 255, 255)'))
        traces.append(nonmnfld_points_scatter)  # nonmnfld points scatter

    # plot normal vectors:
    if n_gt is not None:
        num = 20
        f = ff.create_quiver(clean_points[:num, 0], clean_points[:num, 1], n_gt[:num, 0], n_gt[:num, 1],
                             line=dict(color='rgb(0, 255, 0)'))
        traces.append(f.data[0])
    if n_pred is not None:
        num = 20
        f = ff.create_quiver(clean_points[:num, 0], clean_points[:num, 1], n_pred[:num, 0], n_pred[:num, 1],
                             line=dict(color='rgb(255, 0, 0)'))
        traces.append(f.data[0])

    layout = go.Layout(width=800, height=800,
                       xaxis=dict(side="bottom", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       yaxis=dict(side="left", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                  yaxis=dict(range=[-1, 1], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text=title_text, y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Sherif', size=24, color='red'))
                       )

    fig1 = go.Figure(data=traces, layout=layout)
    offline.plot(fig1, auto_open=False)
    dist_img = utils.plotly_fig2array(fig1)
    print("Finished computing sign distance image")
    return dist_img


def plot_contour_div_props(x_grid, y_grid, z_grid, clean_points,
                           z_diff, eikonal_term, grid_div, grid_curl, hessian_det, eigs,
                           example_idx=0, n_gt=None, n_pred=None,
                           show_ax=False, show_bar=True, title_text='', colorscale='GnBu',
                           nonmnfld_points=None, marker_size=10,
                           plot_second_derivs=False, plot_hessian=True):
    # plot contour and scatter plot given input points, SAL decoder and the point's laten representation
    div_img, curl_img = 0, 0
    if clean_points is not None:
        clean_points = clean_points[example_idx].detach().cpu().numpy()
        if clean_points.shape[0] > 200:
            idx = np.random.choice(clean_points.shape[0], 200, replace=False)
            clean_points = clean_points[idx]
    if n_gt is not None:
        n_gt = n_gt[example_idx].detach().cpu().numpy()
    if n_pred is not None:
        n_pred = n_pred[example_idx].detach().cpu().numpy()

    traces = []
    # plot clean points and noisy points scatters
    if clean_points is not None:
        clean_points_scatter = go.Scatter(x=clean_points[:, 0], y=clean_points[:, 1],
                                          mode='markers',
                                          # marker=dict(size=marker_size, color='rgba(0, 0, 0, 0.3)'))
                                          marker = dict(size=marker_size, color='rgba(255, 255, 255, 0.8)'))
        traces.append(clean_points_scatter)  # clean points scatter

    # plot implicit function contour
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid,
                             colorscale=colorscale,
                             # autocontour=True,
                             contours=dict(
                                 start=-.05,
                                 end=.05,
                                 size=0.005,
                                 # start=-1,
                                 # end=1,
                                 # size=0.025,
                                 # showlabels=True,
                             ), showscale=show_bar,
                             reversescale=True
                             ))  # contour trace
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid,
                             contours=dict(start=0, end=0, coloring='lines'),
                             line=dict(width=3),
                             showscale=False,
                             colorscale=[[0, 'rgb(0, 0, 0)'], [1, 'rgb(0, 0, 0)']]))  # black bold zero line

    if nonmnfld_points is not None:
        nonmnfld_points = nonmnfld_points[example_idx].detach().cpu().numpy()
        nonmnfld_points_scatter = go.Scatter(x=nonmnfld_points[:, 0], y=nonmnfld_points[:, 1],
                                             mode='markers', marker=dict(size=marker_size, color='rgb(255, 255, 255)'))
        traces.append(nonmnfld_points_scatter)  # nonmnfld points scatter

    # plot normal vectors:
    if n_gt is not None:
        f = ff.create_quiver(clean_points[:, 0], clean_points[:, 1], n_gt[:, 0], n_gt[:, 1],
                             line=dict(color='rgb(0, 255, 0)'))
        traces.append(f.data[0])
    if n_pred is not None:
        f = ff.create_quiver(clean_points[:, 0], clean_points[:, 1], n_pred[:, 0], n_pred[:, 1],
                             line=dict(color='rgb(255, 0, 0)'))
        traces.append(f.data[0])

    layout = go.Layout(width=800, height=800,
                       xaxis=dict(side="bottom", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       yaxis=dict(side="left", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                  yaxis=dict(range=[-1, 1], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text=title_text, y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Sherif', size=24, color='red'))
                       )

    fig1 = go.Figure(data=traces, layout=layout)
    offline.plot(fig1, auto_open=False)
    dist_img = utils.plotly_fig2array(fig1)
    print("Finished computing sign distance image")
    # plot eikonal
    if clean_points is not None:
        traces = [clean_points_scatter]
    traces.append(go.Contour(x=x_grid, y=y_grid, z=eikonal_term,
                             colorscale=colorscale,
                             # autocontour=True,
                             contours=dict(
                                 start=-1,
                                 end=1,
                                 size=0.025,
                             ), showscale=show_bar
                             ))  # contour trace
    fig3 = go.Figure(data=traces, layout=layout)
    offline.plot(fig3, auto_open=False)
    eikonal_img = utils.plotly_fig2array(fig3)
    print("Finished computing eikonal image")
    # plot z difference image
    if clean_points is not None:
        traces = [clean_points_scatter]
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_diff,
                             colorscale=colorscale,
                             # autocontour=True,
                             contours=dict(
                                 start=-0.5,
                                 end=0.5,
                                 size=0.025,
                             ), showscale=show_bar,
                             reversescale=True
                             ))  # contour trace
    fig5 = go.Figure(data=traces, layout=layout)
    offline.plot(fig5, auto_open=False)
    z_diff_img = utils.plotly_fig2array(fig5)
    print("Finished computing gt distance image")
    if plot_second_derivs:
        # plot divergence
        if clean_points is not None:
            traces = [clean_points_scatter]
        traces.append(go.Contour(x=x_grid, y=y_grid, z=np.clip(grid_div, -10, 10),
                                 colorscale=colorscale,
                                 # autocontour=True,
                                 contours=dict(
                                     start=-10,
                                     end=10,
                                     size=1,
                                 ),
                                 showscale=show_bar
                                 ))  # contour trace
        fig4 = go.Figure(data=traces, layout=layout)
        offline.plot(fig4, auto_open=False)
        div_img = utils.plotly_fig2array(fig4)
        print("Finished divergence image")
        # plot curl
        # if clean_points is not None:
        #     traces = [clean_points_scatter]
        # traces.append(go.Contour(x=x_grid, y=y_grid, z=np.clip(grid_curl, -1e-4, 1e-4),
        #                            colorscale=colorscale,
        #                            # autocontour=True,
        #                            contours=dict(
        #                                start=-1e-4,
        #                                end=1e-4,
        #                                size=1e-5,
        #                            ),
        #                             showscale=show_bar
        #                            ))  # contour trace

        # fig2 = go.Figure(data=traces, layout=layout)
        # offline.plot(fig2, auto_open=False)
        # curl_img = utils.plotly_fig2array(fig2)
        # print("Finished computing curl image")
        curl_img = np.zeros_like(dist_img)
    else:
        # div_img, z_diff_img, curl_img = np.zeros_like(dist_img), np.zeros_like(dist_img), np.zeros_like(dist_img)
        div_img, curl_img = np.zeros_like(div_img), np.zeros_like(dist_img)
    if plot_hessian:
        # plot hessian
        if clean_points is not None:
            traces = [clean_points_scatter]
        traces.append(go.Heatmap(x=x_grid, y=y_grid, z=hessian_det,
                                 colorscale=colorscale,
                                 showscale=show_bar,
                                 reversescale=True,
                                 ))
        # traces.append(go.Contour(x=x_grid, y=y_grid, z=hessian_det,
        #                          colorscale=colorscale,
        #                          contours=dict(
        #                              start=-1,
        #                              end=1,
        #                              size=0.025,
        #                          ), showscale=show_bar,
        #                          reversescale=True
        #                          ))
        fig5 = go.Figure(data=traces, layout=layout)
        offline.plot(fig5, auto_open=False)
        # hessian_det_img = utils.plotly_fig2array(fig5)
        print("Finished hessian_det image")
    else:
        hessian_det_img = None

    if eigs is not None:
        traces = list()
        if clean_points is not None:
            traces = [clean_points_scatter]
        eigs_max = eigs.max(axis=-1)
        eigs_min = eigs.min(axis=-1)
        # plot eigenvalues
        traces.append(go.Heatmap(x=x_grid, y=y_grid, z=eigs_max,
                                 colorscale=colorscale,
                                 showscale=show_bar,
                                 reversescale=True,
                                 text=eigs_max,
                                 ))
        fig6 = go.Figure(data=traces, layout=layout)
        offline.plot(fig6, auto_open=False)
        # eig_max_img = plotly_fig2array(fig4)

        if clean_points is not None:
            traces = [clean_points_scatter]
        traces.append(go.Heatmap(x=x_grid, y=y_grid, z=eigs_min,
                                 colorscale=colorscale,
                                 showscale=show_bar,
                                 reversescale=True,
                                 text=eigs_min,
                                 ))
        fig7 = go.Figure(data=traces, layout=layout)
        offline.plot(fig7, auto_open=False)
        fig7 = go.Figure(data=traces, layout=layout)
        offline.plot(fig7, auto_open=False)
    else:
        fig6 = None
        fig7 = None

    return dist_img, curl_img, eikonal_img, div_img, z_diff_img, fig5, fig6, fig7


def plot_init_contour_div_props(x_grid, y_grid, z_grid, clean_points,
                                z_diff, eikonal_term, grid_div, grid_curl,
                                example_idx=0, n_gt=None, n_pred=None,
                                show_ax=False, show_bar=True, title_text='', colorscale='RdBu',
                                nonmnfld_points=None, marker_size=10,
                                plot_second_derivs=False):
    import time
    t0 = time.time()
    # plot contour and scatter plot given input points, SAL decoder and the point's laten representation
    div_img, curl_img = 0, 0
    if clean_points is not None:
        clean_points = clean_points[example_idx].detach().cpu().numpy()
    if n_gt is not None:
        n_gt = n_gt[example_idx].detach().cpu().numpy()
    if n_pred is not None:
        n_pred = n_pred[example_idx].detach().cpu().numpy()

    traces = []
    # plot clean points and noisy points scatters
    if clean_points is not None:
        clean_points_scatter = go.Scatter(x=clean_points[:, 0], y=clean_points[:, 1],
                                          mode='markers', marker=dict(size=marker_size, color='rgb(0, 0, 0)'))
        traces.append(clean_points_scatter)  # clean points scatter

    # plot implicit function contour
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid,
                             colorbar=dict(tickfont=dict(size=25)),
                             # contours_coloring='heatmap' # For smoothing
                             colorscale=colorscale,
                             # autocontour=True,
                             contours=dict(
                                 start=-1,
                                 end=1,
                                 #    size=0.05, # Siren
                                 #    size=0.1, # geometric sine
                                 size=0.1,  # mfgi
                             ), showscale=show_bar
                             ))  # contour trace
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid,
                             contours=dict(start=0, end=0, coloring='lines'),
                             line=dict(width=3),
                             showscale=False,
                             colorscale=[[0, 'rgb(100, 100, 100)'], [1, 'rgb(100, 100, 100)']]))  # black bold zero line

    if nonmnfld_points is not None:
        nonmnfld_points = nonmnfld_points[example_idx].detach().cpu().numpy()
        nonmnfld_points_scatter = go.Scatter(x=nonmnfld_points[:, 0], y=nonmnfld_points[:, 1],
                                             mode='markers', marker=dict(size=marker_size, color='rgb(255, 255, 255)'))
        traces.append(nonmnfld_points_scatter)  # nonmnfld points scatter

    # plot normal vectors:
    if n_gt is not None:
        f = ff.create_quiver(clean_points[:, 0], clean_points[:, 1], n_gt[:, 0], n_gt[:, 1],
                             line=dict(color='rgb(0, 255, 0)'))
        traces.append(f.data[0])
    if n_pred is not None:
        f = ff.create_quiver(clean_points[:, 0], clean_points[:, 1], n_pred[:, 0], n_pred[:, 1],
                             line=dict(color='rgb(255, 0, 0)'))
        traces.append(f.data[0])

    layout = go.Layout(width=800, height=800,
                       xaxis=dict(side="bottom", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       yaxis=dict(side="left", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                  yaxis=dict(range=[-1, 1], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text=title_text, y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Sherif', size=24, color='red'))
                       )

    fig1 = go.Figure(data=traces, layout=layout)
    offline.plot(fig1, auto_open=False)
    dist_img = utils.plotly_fig2array(fig1)
    print("Finished computing sign distance image, {:.2f}s".format(time.time() - t0))
    t0 = time.time()
    # plot eikonal
    traces = []
    if clean_points is not None:
        traces = [clean_points_scatter]
    traces.append(go.Contour(x=x_grid, y=y_grid, z=eikonal_term,
                             colorbar=dict(tickfont=dict(size=25)),
                             colorscale=colorscale,
                             # autocontour=True,
                             contours=dict(
                                 start=-1,
                                 end=1,
                                 size=0.1,
                             ), showscale=show_bar
                             ))  # contour trace
    fig3 = go.Figure(data=traces, layout=layout)
    offline.plot(fig3, auto_open=False)
    eikonal_img = utils.plotly_fig2array(fig3)
    print("Finished computing eikonal image, {:.2f}s".format(time.time() - t0))
    t0 = time.time()
    # plot z difference image
    traces = []
    if clean_points is not None:
        traces = [clean_points_scatter]
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_diff,
                             colorbar=dict(tickfont=dict(size=25)),
                             colorscale=colorscale,
                             # autocontour=True,
                             contours=dict(
                                 start=-0.5,
                                 end=0.5,
                                 size=0.05,
                             ), showscale=show_bar
                             ))  # contour trace
    fig5 = go.Figure(data=traces, layout=layout)
    offline.plot(fig5, auto_open=False)
    z_diff_img = utils.plotly_fig2array(fig5)
    print("Finished computing gt distance image, {:.2f}s".format(time.time() - t0))
    t0 = time.time()
    if plot_second_derivs:
        # plot divergence
        traces = []
        if clean_points is not None:
            traces = [clean_points_scatter]
        traces.append(go.Contour(x=x_grid, y=y_grid, z=np.clip(grid_div, -50, 50),
                                 colorbar=dict(tickfont=dict(size=25)),
                                 colorscale=colorscale,
                                 # autocontour=True,
                                 contours=dict(
                                     start=-50,
                                     end=50,
                                     size=1,
                                 ),
                                 contours_coloring='heatmap',
                                 #    showscale=show_bar
                                 ))  # contour trace
        fig4 = go.Figure(data=traces, layout=layout)
        offline.plot(fig4, auto_open=False)
        div_img = utils.plotly_fig2array(fig4)
        print("Finished divergence image, {:.2f}s".format(time.time() - t0))
        t0 = time.time()
        # # plot curl
        # traces = []
        # if clean_points is not None:
        #     traces = [clean_points_scatter]
        # traces.append(go.Contour(x=x_grid, y=y_grid, z=np.clip(grid_curl, -1e-4, 1e-4),
        #                         colorbar=dict(tickfont=dict(size=25)),
        #                            colorscale=colorscale,
        #                            # autocontour=True,
        #                            contours=dict(
        #                                start=-1e-4,
        #                                end=1e-4,
        #                                size=1e-5,
        #                            ),
        #                             showscale=show_bar
        #                            ))  # contour trace

        # fig2 = go.Figure(data=traces, layout=layout)
        # offline.plot(fig2, auto_open=False)
        # curl_img = utils.plotly_fig2array(fig2)
        # print("Finished computing curl image, {:.2f}s".format(time.time()-t0))
        # t0 = time.time()
        curl_img = None
    else:
        div_img, z_diff_img, curl_img = np.zeros_like(dist_img), np.zeros_like(dist_img), np.zeros_like(dist_img)
    return dist_img, curl_img, eikonal_img, div_img, z_diff_img


def plot_gt_contour_div_props(x_grid, y_grid, z_grid, clean_points,
                              z_diff, eikonal_term, grid_div, grid_curl,
                              example_idx=0, n_gt=None, n_pred=None,
                              show_ax=True, show_bar=True, title_text='', colorscale='Geyser',
                              nonmnfld_points=None, marker_size=10,
                              plot_second_derivs=False):
    import time
    t0 = time.time()
    # plot contour and scatter plot given input points, SAL decoder and the point's laten representation
    div_img, curl_img = 0, 0
    if clean_points is not None:
        clean_points = clean_points[example_idx].detach().cpu().numpy()
    if n_gt is not None:
        n_gt = n_gt[example_idx].detach().cpu().numpy()
    if n_pred is not None:
        n_pred = n_pred[example_idx].detach().cpu().numpy()

    traces = []

    # plot implicit function contour
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid,
                             colorscale=colorscale,
                             # autocontour=True,
                             contours=dict(
                                 start=-1,
                                 end=1,
                                 size=0.025,
                             ), showscale=show_bar
                             ))  # contour trace
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid,
                             contours=dict(start=0, end=0, coloring='lines'),
                             line=dict(width=3),
                             showscale=False,
                             colorscale=[[0, 'rgb(100, 100, 100)'], [1, 'rgb(100, 100, 100)']]))  # black bold zero line

    # plot clean points and noisy points scatters
    if clean_points is not None:
        clean_points_scatter = go.Scatter(x=clean_points[:, 0], y=clean_points[:, 1],
                                          mode='markers', marker=dict(size=marker_size, color='rgb(0, 0, 0)'))
        traces.append(clean_points_scatter)  # clean points scatter

    if nonmnfld_points is not None:
        nonmnfld_points = nonmnfld_points[example_idx].detach().cpu().numpy()
        nonmnfld_points_scatter = go.Scatter(x=nonmnfld_points[:, 0], y=nonmnfld_points[:, 1],
                                             mode='markers', marker=dict(size=marker_size, color='rgb(255, 255, 255)'))
        traces.append(nonmnfld_points_scatter)  # nonmnfld points scatter

    # plot normal vectors:
    if n_gt is not None:
        f = ff.create_quiver(clean_points[:, 0], clean_points[:, 1], n_gt[:, 0], n_gt[:, 1],
                             line=dict(color='rgb(0, 255, 0)'))
        traces.append(f.data[0])
    if n_pred is not None:
        f = ff.create_quiver(clean_points[:, 0], clean_points[:, 1], n_pred[:, 0], n_pred[:, 1],
                             line=dict(color='rgb(255, 0, 0)'))
        traces.append(f.data[0])

    layout = go.Layout(width=800, height=800,
                       xaxis=dict(side="bottom", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       yaxis=dict(side="left", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                  yaxis=dict(range=[-1, 1], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text=title_text, y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Sherif', size=24, color='red'))
                       )

    fig1 = go.Figure(data=traces, layout=layout)
    offline.plot(fig1, auto_open=False)
    dist_img = utils.plotly_fig2array(fig1)
    print("Finished computing sign distance image, {:.2f}s".format(time.time() - t0))
    t0 = time.time()
    # plot eikonal
    traces = []
    if clean_points is not None:
        traces = [clean_points_scatter]
    traces.append(go.Contour(x=x_grid, y=y_grid, z=eikonal_term,
                             colorscale=colorscale,
                             # autocontour=True,
                             contours=dict(
                                 start=0,
                                 end=2,
                                 size=0.025,
                             ), showscale=show_bar
                             ))  # contour trace
    fig3 = go.Figure(data=traces, layout=layout)
    offline.plot(fig3, auto_open=False)
    eikonal_img = utils.plotly_fig2array(fig3)
    print("Finished computing eikonal image, {:.2f}s".format(time.time() - t0))
    t0 = time.time()
    # plot z difference image
    traces = []
    if clean_points is not None:
        traces = [clean_points_scatter]
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_diff,
                             colorscale=colorscale,
                             # autocontour=True,
                             contours=dict(
                                 start=-0.5,
                                 end=0.5,
                                 size=0.025,
                             ), showscale=show_bar
                             ))  # contour trace
    fig5 = go.Figure(data=traces, layout=layout)
    offline.plot(fig5, auto_open=False)
    z_diff_img = utils.plotly_fig2array(fig5)
    print("Finished computing gt distance image, {:.2f}s".format(time.time() - t0))
    t0 = time.time()
    if plot_second_derivs:
        # plot divergence
        traces = []
        if clean_points is not None:
            traces = [clean_points_scatter]
        traces.append(go.Contour(x=x_grid, y=y_grid, z=np.clip(grid_div, -50, 50),
                                 colorscale=colorscale,
                                 # autocontour=True,
                                 contours=dict(
                                     start=-50,
                                     end=50,
                                     size=1,
                                 ),
                                 contours_coloring='heatmap',
                                 #    showscale=show_bar
                                 ))  # contour trace
        fig4 = go.Figure(data=traces, layout=layout)
        offline.plot(fig4, auto_open=False)
        div_img = utils.plotly_fig2array(fig4)
        print("Finished divergence image, {:.2f}s".format(time.time() - t0))
        t0 = time.time()
        # plot curl
        traces = []
        if clean_points is not None:
            traces = [clean_points_scatter]
        traces.append(go.Contour(x=x_grid, y=y_grid, z=np.clip(grid_curl, -1e-4, 1e-4),
                                 colorscale=colorscale,
                                 # autocontour=True,
                                 contours=dict(
                                     start=-1e-4,
                                     end=1e-4,
                                     size=1e-5,
                                 ),
                                 showscale=show_bar
                                 ))  # contour trace

        fig2 = go.Figure(data=traces, layout=layout)
        offline.plot(fig2, auto_open=False)
        curl_img = utils.plotly_fig2array(fig2)
        print("Finished computing curl image, {:.2f}s".format(time.time() - t0))
        t0 = time.time()
        # curl_img = None
    else:
        div_img, z_diff_img, curl_img = np.zeros_like(dist_img), np.zeros_like(dist_img), np.zeros_like(dist_img)
    return dist_img, curl_img, eikonal_img, div_img, z_diff_img


def plot_paper_teaser_images(x_grid, y_grid, z_grid, clean_points, grid_normals, colorscale='Geyser', show_ax=False,
                             output_path='/mnt/3.5TB_WD/PycharmProjects/DiBS/sanitychecks/vis/'):
    os.makedirs(output_path, exist_ok=True)
    layout1 = go.Layout(width=1200, height=1200,
                        xaxis=dict(side="bottom", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                        yaxis=dict(side="left", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                        scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                   yaxis=dict(range=[-1, 1], autorange=False),
                                   aspectratio=dict(x=1, y=1)),
                        showlegend=False,
                        title=dict(text='', y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                   font=dict(family='Sherif', size=24, color='red'))
                        )
    layout2 = go.Layout(width=1200, height=1200, plot_bgcolor='rgb(255, 255, 255)',
                        xaxis=dict(side="bottom", range=[-1, 1], showgrid=True, zeroline=show_ax, visible=True,
                                   gridcolor='LightGrey'),
                        yaxis=dict(side="left", range=[-1, 1], showgrid=True, zeroline=show_ax, visible=True,
                                   gridcolor='LightGrey'),
                        scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                   yaxis=dict(range=[-1, 1], autorange=False),
                                   aspectratio=dict(x=1, y=1)),
                        showlegend=False,
                        title=dict(text='', y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                   font=dict(family='Sherif', size=24, color='red'))
                        )
    traces2 = []
    traces = []
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid, colorscale=colorscale,
                             contours=dict(start=-3, end=3, size=0.1), showscale=False))  # contour trace
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid,
                             contours=dict(start=0, end=0, coloring='lines'),
                             line=dict(width=5, dash="dash"),
                             showscale=False,
                             colorscale=[[0, 'rgb(100, 100, 100)'], [1, 'rgb(100, 100, 100)']]))  # contour trace (zls)

    # plot clean points and noisy points scatters
    scatter_points = go.Scatter(x=clean_points[:, 0], y=clean_points[:, 1],
                                mode='markers', marker=dict(size=40, color='rgb(50, 50, 50)'))
    traces.append(scatter_points)
    traces2.append(scatter_points)
    traces.append(
        ff.create_quiver(x_grid, y_grid, grid_normals[:, 0], grid_normals[:, 1], line=dict(color='rgb(0, 0, 255)'),
                         scale=.04).data[0])

    fig1 = go.Figure(data=traces, layout=layout1)
    fig2 = go.Figure(data=traces2, layout=layout2)
    # fig2.show()
    offline.plot(fig1, auto_open=False)
    offline.plot(fig2, auto_open=False)
    img1 = utils.plotly_fig2array(fig1)
    img2 = utils.plotly_fig2array(fig2)
    im = Image.fromarray(img1)
    im.save(os.path.join(output_path, "teaser1.png"))
    im = Image.fromarray(img2)
    im.save(os.path.join(output_path, "teaser2.png"))


def plot_shape_data(x_grid, y_grid, z_grid, clean_points, n_gt=None, show_ax=True, show_bar=True,
                    title_text='', colorscale='Geyser', nonmnfld_points=None, divergence=None, grid_normals=None):
    # plot contour and scatter plot given input points

    traces = []
    # plot implicit function contour
    if divergence is None:
        traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid, colorscale=colorscale,
                                 contours=dict(start=-1, end=1, size=0.025), showscale=show_bar))  # contour trace
    else:
        traces.append(go.Contour(x=x_grid, y=y_grid, z=divergence, colorscale=colorscale,
                                 contours=dict(start=-3, end=3, size=0.25), showscale=show_bar))  # contour trace

    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid,
                             contours=dict(start=0, end=0, coloring='lines'),
                             line=dict(width=5),
                             showscale=False,
                             colorscale=[[0, 'rgb(100, 100, 100)'], [1, 'rgb(100, 100, 100)']]))  # contour trace

    # plot clean points and noisy points scatters
    clean_points_scatter = go.Scatter(x=clean_points[:, 0], y=clean_points[:, 1],
                                      mode='markers', marker=dict(size=16, color='rgb(0, 0, 0)'))
    traces.append(clean_points_scatter)  # clean points scatter

    if nonmnfld_points is not None:
        # plot clean points and noisy points scatters
        nonmnfld_points_scatter = go.Scatter(x=nonmnfld_points[:, 0], y=nonmnfld_points[:, 1],
                                             mode='markers', marker=dict(size=16, color='rgb(255, 255, 255)'))
        traces.append(nonmnfld_points_scatter)  # clean points scatter

    # plot normal vectors:
    if n_gt is not None:
        f = ff.create_quiver(clean_points[:, 0], clean_points[:, 1], n_gt[:, 0], n_gt[:, 1],
                             line=dict(color='rgb(0, 255, 0)'))
        traces.append(f.data[0])

    if grid_normals is not None:
        f = ff.create_quiver(x_grid, y_grid, grid_normals[:, 0], grid_normals[:, 1], line=dict(color='rgb(0, 0, 255)'),
                             scale=.01)
        traces.append(f.data[0])

    layout = go.Layout(width=1200, height=1200,
                       xaxis=dict(side="bottom", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       yaxis=dict(side="left", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                  yaxis=dict(range=[-1, 1], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text=title_text, y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Sherif', size=24, color='red'))
                       )

    fig = go.Figure(data=traces, layout=layout)
    fig.show()


def plot_sdf_indicator(points, x_grid, y_grid, sdf, title_text='', show_ax=False, output_path='./'):
    points = np.concatenate([points, points[0, :][None, :]], axis=0)
    indicator = -np.ones_like(sdf)
    indicator[sdf > 0] = 1

    layout = go.Layout(width=1200, height=1200,
                       xaxis=dict(side="bottom", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=True,
                                  mirror=True, showline=True, linewidth=2, linecolor='black', showticklabels=False),
                       yaxis=dict(side="left", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=True,
                                  mirror=True, showline=True, linewidth=2, linecolor='black', showticklabels=False),
                       scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                  yaxis=dict(range=[-1, 1], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text=title_text, y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Sherif', size=24, color='red')),
                       paper_bgcolor='rgb(255,255,255)',
                       plot_bgcolor='rgb(255, 255, 255)'
                       )

    # plot clean points and noisy points scatters
    clean_points_scatter = go.Scatter(x=points[:, 0], y=points[:, 1],
                                      mode='markers', marker=dict(size=52, color='rgb(0, 0, 0)'))
    fig = go.Figure(data=clean_points_scatter, layout=layout)
    fig.write_image(os.path.join(output_path, "points.png"))
    fig.data = []
    line_plot = go.Scatter(x=points[:, 0], y=points[:, 1],
                           mode='markers+lines', marker=dict(size=52, color='rgb(0, 0, 0)'),
                           fill="toself", fillcolor='rgb(255,255,255)',
                           line=dict(width=6, color='rgb(0, 0, 0)'))

    fig.add_traces(line_plot)
    fig.write_image(os.path.join(output_path, "lines.png"))
    fig.data = []

    sdf_trace = go.Contour(x=x_grid, y=y_grid, z=sdf, colorscale='Geyser',
                           contours=dict(start=-1, end=1, size=0.025), showscale=False)
    sdf_zls_trace = go.Contour(x=x_grid, y=y_grid, z=sdf,
                               contours=dict(start=0, end=0, coloring='lines'),
                               line=dict(width=5),
                               showscale=False,
                               colorscale=[[0, 'rgb(100, 100, 100)'], [1, 'rgb(100, 100, 100)']])  # contour trace
    fig.add_traces([sdf_trace, sdf_zls_trace, clean_points_scatter])
    fig.write_image(os.path.join(output_path, "sdf.png"))
    fig.data = []

    indicator_trace = go.Contour(x=x_grid, y=y_grid, z=indicator, colorscale='Geyser',
                                 contours=dict(start=-1, end=1, size=0.025), showscale=False)
    indicator_zls_trace = go.Contour(x=x_grid, y=y_grid, z=indicator,
                                     contours=dict(start=0, end=0, coloring='lines'),
                                     line=dict(width=5),
                                     showscale=False,
                                     colorscale=[[0, 'rgb(100, 100, 100)'], [1, 'rgb(100, 100, 100)']])  # contour trace
    fig.add_traces([indicator_trace, indicator_zls_trace, clean_points_scatter])
    fig.write_image(os.path.join(output_path, "indicator.png"))

    # fig = go.Figure(data=traces, layout=layout)
    # fig.show()


def plot_distance_map(x_grid, y_grid, z_grid, title_text='', colorscale='Geyser', show_ax=True):
    # plot contour and scatter plot given input points

    traces = []
    # plot implicit function contour
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid, colorscale=colorscale,
                             contours=dict(start=-1, end=1, size=0.025), showscale=True))
    traces.append(go.Contour(x=x_grid, y=y_grid, z=z_grid,
                             contours=dict(start=0, end=0, coloring='lines'),
                             line=dict(width=5),
                             showscale=False,
                             colorscale=[[0, 'rgb(100, 100, 100)'], [1, 'rgb(100, 100, 100)']]))  # contour trace

    layout = go.Layout(width=1200, height=1200,
                       xaxis=dict(side="bottom", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       yaxis=dict(side="left", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                  yaxis=dict(range=[-1, 1], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text=title_text, y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Sherif', size=24, color='red'))
                       )

    fig = go.Figure(data=traces, layout=layout)
    fig.show()


def plot_points(points, title_txt='circle', example_idx=0, show_ax=True):
    points = points[example_idx]

    traces = []
    traces.append(go.Scatter(x=points[:, 0], y=points[:, 1],
                             mode='markers', marker=dict(size=16, color='rgb(0, 0, 0)')))  # points scatter

    # plot jet
    layout = go.Layout(width=1200, height=1200,
                       xaxis=dict(side="bottom", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       yaxis=dict(side="left", range=[-1, 1], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
                                  yaxis=dict(range=[-1, 1], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text='Shape: ' + title_txt, y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Sherif', size=24, color='red'))
                       )
    fig = go.Figure(data=traces, layout=layout)
    fig.show()


def plot_mesh(mesh_trace, mesh=None, output_ply_path=None, show_ax=False, title_txt='', show=True):
    # plot mesh and export ply file for mesh

    # export mesh
    if not (mesh is None or output_ply_path is None):
        mesh.export(output_ply_path, vertex_normal=True)

    # plot mesh
    layout = go.Layout(width=1200, height=1200,
                       xaxis=dict(side="bottom", range=[-2, 2], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       yaxis=dict(side="left", range=[-2, 2], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       scene=dict(xaxis=dict(range=[-2, 2], autorange=False),
                                  yaxis=dict(range=[-2, 2], autorange=False),
                                  zaxis=dict(range=[-2, 2], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text='Shape: ' + title_txt, y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Sherif', size=24, color='red'))
                       )
    fig = go.Figure(data=mesh_trace, layout=layout)
    if show:
        fig.show()


def plot_sdf_surface(x, y, z, show=True, show_ax=True, title_txt=''):
    # plot mesh
    layout = go.Layout(width=1200, height=1200,
                       xaxis=dict(side="bottom", range=[-2, 2], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       yaxis=dict(side="left", range=[-2, 2], showgrid=show_ax, zeroline=show_ax, visible=show_ax),
                       scene=dict(xaxis=dict(range=[-2, 2], autorange=False),
                                  yaxis=dict(range=[-2, 2], autorange=False),
                                  zaxis=dict(range=[-2, 2], autorange=False),
                                  aspectratio=dict(x=1, y=1)),
                       showlegend=False,
                       title=dict(text=title_txt, y=0.95, x=0.5, xanchor='center', yanchor='middle',
                                  font=dict(family='Sherif', size=24, color='red'))
                       )
    surf_trace = go.Surface(z=z, x=x, y=y)
    fig = go.Figure(data=surf_trace, layout=layout)
    if show:
        fig.show()

    return fig

utils/visualizations.py
bearprin/Neural-Singular-Hessian