# This file is partly based on DiGS: https://github.com/Chumbyte/DiGS

import numpy as np
import torch
import torch.nn as nn
from models.SIREN import FCBlock

class Decoder(nn.Module):
    def forward(self, *args, **kwargs):
        return self.fc_block(*args, **kwargs)

class Network(nn.Module):
    def __init__(self, latent_size, in_dim=3, decoder_hidden_dim=256, nl='sine', encoder_type=None,
                 decoder_n_hidden_layers=8, init_type='siren', neuron_type='quadratic', sphere_init_params=[1.6, 1.0]):
        super().__init__()
        self.encoder_type = encoder_type
        self.init_type = init_type
        if encoder_type == 'autodecoder':
            # latent_size will stay as input latent size
            pass
        elif encoder_type == 'none':
            latent_size = 0
        else:
            raise ValueError("unsupported encoder type")
        self.decoder = Decoder()
        if neuron_type == 'linear':
            self.decoder.fc_block = FCBlock(in_dim + latent_size, 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
        elif neuron_type == 'quadratic':
            self.decoder.fc_block = QuaNet(d_in=in_dim + latent_size, nl=nl, n_layers=decoder_n_hidden_layers, d_hidden=decoder_hidden_dim,
                                                    init_type=init_type, sphere_init_params=sphere_init_params)
        else:
            raise ValueError("unsupported neuron type")
                                                    

    def forward(self, non_mnfld_pnts, mnfld_pnts=None):
        # shape is (bs, npoints, in_dim+latent_size) for both inputs, npoints could be different sizes
        batch_size = non_mnfld_pnts.shape[0]
        if not mnfld_pnts is None and self.encoder_type == 'autodecoder':
            # Assume inputs have latent vector concatted with [xyz, latent]
            latent = non_mnfld_pnts[:,:,3:]
            latent_reg = latent.norm(dim=-1).mean()
            manifold_pnts_pred = self.decoder(mnfld_pnts.view(-1, mnfld_pnts.shape[-1])).reshape(batch_size, -1)
        elif mnfld_pnts is not None:
            manifold_pnts_pred = self.decoder(mnfld_pnts.view(-1, mnfld_pnts.shape[-1])).reshape(batch_size, -1)
            latent = None
            latent_reg = None
        else:
            manifold_pnts_pred = None
            latent = None
            latent_reg = None

        # Off manifold points
        nonmanifold_pnts_pred = self.decoder(non_mnfld_pnts.view(-1, non_mnfld_pnts.shape[-1])).reshape(batch_size, -1)

        return {"manifold_pnts_pred": manifold_pnts_pred,
                "nonmanifold_pnts_pred": nonmanifold_pnts_pred,
                "latent_reg": latent_reg,
                "latent": latent}


class QuaNet(nn.Module):
    def __init__(
            self,
            d_in=3,
            d_out=1,
            d_hidden=256,
            n_layers=8,
            nl='softplus',
            sphere_init_params=[1.6, 1.0],
            init_type='siren',
    ):
        super().__init__()
        dims = [d_in] + [d_hidden for _ in range(n_layers)] + [d_out]
        self.num_layers = len(dims)
        self.init_type = init_type
        self.sphere_init_params = sphere_init_params

        for l in range(self.num_layers - 1):
            if l == 0:
                qua = QuadraticLayer(dims[l], dims[l + 1])
            else:
                qua = QuadraticLayer(dims[l], dims[l + 1])
            setattr(self, "qua" + str(l), qua)
        if self.init_type == "mfgi":
            self.apply(geom_sine_init)
            self.qua0.apply(first_layer_mfgi_init)
            self.qua1.apply(second_layer_mfgi_init)
            getattr(self, "qua" + str(self.num_layers - 3)).apply(second_last_layer_geom_sine_init)
            getattr(self, "qua" + str(self.num_layers - 2)).apply(last_layer_geom_sine_init)
        elif self.init_type == "siren":
            self.apply(sine_init)
            self.qua0.apply(first_layer_sine_init)
        elif init_type == 'geometric_sine':
            self.apply(geom_sine_init)
            self.qua0.apply(first_layer_geom_sine_init)
            getattr(self, "qua" + str(self.num_layers - 3)).apply(second_last_layer_geom_sine_init)
            getattr(self, "qua" + str(self.num_layers - 2)).apply(last_layer_geom_sine_init)

        if nl == 'softplus':
            self.activation = nn.Softplus(beta=100)
        elif nl == 'tanh':
            self.activation = nn.Tanh()
        elif nl == "sine":
            self.activation = Sine()
        elif nl == "relu":
            self.activation = nn.ReLU()

    def forward(self, inputs):
        x = inputs
        for l in range(self.num_layers - 1):
            qua = getattr(self, "qua" + str(l))
            if l == self.num_layers - 2:
                x = qua(x)
            else:
                x = qua(x)
                x = self.activation(x)

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

    def gradient(self, x):
        x.requires_grad_(True)
        y = self.sdf(x)
        d_output = torch.ones_like(y, requires_grad=False, device=y.device)
        gradients = torch.autograd.grad(
            outputs=y,
            inputs=x,
            grad_outputs=d_output,
            create_graph=True,
            retain_graph=True,
            only_inputs=True,
        )[0]
        return gradients.unsqueeze(1)


class QuadraticLayer(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.lin1 = nn.Linear(d_in, d_out)
        self.lin2 = nn.Linear(d_in, d_out)
        self.lin3 = nn.Linear(d_in, d_out)

    def forward(self, x):
        return torch.mul(self.lin1(x), self.lin2(x)) + self.lin3(torch.square(x))


def init_lin2_lin3(m):
    nn.init.normal_(m.lin2.weight, mean=0.0, std=1e-5)
    nn.init.ones_(m.lin2.bias)
    nn.init.normal_(m.lin3.weight, mean=0.0, std=1e-5)
    nn.init.zeros_(m.lin3.bias)


################################# SIREN's initialization ###################################
# This file is borrowed from DiGS: https://github.com/Chumbyte/DiGS
def sine_init(m):
    with torch.no_grad():
        if hasattr(m, "lin1") and hasattr(m, "lin2") and hasattr(m, "lin3"):
            num_input = m.lin1.weight.size(-1)
            # See SIREN paper supplement Sec. 1.5 for discussion of factor 30
            m.lin1.weight.uniform_(-np.sqrt(6 / num_input) / 30, np.sqrt(6 / num_input) / 30)
            init_lin2_lin3(m)


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


################################# sine geometric initialization ###################################
# This file is borrowed from DiGS: https://github.com/Chumbyte/DiGS

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


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


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


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


################################# multi frequency geometric initialization ###################################
# This file is borrowed from DiGS: https://github.com/Chumbyte/DiGS
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, "lin1") and hasattr(m, "lin2") and hasattr(m, "lin3"):
            num_input = m.lin1.weight.size(-1)
            num_output = m.lin1.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.lin1.weight.data = W0_new
            init_lin2_lin3(m)


def second_layer_mfgi_init(m):
    global portion_per_period
    with torch.no_grad():
        if hasattr(m, "lin1") and hasattr(m, "lin2") and hasattr(m, "lin3"):
            num_input = m.lin1.weight.size(-1)
            assert m.lin1.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.lin1.weight.data = W1_new
            init_lin2_lin3(m)


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)



models/Net.py
sunyx523/StEik
# This file is borrowed from DiGS: https://github.com/Chumbyte/DiGS
import numpy as np
import torch
import torch.nn as nn
from models.torchmeta.modules.module import MetaModule
from models.torchmeta.modules.container import MetaSequential
from models.torchmeta.modules.utils import get_subdict
from collections import OrderedDict


class FCBlock(MetaModule):
    '''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(MetaSequential(BatchLinear(in_features, hidden_features), nl))

        for i in range(num_hidden_layers):
            self.net.append(MetaSequential(BatchLinear(hidden_features, hidden_features), nl))

        if outermost_linear:
            self.net.append(MetaSequential(BatchLinear(hidden_features, out_features)))
        else:
            self.net.append(MetaSequential(BatchLinear(hidden_features, out_features), nl))

        self.net = MetaSequential(*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, params=None, **kwargs):
        if params is None:
            params = OrderedDict(self.named_parameters())

        output = self.net(coords, params=get_subdict(params, 'net'))

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

        return output


class BatchLinear(nn.Linear, MetaModule):
    '''A linear meta-layer that can deal with batched weight matrices and biases, as for instance output by a
    hypernetwork.'''
    __doc__ = nn.Linear.__doc__

    def forward(self, input, params=None):
        if params is None:
            params = OrderedDict(self.named_parameters())

        bias = params.get('bias', None)
        weight = params['weight']

        output = input.matmul(weight.permute(*[i for i in range(len(weight.shape) - 2)], -1, -2))
        output += bias.unsqueeze(-2)
        return output


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/SIREN.py
sunyx523/StEik
# This file is partly based on DiGS: https://github.com/Chumbyte/DiGS

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
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 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


def directional_div(points, grads):
    dot_grad = (grads * grads).sum(dim=-1, keepdim=True)
    hvp = torch.ones_like(dot_grad)
    hvp = 0.5 * torch.autograd.grad(dot_grad, points, hvp, retain_graph=True, create_graph=True)[0]
    div = (grads * hvp).sum(dim=-1) / (torch.sum(grads ** 2, dim=-1) + 1e-5)
    return div

def full_div(points, grads):
    dx = utils.gradient(points, grads[:, :, 0])
    dy = utils.gradient(points, grads[:, :, 1])
    if points.shape[-1] == 3:
        dz = utils.gradient(points, grads[:, :, 2])
        div = dx[:, :, 0] + dy[:, :, 1] + dz[:, :, 2]
    else:
        div = dx[:, :, 0] + dy[:, :, 1]
    div[div.isnan()] = 0
    return div

class Loss(nn.Module):
    def __init__(self, weights=[3e3, 1e2, 1e2, 5e1, 1e2], loss_type='siren', div_decay='none', div_type='dir_l1'):
        super().__init__()
        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_div = True if 'div' in self.loss_type else False

    def forward(self, output_pred, mnfld_points, nonmnfld_points, mnfld_n_gt=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"]
        latent = output_pred["latent"]

        div_loss = 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)

        # div_term
        if self.use_div and self.weights[4] > 0.0:

            if self.div_type == 'full_l2':
                nonmnfld_divergence = full_div(nonmnfld_points, nonmnfld_grad)
                nonmnfld_divergence_term = torch.clamp(torch.square(nonmnfld_divergence), 0.1, 50)
            elif self.div_type == 'full_l1':
                nonmnfld_divergence = full_div(nonmnfld_points, nonmnfld_grad)
                nonmnfld_divergence_term = torch.clamp(torch.abs(nonmnfld_divergence), 0.1, 50)
            elif self.div_type == 'dir_l2':
                nonmnfld_divergence = directional_div(nonmnfld_points, nonmnfld_grad)
                nonmnfld_divergence_term = torch.square(nonmnfld_divergence)
            elif self.div_type == 'dir_l1':
                nonmnfld_divergence = directional_div(nonmnfld_points, nonmnfld_grad)
                nonmnfld_divergence_term = torch.abs(nonmnfld_divergence)
            else:
                raise Warning("unsupported divergence type. only suuports dir_l1, dir_l2, full_l1, full_l2")

            div_loss = nonmnfld_divergence_term.mean() #+ mnfld_divergence_term.mean()

        # eikonal term
        eikonal_term = eikonal_loss(nonmnfld_grad, mnfld_grad=mnfld_grad, eikonal_type='abs')

        # 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()

        # 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_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 == '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 == '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_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:
            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 is not None and latent_reg is not None:
            loss += self.weights[5] * 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
                }, mnfld_grad

    def update_div_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[4] = w0
            elif  current_iteration >= s*n_iterations and current_iteration < e*n_iterations:
                self.weights[4] = w0 + (we - w0) * (current_iteration/n_iterations - s) / (e - s)
            else:
                self.weights[4] = we
        elif self.div_decay == 'quintic': # linearly decrease weight from iter s to iter e
            if current_iteration < s*n_iterations:
                self.weights[4] = w0
            elif  current_iteration >= s*n_iterations and current_iteration < e*n_iterations:
                self.weights[4] = w0 + (we - w0) * (1 - (1 -(current_iteration/n_iterations - s) / (e - s))**5)
            else:
                self.weights[4] = we
        elif self.div_decay == 'step': # change weight at s
            if current_iteration < s*n_iterations:
                self.weights[4] = w0
            else:
                self.weights[4] = we
        elif self.div_decay == 'none':
            pass
        else:
            raise Warning("unsupported div decay value")
models/losses.py
sunyx523/StEik
import torch.nn as nn
import torch
from models.torchmeta.modules.module import MetaModule
from models.torchmeta.modules.utils import get_subdict

class MetaSequential(nn.Sequential, MetaModule):
    __doc__ = nn.Sequential.__doc__

    def forward(self, input, params=None):
        for name, module in self._modules.items():
            if isinstance(module, MetaModule):
                input = module(input, params=get_subdict(params, name))
            elif isinstance(module, nn.Module):
                input = module(input)
            else:
                raise TypeError('The module must be either a torch module '
                    '(inheriting from `nn.Module`), or a `MetaModule`. '
                    'Got type: `{0}`'.format(type(module)))
        return input
models/torchmeta/modules/container.py
sunyx523/StEik
# This file is partly based on DiGS: https://github.com/Chumbyte/DiGS
import torch.utils.data as data
import numpy as np
import scipy.spatial as spatial
import utils.visualizations as vis
from abc import ABC, abstractmethod
from matplotlib.path import Path
import torch


class BasicShape2D(data.Dataset):
    # A class to generate synthetic examples of basic shapes.
    # Generates clean and noisy point clouds sampled on Jets + samples no a grid with their distance to the surface
    def __init__(self, n_points, n_samples=128, res=128, sample_type='grid', sapmling_std=0.005,
                 grid_range=1.2):


        self.grid_range = grid_range
        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
        # Generate shape

        self.points = self.get_mnfld_points()

        # generate grid points and find distance to closest point on the line
        x, y = np.linspace(-grid_range, grid_range, self.grid_res), np.linspace(-grid_range, grid_range, self.grid_res)
        xx, yy = np.meshgrid(x, y)
        xx, yy = xx.ravel(), yy.ravel()
        self.grid_points = np.stack([xx, yy], axis=1).astype('f')
        self.nonmnfld_points = self.get_nonmnfld_points()

        # Compute gt mnfld normals
        self.mnfld_n = self.get_mnfld_n()

        self.grid_dist, self.grid_n = self.get_points_distances_and_normals(self.grid_points)
        self.nonmnfld_dist, self.nonmnfld_n = self.get_points_distances_and_normals(self.nonmnfld_points)
        self.dist_img = np.reshape(self.grid_dist, [self.grid_res, self.grid_res])

        self.point_idxs = np.arange(self.points.shape[1])
        self.grid_points_idxs = np.arange(self.grid_points.shape[0])
        self.nonmnfld_points_idxs = np.arange(self.nonmnfld_points.shape[0])
        self.sample_probs = np.ones_like(self.grid_points_idxs) / self.grid_points.shape[0]

        self.generate_batch_indices()

    @abstractmethod
    def get_mnfld_points(self):
        # implement a function that returns points on the manifold
        pass

    @abstractmethod
    def get_mnfld_n(self):
        #implement a function that returns normal vectors for points on the manifold
        pass

    @abstractmethod
    def get_points_distances_and_normals(self, points):
        # implement a function that computes the distance and normal vectors of nonmanifold points.
        # default implementation finds the nearest neighbor and return its normal and the distance to it.
        # which is a coarse approxiamation

        distances = []
        normals = []
        # compute distance and normal (general case)
        for i, point_cloud in enumerate(self.points):
            kdtree = spatial.cKDTree(point_cloud)
            distances, nn_idx = kdtree.query(points, k=1)
            signs = np.sign(np.einsum('ij,ij->i', points - point_cloud[nn_idx],
                                      self.mnfld_n[i, nn_idx]))
            normals.append(self.mnfld_n[i, nn_idx])
            distances.append(signs*distances)

        distances = np.stack(distances).astype('f')
        normals = np.stack(normals).astype('f')
        return distances, normals

    def get_grid_divergence(self):
        # 2D implementation
        n_img = np.reshape(self.grid_n, [self.grid_res, self.grid_res, -1])
        frac_45 = 1./np.sqrt(2)
        filter = np.array([[[frac_45, -frac_45], [1., 0.], [frac_45, frac_45]],
                           [[0., -1.], [0., 0.], [0., 1.]],
                  [[-frac_45, -frac_45], [ -1., 0.], [-frac_45, frac_45]]])  # [y, x]
        padding = self.get_padding(n_img, filter, strides=[1, 1])
        n_img = torch.nn.functional.pad(torch.tensor(n_img, dtype=torch.float32), padding)
        div_img = torch.nn.functional.conv2d(n_img.permute([2, 0, 1]).unsqueeze(0),
                                             torch.tensor(filter, dtype=torch.float32).permute([2, 0, 1]).unsqueeze(0),
                                             ).squeeze().numpy()

        return div_img.flatten()


    def get_offgrid_divergnce(self, off_grid_points, method='nn'):
        #TODO implement interpulation method?
        if method == 'nn':
            # find the nearest grid point and return its divergence
            kdtree = spatial.cKDTree(self.grid_points)
            _, nn_idx = kdtree.query(off_grid_points, k=1)
        else:
            raise Warning('unsupported offgrid div computeation method')
        return self.grid_div[nn_idx]


    def get_padding(self, img, filter, strides=[1, 1]):
        # from https://discuss.pytorch.org/t/same-padding-equivalent-in-pytorch/85121/3
        in_height, in_width, _ = img.shape
        filter_height, filter_width, _ = filter.shape
        # The total padding applied along the height and width is computed as:
        if (in_height % strides[0] == 0):
            pad_along_height = max(filter_height - strides[0], 0)
        else:
            pad_along_height = max(filter_height - (in_height % strides[0]), 0)
        if (in_width % strides[1] == 0):
            pad_along_width = max(filter_width - strides[1], 0)
        else:
            pad_along_width = max(filter_width - (in_width % strides[1]), 0)

        # Finally, the padding on the top, bottom, left and right are:
        pad_top = pad_along_height // 2
        pad_bottom = pad_along_height - pad_top
        pad_left = pad_along_width // 2
        pad_right = pad_along_width - pad_left

        return (0, 0, pad_left, pad_right, pad_top, pad_bottom)


    def get_nonmnfld_points(self):
        if self.sample_type == 'grid':
            nonmnfld_points = self.grid_points
        elif self.sample_type == 'uniform':
            nonmnfld_points = np.random.uniform(-self.grid_range, self.grid_range,
                                                size=(self.grid_res * self.grid_res , 2)).astype(np.float32)
        elif self.sample_type == 'gaussian':
            nonmnfld_points = self.sample_gaussian_noise_around_shape()
            idx = np.random.choice(range(nonmnfld_points.shape[1]), self.grid_res * self.grid_res)
            sample_idx = np.random.choice(range(nonmnfld_points.shape[0]), self.grid_res * self.grid_res)
            nonmnfld_points = nonmnfld_points[sample_idx, idx]
        elif self.sample_type == 'combined':
            nonmnfld_points1 = self.sample_gaussian_noise_around_shape()
            nonmnfld_points2 = self.grid_points
            idx1 = np.random.choice(range(nonmnfld_points1.shape[1]), int(np.ceil(self.grid_res * self.grid_res / 2)))
            idx2 = np.random.choice(range(nonmnfld_points2.shape[0]), int(np.floor(self.grid_res * self.grid_res / 2)))
            sample_idx = np.random.choice(range(nonmnfld_points1.shape[0]), int(np.ceil(self.grid_res * self.grid_res / 2)))

            nonmnfld_points = np.concatenate([nonmnfld_points1[sample_idx, idx1], nonmnfld_points2[idx2]], axis=0)
        else:
            raise Warning("Unsupported non manfold sampling type {}".format(self.sample_type))
        return nonmnfld_points

    def sample_gaussian_noise_around_shape(self):
        n_noisy_points = int(np.round(self.grid_res * self.grid_res / self.n_points))
        noise = np.random.multivariate_normal([0, 0], [[self.sampling_std, 0], [0, self.sampling_std]],
                                              size=(self.n_samples, self.n_points, n_noisy_points)).astype(np.float32)
        nonmnfld_points = np.tile(self.points[:, :, None, :], [1, 1, n_noisy_points, 1]) + noise
        nonmnfld_points = nonmnfld_points.reshape([nonmnfld_points.shape[0], -1, nonmnfld_points.shape[-1]])
        return nonmnfld_points

    def generate_batch_indices(self):
        mnfld_idx = []
        nonmnfld_idx = []
        for i in range(self.n_samples):
            mnfld_idx.append(np.random.choice(self.point_idxs, self.n_points))
            nonmnfld_idx.append(np.random.choice(self.nonmnfld_points_idxs, self.n_points))
        self.mnfld_idx = np.array(mnfld_idx)
        self.nonmnfld_idx = np.array(nonmnfld_idx)

    def __getitem__(self, index):
        nonmnfld_idx = self.nonmnfld_idx[index]
        mnfld_idx = self.mnfld_idx[index]
        if self.nonmnfld_dist is not None:
            nonmnfld_dist = self.nonmnfld_dist[nonmnfld_idx]
        else:
            nonmnfld_dist = torch.tensor(0)

        return {'points' : self.points[index, mnfld_idx, :],  'mnfld_n': self.mnfld_n[index, mnfld_idx, :],  \
                'nonmnfld_dist': nonmnfld_dist, 'nonmnfld_n': self.nonmnfld_n[nonmnfld_idx],
                'nonmnfld_points': self.nonmnfld_points[nonmnfld_idx],
                }

    def __len__(self):
        return self.n_samples


class Circle(BasicShape2D):
    def __init__(self, *args, r=0.5):
        self.r = r
        BasicShape2D.__init__(self, *args)


    def get_mnfld_points(self):
        theta = np.random.uniform(0, 2*np.pi, size=(self.n_samples, self.n_points)).astype('f')
        x = self.r * np.sin(theta)
        y = self.r * np.cos(theta)

        points = np.stack([x, y], axis=2)
        return points

    def get_mnfld_n(self):
        return self.points / np.linalg.norm(self.points, axis=2, keepdims=True)


    def get_points_distances_and_normals(self, points):
        point_dist = np.linalg.norm(points, axis=1, keepdims=True)
        distances = point_dist - self.r
        normals = points / point_dist
        return distances, normals


class Polygon(BasicShape2D):
    def __init__(self, *args, vertices=[], line_sample_type='uniform'):

        # vertices: x,y points specifying the polygon
        self.vertices = np.array(vertices)
        self.lines = self.get_line_props()
        self.line_sample_type = line_sample_type
        BasicShape2D.__init__(self, *args)

    def get_mnfld_points(self):
        # sample points on the lines
        n_points_to_sample = self.n_points - len(self.vertices)
        if n_points_to_sample < 0:
            raise Warning("Fewer points to sample than polygon vertices. Please change the number of points")
        sample_prob = self.lines['line_length'] / np.sum(self.lines['line_length'])
        points_per_segment = np.floor(n_points_to_sample * sample_prob).astype(np.int32)
        points_leftover = int(n_points_to_sample - points_per_segment.sum())
        if not points_leftover == 0:
            for j in range(points_leftover):
                actual_prob = points_per_segment / points_per_segment.sum()
                prob_diff = sample_prob - actual_prob
                add_idx = np.argmax(prob_diff)
                points_per_segment[add_idx] = points_per_segment[add_idx] + 1

        points = []
        self.point_normal = []
        for point_idx, point in enumerate(self.vertices):
            l1_idx = len(self.vertices) - 1 if point_idx == 0 else point_idx - 1
            l2_idx = point_idx
            n = self.lines['nl'][l1_idx] + self.lines['nl'][l2_idx]
            self.point_normal.append(n / np.linalg.norm(n))
            points.append(point)
        points = np.repeat(np.array(points)[None, :], self.n_samples, axis=0)
        self.point_normal = np.repeat(np.array(self.point_normal)[None, :], self.n_samples, axis=0)

        for line_idx in range(len(self.lines['A'])):
            if self.line_sample_type == 'uniform':
                t = np.linspace(0, 1, points_per_segment[line_idx] + 1, endpoint=False)[1:]
                t = np.repeat(t[None, :], self.n_samples, axis=0)
            else:
                t = np.random.uniform(0, 1, [self.n_samples, points_per_segment[line_idx]])
            p1 = np.array(self.vertices[self.lines['start_idx'][line_idx]])
            p2 = np.array(self.vertices[self.lines['end_idx'][line_idx]])
            points = np.concatenate([points, p1 + t[:, :, None]*(p2 - p1)], axis=1)
            self.point_normal = np.concatenate([self.point_normal,
                                                np.tile(self.lines['nl'][line_idx][None, None, :],
                                                        [self.n_samples, points_per_segment[line_idx], 1])], axis=1)
        return points.astype('f')

    def get_mnfld_n(self):
        return self.point_normal


    def get_points_distances_and_normals(self, points):
        # iterate over all the lines and  finds the minimum distance between all points and line segments
        # good explenation ref : https://stackoverflow.com/questions/10983872/distance-from-a-point-to-a-polygon
        n_grid_points = len(points)
        p1x = np.vstack(self.vertices[self.lines['start_idx']][:, 0])
        p1y = np.vstack(self.vertices[self.lines['start_idx']][:, 1])
        p2x = np.vstack(self.vertices[self.lines['end_idx']][:, 0])
        p2y = np.vstack(self.vertices[self.lines['end_idx']][:, 1])
        p1p2 = np.array(self.lines['direction'])
        px = points[:, 0]
        py = points[:, 1]
        pp1 = np.vstack([px - np.tile(p1x, [1, 1, n_grid_points]), py - np.tile(p1y, [1, 1, n_grid_points])])
        pp2 = np.vstack([px - np.tile(p2x, [1, 1, n_grid_points]), py - np.tile(p2y, [1, 1, n_grid_points])])

        r = (p1p2[:, 0, None] * pp1[0, :, :] + p1p2[:, 1, None] * pp1[1, :, :]) / np.array(self.lines['line_length'])[:, None]

        d1 = np.linalg.norm(pp1, axis=0)
        d2 = np.linalg.norm(pp2, axis=0)
        dp = np.sqrt(np.square(d1) - np.square(r * np.array(self.lines['line_length'])[:, None]))
        d = np.where(r < 0, d1, np.where(r > 1, d2, dp))
        distances = np.min(d, axis=0)
        idx = np.argmin(d, axis=0)
        # compute normal vector
        polygon_path = Path(self.vertices)
        point_in_polygon = polygon_path.contains_points(points)
        point_sign = np.where(point_in_polygon, -1, 1)

        n = np.where(r < 0, pp1, np.where(r > 1, pp2, point_sign *
                                          np.tile(np.array(self.lines['nl']).transpose()[:, :, None],
                                          [1, 1, n_grid_points])))
        normals = np.take_along_axis(n, idx[None, None, :], axis=1).squeeze().transpose()
        normals = point_sign[:, None] * normals / np.linalg.norm(normals, axis=1, keepdims=True)
        distances = point_sign * distances

        return distances, normals


    def get_line_props(self):
        lines = {'A': [], 'B': [], 'C': [], 'nl': [], 'line_length': [], 'start_idx': [], 'end_idx': [], 'direction': []}
        for start_idx, start_point in enumerate(self.vertices):
            end_idx = 0 if start_idx == len(self.vertices)-1 else start_idx + 1
            end_point = self.vertices[end_idx]
            # Compute standard form coefficients

            A = start_point[1] - end_point[1]
            B = end_point[0] - start_point[0]
            C = - (A * start_point[0] + B * start_point[1])
            line_length = np.sqrt(np.square(A) + np.square(B))
            direction = [end_point[0] - start_point[0], end_point[1] - start_point[1]] / line_length
            nl = [A, B]
            nl = nl / np.linalg.norm(nl)
            line_props = {'A': A, 'B': B, 'C': C, 'nl': nl, 'line_length': line_length,
                          'start_idx': start_idx, 'end_idx': end_idx, 'direction': direction}
            for key in lines.keys():
                lines[key].append(line_props[key])

        return lines


def koch_line(start, end, factor):
    """
    Segments a line to Koch line, creating fractals.


    :param tuple start:  (x, y) coordinates of the starting point
    :param tuple end: (x, y) coordinates of the end point
    :param float factor: the multiple of sixty degrees to rotate
    :returns tuple: tuple of all points of segmentation
    """

    # coordinates of the start
    x1, y1 = start[0], start[1]

    # coordinates of the end
    x2, y2 = end[0], end[1]

    # the length of the line
    l = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)

    # first point: same as the start
    a = (x1, y1)

    # second point: one third in each direction from the first point
    b = (x1 + (x2 - x1) / 3., y1 + (y2 - y1) / 3.)

    # third point: rotation for multiple of 60 degrees
    c = (b[0] + l / 3. * np.cos(factor * np.pi / 3.), b[1] + l / 3. * np.sin(factor * np.pi / 3.))

    # fourth point: two thirds in each direction from the first point
    d = (x1 + 2. * (x2 - x1) / 3., y1 + 2. * (y2 - y1) / 3.)

    # the last point
    e = end

    return {'a': a, 'b': b, 'c': c, 'd': d, 'e': e, 'factor': factor}


def koch_snowflake(degree, s=1.0):
    """Generates all lines for a Koch Snowflake with a given degree.
    code from: https://github.com/IlievskiV/Amusive-Blogging-N-Coding/blob/master/Visualizations/snowflake.ipynb
    :param int degree: how deep to go in the branching process
    :param float s: the length of the initial equilateral triangle
    :returns list: list of all lines that form the snowflake
    """
    # all lines of the snowflake
    lines = []

    # we rotate in multiples of 60 degrees
    sixty_degrees = np.pi / 3.

    # vertices of the initial equilateral triangle
    A = (0., 0.)
    B = (s, 0.)
    C = (s * np.cos(sixty_degrees), s * np.sin(sixty_degrees))

    # set the initial lines
    if degree == 0:
        lines.append(koch_line(A, B, 0))
        lines.append(koch_line(B, C, 2))
        lines.append(koch_line(C, A, 4))
    else:
        lines.append(koch_line(A, B, 5))
        lines.append(koch_line(B, C, 1))
        lines.append(koch_line(C, A, 3))

    for i in range(1, degree):
        # every lines produce 4 more lines
        for _ in range(3 * 4 ** (i - 1)):
            line = lines.pop(0)
            factor = line['factor']

            lines.append(koch_line(line['a'], line['b'], factor % 6))  # a to b
            lines.append(koch_line(line['b'], line['c'], (factor - 1) % 6))  # b to c
            lines.append(koch_line(line['c'], line['d'], (factor + 1) % 6))  # d to c
            lines.append(koch_line(line['d'], line['e'], factor % 6))  # d to e

    return lines

def get_koch_points(degree, s=1.0):
    lines = koch_snowflake(degree, s=s)
    points = []
    for line in lines:
        for key in line.keys():
            if not key == 'factor' and not key == 'e':
                points.append(line[key])
    points = np.array(points) - np.array([s/2, (s/2)*np.tan(np.pi/6)])
    points = np.flipud(points) #reorder the points clockwise
    return points

def get2D_dataset(*args, shape_type='circle'):

    if shape_type == 'circle':
        out_shape = Circle(*args)
    elif shape_type == 'L':
        out_shape = Polygon(*args, vertices=[[0., 0.], [0.5, 0.], [0.5, -0.5],
                                 [-0.5, -0.5], [-0.5, 0.5], [0, 0.5]])
    elif shape_type == 'square':
        out_shape = Polygon(*args, vertices=[[-0.5, 0.5], [0.5, 0.5], [0.5, -0.5], [-0.5, -0.5]])
    elif shape_type == 'snowflake':
        vertices = get_koch_points(degree=2, s=1.0)
        out_shape = Polygon(*args, vertices=vertices)
    else:
        raise Warning("Unsupportaed shape")

    return out_shape


if __name__ == "__main__":
    np.random.seed(0)
    shape_type = 'L'
    res = 128 # has to be even
    example_idx = 0
    sample_type = 'grid'
    n_samples = 2
    n_points = 24
    dataset = get2D_dataset(n_points, n_samples, res, sample_type, 0.005, shape_type=shape_type)  # BasicShape2D(100, 20, res=50)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=False, num_workers=3, pin_memory=True)
    data = next(iter(dataloader))
    clean_points_gt = data['points'][example_idx].detach().cpu().numpy()
    n_gt = data['mnfld_n'][example_idx].detach().cpu().numpy()
    nonmnfld_points = data['nonmnfld_points'][example_idx].detach().cpu().numpy()
    grid_normals = dataset.grid_n

    vis.plot_sdf_indicator(dataset.vertices, dataset.grid_points[:, 0], dataset.grid_points[:, 1],
                           dataset.dist_img.flatten(), title_text='', show_ax=False, output_path='./vis/') # plot sdf, indicator function, and points

    # vis.plot_shape_data(dataset.grid_points[:, 0], dataset.grid_points[:, 1], dataset.dist_img.flatten(),
    #                     clean_points_gt, n_gt=n_gt, show_ax=True, show_bar=True,
    #                     title_text='', colorscale='Geyser', nonmnfld_points=nonmnfld_points, divergence=None,
    #                     grid_normals=grid_normals) # plot shape, sdf and other data

    # vis.plot_paper_teaser_images(dataset.grid_points[:, 0], dataset.grid_points[:, 1] ,dataset.dist_img.flatten(),
    #                              clean_points_gt, grid_normals) # plot images for DiGS paper teaser image
sanitychecks/basic_shape_dataset2d.py
sunyx523/StEik
# This file is partly based on DiGS: https://github.com/Chumbyte/DiGS
# This code initializes a 2D SIREN and generates images of SDF (image + 3D surface), eikonal, curl, divergence and distance to GT difference
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import basic_shape_dataset2d
import torch
import utils.visualizations as vis
import numpy as np
import models.Net as Net
import utils.utils as utils
from PIL import Image
import argparse

parser = argparse.ArgumentParser(description='Initialization visualization test')
parser.add_argument('--gpu_idx', type=int, default=0, help='set < 0 to use CPU')
parser.add_argument('--output_dir', type=str, default='./log/initialization_vis/', help='path to store output images')
# parser.add_argument('--plot_second_derivs', type=bool, default=False, help='toggle on to plot second derivative images')
parser.add_argument('--plot_second_derivs', type=bool, default=True, help='toggle on to plot second derivative images')
parser.add_argument('--nl', type=str, default='sine', help='specify nonlinearity sine | relu')
parser.add_argument('--init_type', type=str, default='mfgi', help='specify type of initialization mfgi | geometric_relu | geometric_sine | siren')
# parser.add_argument('--init_type', type=str, default='geometric_sine', help='specify type of initialization mfgi | geometric_relu | geometric_sine | siren')
# parser.add_argument('--init_type', type=str, default='siren', help='specify type of initialization mfgi | geometric_relu | geometric_sine | siren')
parser.add_argument('--decoder_n_hidden_layers', type=int, default=4, help='number of MLP layers')
parser.add_argument('--decoder_hidden_dim', type=int, default=128, help='number units in the decoder')
parser.add_argument('--grid_res', type=int, default=256, help='evaluation grid resolution')
args = parser.parse_args()

seed = 0
# get data loaders
torch.manual_seed(seed)
np.random.seed(seed)
# [n_points, n_samples, res, sample_type, sapmling_std], shape_type
test_set = basic_shape_dataset2d.get2D_dataset(128, 1, args.grid_res, 'grid', 'circle')
test_dataloader = torch.utils.data.DataLoader(test_set, batch_size=1, shuffle=False, num_workers=2,
                                              pin_memory=True)
# get model
device = torch.device("cuda:" + str(args.gpu_idx) if (torch.cuda.is_available()) else "cpu")

SINR = Net.Network(latent_size=0, in_dim=2, decoder_hidden_dim=args.decoder_hidden_dim,
                              nl=args.nl, encoder_type='none',
                              decoder_n_hidden_layers=args.decoder_n_hidden_layers, init_type=args.init_type)

os.makedirs(args.output_dir, exist_ok=True)

_, test_data = next(enumerate(test_dataloader))
SINR.eval()
nonmnfld_points = test_data['nonmnfld_points'].to(device)
SINR.to(device)

print("Visualizing initialization output {}".format(args.init_type))
# initialize and evaluate the model
nonmnfld_points.requires_grad_()
output_pred = SINR(nonmnfld_points, mnfld_pnts=None)

x_grid, y_grid, z_grid, z_diff, eikonal_term, grid_div, grid_curl = \
    utils.compute_deriv_props(SINR.decoder, output_pred["latent"], z_gt=test_set.dist_img, device=device)

contour_img, curl_img, eikonal_img, div_image, z_diff_img =\
    vis.plot_init_contour_div_props(x_grid, y_grid, z_grid, clean_points=None,
                            z_diff=z_diff, eikonal_term=eikonal_term, grid_div=grid_div, grid_curl=grid_curl,
                            example_idx=0, n_gt=None, n_pred=None,
                            nonmnfld_points=None, 
                            title_text="",
                            # title_text='initialization:  ' + str(args.init_type),
                            plot_second_derivs=args.plot_second_derivs)

# print(z_grid.min(), z_grid.max())

# Output visualizations
fig = vis.plot_sdf_surface(x_grid, y_grid, z_grid, show=False, show_ax=True, title_txt='sdf initialization surface')
fig.write_html(os.path.join(args.output_dir, "sdf_surface" + str(args.decoder_hidden_dim) + '_' + args.init_type + ".html"))
im = Image.fromarray(contour_img)
im.save(os.path.join(args.output_dir, "sdf_" + str(args.decoder_hidden_dim) + '_' + args.init_type + "_init.png"))
im = Image.fromarray(eikonal_img)
im.save(os.path.join(args.output_dir, "eikonal_" + str(args.decoder_hidden_dim) + '_' + args.init_type + "_init.png"))
im = Image.fromarray(z_diff_img)
im.save(os.path.join(args.output_dir, "zdiff_" + str(args.decoder_hidden_dim) + '_' + args.init_type + "_init.png"))
if args.plot_second_derivs:
    # im = Image.fromarray(curl_img)
    # im.save(os.path.join(args.output_dir, "curl_" + str(args.decoder_hidden_dim) + '_' + args.init_type + "_init.png"))
    im = Image.fromarray(div_image)
    im.save(os.path.join(args.output_dir, "div_" + str(args.decoder_hidden_dim) + '_' + args.init_type + "_init.png"))






sanitychecks/initialization_visualization.py
sunyx523/StEik
# This file is partly based on DiGS: https://github.com/Chumbyte/DiGS
import argparse
import torch
import os
import numpy as np

def add_args(parser):

    parser.add_argument('--gpu_idx', type=int, default=1, help='set < 0 to use CPU')
    parser.add_argument('--logdir', type=str, default='./log/siren_square', help='log directory')
    parser.add_argument('--seed', type=int, default=3627473, help='random seed')
    parser.add_argument('--shape_type', type=str, default='L', help='shape dataset to load circle | square | L ')

    parser.add_argument('--n_samples', type=int, default=10000,
                        help='number of samples in the generated train and test set')
    parser.add_argument('--n_point_total', type=int, default=30000,
                        help='the total number of points sampled on the manifold')
    parser.add_argument('--refine_epoch', type=int, default=0, help='refine model from this epoch, '
                                                                    '0 for training from scratch')
    parser.add_argument('--model_dirpath', type=str,
                        default='/mnt/3.5TB_WD/PycharmProjects/ZeroCurlShapes/models',
                        help='path to model directory for backup')
    parser.add_argument('--parallel', type=int, default=False, help='use data parallel')

    # training paramaeters
    parser.add_argument('--num_epochs', type=int, default=1, help='number of training epochs')
    parser.add_argument('--lr', type=float, default=1e-4, 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('--n_points', type=int, default=15000, help='number of points in each point cloud')
    # parser.add_argument('--delta', type=float, default=0.25, help='clamping value')

    # Network architecture and loss
    parser.add_argument('--decoder_hidden_dim', type=int, default=128, help='length of decoder hidden dim')
    parser.add_argument('--encoder_hidden_dim', type=int, default=128, help='length of encoder hidden dim')
    parser.add_argument('--decoder_n_hidden_layers', type=int, default=4, help='number of layers in the decoder')
    parser.add_argument('--nl', type=str, default='sine', help='type of non linearity sine | relu')
    parser.add_argument('--latent_size', type=int, default=0, help='number of elements in the latent vector')
    parser.add_argument('--sphere_init_params', nargs='+', type=float, default=[1.6, 1.0],
                    help='radius and scaling')
    parser.add_argument('--neuron_type', type=str, default='quadratic', help='type of neuron')

    parser.add_argument('--normalize_normal_loss', type=int, default=False, help='normal loss normalization flag')
    parser.add_argument('--unsigned_n', type=int, default=True, help='flag for unsigned normal loss')
    parser.add_argument('--unsigned_d', type=int, default=False, help='flag for unsigned distance loss')

    parser.add_argument('--encoder_type', type=str, default='none', help='type of encoder None | pointnet')
    parser.add_argument('--export_vis', type=int, default=False, help='export levelset visualization while training')

    parser.add_argument('--loss_type', type=str, default='siren_wo_n_w_div', help='loss type to use: siren | siren_wo_n | siren_wo_n_w_div | siren_wo_n_w_div_w_reg')
    parser.add_argument('--n_loss_type', type=str, default='cos', help='type of normal los cos | ')
    parser.add_argument('--inter_loss_type', type=str, default='exp', help='type of inter los exp | unsigned_diff | signed_diff')
    parser.add_argument('--sampler_prob', type=str, default='none',
                        help='type of sampler probability for non manifold points on the grid none | div | curl')

    parser.add_argument('--div_type', type=str, default='dir_l1', help='divergence term norm dir_l1 | dir_l2 | full_l1 | full_l2')
    parser.add_argument('--grid_res', type=int, default=256,
                        help='uniform grid resolution')
    parser.add_argument('--div_clamp', type=float, default=50, help='divergence clamping value')
    parser.add_argument('--nonmnfld_sample_type', type=str, default='grid',
                        help='how to sample points off the manifold - grid | gaussian | combined')
    parser.add_argument('--div_decay', type=str, default='none', help='divergence term importance decay none | step | linear')
    parser.add_argument('--div_decay_params', nargs='+', type=float, default=[0.0, 0.5, 0.75],
                        help='epoch number to evaluate')
    parser.add_argument('--init_type', type=str, default='mfgi',
                        help='initialization type siren | geometric_sine | geometric_relu |pt_circle | mfgi')
    parser.add_argument('--loss_weights', nargs='+', type=float, default=[3e3, 1e2, 1e2, 5e1, 1e1],
                        help='loss terms weights sdf | inter | normal | eikonal | div')
    return parser


def get_parser():
    # command line args
    parser = argparse.ArgumentParser(description='Local implicit functions experiment')
    parser = add_args(parser)
    return parser


def get_args():
    parser = get_parser()
    args = parser.parse_args()
    return args

def get_test_args():
    parser = argparse.ArgumentParser(description='Local implicit functions test experiment')
    parser.add_argument('--gpu_idx', type=int, default=0, help='set < 0 to use CPU')
    parser.add_argument('--logdir', type=str, default='./log/siren_square/', help='log directory')
    parser.add_argument('--n_points', type=int, default=0, help='number of points in each point cloud, '
                                                                  'if 0 use training options')
    parser.add_argument('--batch_size', type=int, default=0, help='number of samples in a minibatch, if 0 use training')
    parser.add_argument('--grid_res', type=int, default=256, help='evaluation grid resolution')
    parser.add_argument('--epoch_n', type=int, nargs='+', default=np.arange(0, 10000, 100).tolist(), help='epoch number to evaluate')
    # parser.add_argument('--epoch_n', type=int, default=29900, help='epoch to evaluate')
    parser.add_argument('--std', type=float, default=0.12, help='noise standard deviation')
    parser.add_argument('--noise_type', type=str, default='gaussian', help='type of noise to inject the data')
    parser.add_argument('--normal_vis', type=int, default=False, help='flag to visualize normals in plot')
    parser.add_argument('--n_example_plots', type=int, default=1, help='number of examples to plot')
    test_opt = parser.parse_args()

    param_filename = os.path.join(test_opt.logdir, 'trained_models/', 'params.pth')
    train_opt = torch.load(param_filename)

    test_opt.nl, test_opt.latent_size, test_opt.encoder_type, test_opt.n_samples, test_opt.seed, \
        test_opt.decoder_hidden_dim, test_opt.encoder_hidden_dim, test_opt.n_loss_type,\
        test_opt.normalize_normal_loss, test_opt.unsigned_n, test_opt.unsigned_d, test_opt.loss_type, \
        test_opt.model_dirpath, test_opt.inter_loss_type, test_opt.shape_type, \
        test_opt.div_decay,  test_opt.div_type, test_opt.div_clamp, test_opt.decoder_n_hidden_layers, \
         = train_opt.nl, train_opt.latent_size, train_opt.encoder_type,  train_opt.n_samples,  train_opt.seed,\
          train_opt.decoder_hidden_dim, train_opt.encoder_hidden_dim, train_opt.n_loss_type,  \
          train_opt.normalize_normal_loss, train_opt.unsigned_n, train_opt.unsigned_d, train_opt.loss_type, \
          train_opt.model_dirpath, train_opt.inter_loss_type, train_opt.shape_type,\
          train_opt.div_decay, train_opt.div_type, train_opt.div_clamp, train_opt.decoder_n_hidden_layers

    test_opt.n_point_total = train_opt.n_points

    if test_opt.n_points == 0:
        test_opt.n_points = train_opt.n_points
    if test_opt.batch_size == 0:
        test_opt.batch_size = train_opt.batch_size
    if "parallel" in train_opt:
        test_opt.parallel = train_opt.parallel
    else:
        test_opt.parallel = False
    if 'init_type' in train_opt:
        test_opt.init_type = train_opt.init_type
    else:
        test_opt.init_type = 'siren'
    test_opt.nonmnfld_sample_type = 'grid'
    return test_opt

sanitychecks/sc_args.py
sunyx523/StEik
# This script converts the output test images into mp4 videos
import cv2
import os
image_path = '/mnt/IronWolf/logs/DiBS/2D/grid_sampling_256/snowflake/siren_origin/vis/' # path to image output directory
output_path = os.path.join(image_path, 'vids')
os.makedirs(output_path, exist_ok=True)
image_siz = 800
fps = 15
fourcc_encoder = cv2.VideoWriter_fourcc(*'mp4v')
extension = '.mp4'
out_sdf_eikonal_div = cv2.VideoWriter(os.path.join(output_path, 'sdf_eikonal_div' + extension), fourcc_encoder, fps, (image_siz*3, image_siz))
out_sdf_eikonal_div_zdiff = cv2.VideoWriter(os.path.join(output_path, 'sdf_eikonal_div_zdiff' + extension), fourcc_encoder, fps, (image_siz*4, image_siz))
out_sdf = cv2.VideoWriter(os.path.join(output_path, 'sdf' + extension), fourcc_encoder, fps, (image_siz, image_siz))
out_eikonal= cv2.VideoWriter(os.path.join(output_path, 'eikonal' + extension), fourcc_encoder, fps, (image_siz, image_siz))
out_div = cv2.VideoWriter(os.path.join(output_path, 'div' + extension), fourcc_encoder, fps, (image_siz, image_siz))
out_curl = cv2.VideoWriter(os.path.join(output_path, 'curl' + extension), fourcc_encoder, fps, (image_siz, image_siz))
out_zdiff = cv2.VideoWriter(os.path.join(output_path, 'zdiff' + extension), fourcc_encoder, fps, (image_siz, image_siz))

sdf_file_list = [file for file in os.listdir(image_path) if file.endswith('.png') and 'sdf' in file]
eikonal_file_list = [file for file in os.listdir(image_path) if file.endswith('.png') and 'eikonal' in file]
div_file_list = [file for file in os.listdir(image_path) if file.endswith('.png') and 'div' in file]
curl_file_list = [file for file in os.listdir(image_path) if file.endswith('.png') and 'curl' in file]
zdiff_file_list = [file for file in os.listdir(image_path) if file.endswith('.png') and 'zdiff' in file]
sdf_file_list.sort(); eikonal_file_list.sort(); div_file_list.sort(); curl_file_list.sort(); zdiff_file_list.sort()

n_files = sum('sdf' in s for s in sdf_file_list)
for i in range(n_files):
    sdf_img = cv2.imread(os.path.join(image_path, sdf_file_list[i]))
    eikonal_img = cv2.imread(os.path.join(image_path, eikonal_file_list[i]))
    div_img = cv2.imread(os.path.join(image_path, div_file_list[i]))
    curl_img = cv2.imread(os.path.join(image_path, curl_file_list[i]))
    zdiff_img = cv2.imread(os.path.join(image_path, zdiff_file_list[i]))
    out_sdf.write(sdf_img); out_eikonal.write(eikonal_img); out_div.write(div_img); out_curl.write(curl_img); out_zdiff.write(zdiff_img);
    out_sdf_eikonal_div.write(cv2.hconcat([sdf_img, eikonal_img, div_img]))
    out_sdf_eikonal_div_zdiff.write(cv2.hconcat([sdf_img, eikonal_img, div_img, zdiff_img]))

out_sdf_eikonal_div.release(); out_sdf.release(); out_eikonal.release(); out_div.release(); out_curl.release(); out_zdiff.release(); out_sdf_eikonal_div_zdiff.release();
sanitychecks/scripts/convert_images_to_video.py
sunyx523/StEik
# This file is partly based on DiGS: https://github.com/Chumbyte/DiGS
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import basic_shape_dataset2d
import torch
import utils.visualizations as vis
import numpy as np
import models.Net as Net
from models.losses import Loss
import torch.nn.parallel
import utils.utils as utils
import sc_args

from PIL import Image

# get training parameters
args = sc_args.get_test_args()
gpu_idx, nl, n_points, batch_size, n_samples, latent_size, logdir, \
n_loss_type, normalize_normal_loss, unsigned_n, unsigned_d, loss_type, seed, encoder_type,\
    model_dirpath, inter_loss_type =\
    args.gpu_idx, args.nl, args.n_points, args.batch_size, args.n_samples, args.latent_size, \
    args.logdir, args.n_loss_type, \
    args.normalize_normal_loss, args.unsigned_n, args.unsigned_d, args.loss_type, args.seed, args.encoder_type, \
    args.model_dirpath, args.inter_loss_type

n_samples = 1
n_points = 128
args.n_point_total = 1024
plot_second_derivs = True
# plot_second_derivs = False

# get data loaders
torch.manual_seed(seed)
np.random.seed(seed)
test_set = basic_shape_dataset2d.get2D_dataset(n_points, n_samples, args.grid_res,
                                               args.nonmnfld_sample_type, shape_type=args.shape_type)
test_dataloader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=2,
                                              pin_memory=True)
# get model
device = torch.device("cuda:" + str(gpu_idx) if (torch.cuda.is_available()) else "cpu")

SINR = Net.Network(latent_size=latent_size, in_dim=2, decoder_hidden_dim=args.decoder_hidden_dim,
                              nl=args.nl, encoder_type='none',
                              decoder_n_hidden_layers=args.decoder_n_hidden_layers, init_type=args.init_type)
if args.parallel:
    if (device.type == 'cuda'):
        SINR = torch.nn.DataParallel(SINR)


model_dir = os.path.join(logdir, 'trained_models')
output_dir = os.path.join(logdir, 'vis')
os.makedirs(output_dir, exist_ok=True)
# get loss
criterion = Loss(weights=[3e3, 1e2, 1e2, 5e1, 1e1], #sdf, intern, normal, eikonal, div
                               loss_type=loss_type, div_decay=args.div_decay, div_type=args.div_type)

_, test_data = next(enumerate(test_dataloader))
SINR.eval()
mnfld_points, normals_gt, nonmnfld_dist_gt, nonmnfld_points, nonmnfld_n_gt = \
    test_data['points'].to(device), test_data['mnfld_n'].to(device), \
    test_data['nonmnfld_dist'].to(device), test_data['nonmnfld_points'].to(device), \
    test_data['nonmnfld_n'].to(device)

grid_points = test_set.grid_points

for epoch in args.epoch_n:
    model_filename = os.path.join(model_dir, 'model_%d.pth' % (epoch))
    SINR.load_state_dict(torch.load(model_filename, map_location=device))
    SINR.to(device)

    print("Converting implicit to level set for shape {} epoch {}".format(args.shape_type, epoch))

    mnfld_points.requires_grad_()
    nonmnfld_points.requires_grad_()
    output_pred = SINR(nonmnfld_points, mnfld_points)
    loss_dict_test, n_pred = criterion(output_pred=output_pred, mnfld_points=mnfld_points, 
                                    nonmnfld_points=nonmnfld_points, mnfld_n_gt=normals_gt)

    x_grid, y_grid, z_grid, z_diff, eikonal_term, grid_div, grid_curl = \
        utils.compute_deriv_props(SINR.decoder, output_pred["latent"], z_gt=test_set.dist_img, device=device)

    contour_img, curl_img, eikonal_img, div_image, z_diff_img =\
        vis.plot_contour_div_props(x_grid, y_grid, z_grid, mnfld_points,
                                   z_diff, eikonal_term, grid_div, grid_curl, example_idx=0,
                                   n_gt=normals_gt, n_pred=n_pred,
                                   nonmnfld_points=None, title_text='Epoch ' + str(epoch), plot_second_derivs=plot_second_derivs)
    # save the generated images
    im = Image.fromarray(contour_img)
    im.save(os.path.join(output_dir, "sdf_" + str(epoch).zfill(6) + ".png"))
    im = Image.fromarray(eikonal_img)
    im.save(os.path.join(output_dir, "eikonal_" + str(epoch).zfill(6) + ".png"))
    im = Image.fromarray(z_diff_img)
    im.save(os.path.join(output_dir, "zdiff_" + str(epoch).zfill(6) + ".png"))    
    if plot_second_derivs:
        # im = Image.fromarray(curl_img)
        # im.save(os.path.join(output_dir, "curl_" + str(epoch).zfill(6) + ".png"))
        im = Image.fromarray(div_image)
        im.save(os.path.join(output_dir, "div_" + str(epoch).zfill(6) + ".png"))






sanitychecks/test_basic_shape.py
sunyx523/StEik
# This file is partly based on DiGS: https://github.com/Chumbyte/DiGS
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import basic_shape_dataset2d
import torch
import utils.visualizations as vis
import numpy as np
import models.Net as Net
from models.losses import Loss
import torch.nn.parallel
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
import utils.utils as utils
import sc_args

# get training parameters
args = sc_args.get_args()
gpu_idx, nl, n_points, batch_size, n_samples, latent_size, lr, num_epochs, logdir, \
n_loss_type, normalize_normal_loss, unsigned_n, unsigned_d, loss_type, seed, encoder_type,\
    model_dirpath, inter_loss_type =\
    args.gpu_idx, args.nl, args.n_points, args.batch_size, args.n_samples, args.latent_size, \
    args.lr, args.num_epochs, args.logdir, args.n_loss_type, \
    args.normalize_normal_loss, args.unsigned_n, args.unsigned_d, args.loss_type, args.seed, args.encoder_type, \
    args.model_dirpath, args.inter_loss_type

# set up backup and 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' % ('../models/Net.py', logdir))  # backup the models files
os.system('cp %s %s' % ('../models/losses.py', logdir))  # backup the losses files

# get data loaders
torch.manual_seed(0)  #change random seed for training set (so it will be different from test set
np.random.seed(0)
train_set = basic_shape_dataset2d.get2D_dataset(n_points, n_samples, args.grid_res, args.nonmnfld_sample_type, shape_type=args.shape_type)
train_dataloader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2,
                                               pin_memory=True)
torch.manual_seed(seed)
np.random.seed(seed)

# get model
device = torch.device("cuda:" + str(gpu_idx) if (torch.cuda.is_available()) else "cpu")

SINR = Net.Network(latent_size=latent_size, in_dim=2, decoder_hidden_dim=args.decoder_hidden_dim,
                              nl=args.nl, encoder_type='none', neuron_type=args.neuron_type,
                              decoder_n_hidden_layers=args.decoder_n_hidden_layers, init_type=args.init_type)
if args.parallel:
    if (device.type == 'cuda'):
        SINR = torch.nn.DataParallel(SINR)
n_parameters = utils.count_parameters(SINR)
utils.log_string("Number of parameters in the current model:{}".format(n_parameters), log_file)
# Setup Adam optimizers
optimizer = optim.Adam(SINR.parameters(), lr=lr, betas=(0.9, 0.999))
scheduler = lr_scheduler.ExponentialLR(optimizer, gamma=1.0) # Does nothing
# scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=np.arange(2000, args.n_samples*num_epochs, 2000), gamma=0.5)  # milestones in number of optimizer iterations

if not args.refine_epoch == 0:
    refine_model_filename = os.path.join(model_outdir,
                                         'model_%d.pth' % (args.refine_epoch))
    SINR.load_state_dict(torch.load(refine_model_filename, map_location=device))
    optimizer.step()

SINR.to(device)

# get loss
criterion = Loss(weights=[3e3, 1e2, 1e2, 5e1, 1e2],
                               loss_type=loss_type, div_decay=args.div_decay, div_type=args.div_type)

num_batches = len(train_dataloader)
refine_flag = True
grid_points = train_set.grid_points

# Train the shape implicit neural representation
train_time_running_mean = 0
train_time_running_sum = 0
for epoch in range(num_epochs):
    if epoch <= args.refine_epoch and refine_flag and not args.refine_epoch == 0:
        scheduler.step()
        continue
    else:
        refine_flag = False

    for batch_idx, data in enumerate(train_dataloader):

        # save model before update
        if batch_idx % 50 == 0 :
            utils.log_string("saving model to file :{}".format('model_%d.pth' % (batch_idx)),
                             log_file)
            torch.save(SINR.state_dict(),
                       os.path.join(model_outdir, 'model_%d.pth' % (batch_idx)))
        # train model
        SINR.zero_grad()
        # with torch.cuda.amp.autocast():
        SINR.train()

        mnfld_points, normals_gt, nonmnfld_dist_gt, nonmnfld_points, nonmnfld_n_gt  =  \
            data['points'].to(device), data['mnfld_n'].to(device),  data['nonmnfld_dist'].to(device), \
            data['nonmnfld_points'].to(device), data['nonmnfld_n'].to(device),

        mnfld_points.requires_grad_()
        nonmnfld_points.requires_grad_()

        output_pred = SINR(nonmnfld_points, mnfld_points)

        loss_dict, _ = criterion(output_pred=output_pred, mnfld_points=mnfld_points, 
                                    nonmnfld_points=nonmnfld_points, mnfld_n_gt=normals_gt)

        lr = torch.tensor(optimizer.param_groups[0]['lr'])
        loss_dict["lr"] = lr
        utils.log_losses(log_writer_train, batch_idx, num_batches, loss_dict)
        # utils.log_weight_hist(log_writer_train, epoch, batch_idx, num_batches, SINR.decoder.fc_block.net[:], batch_size)
        if batch_idx % 25 == 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}'.format(
                epoch, batch_idx * 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()), log_file)
            utils.log_string('Epoch: {} [{:4d}/{} ({:.0f}%)] Unweighted L_s : L_Mnfld: {:.5f},  '
                    'L_NonMnfld: {:.5f},  L_Nrml: {:.5f},  L_Eknl: {:.5f},  L_Div: {:.5f}'.format(
                epoch, batch_idx * 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["div_loss"].item()), log_file)
            utils.log_string('', log_file)


        loss_dict["loss"].backward()
        optimizer.step()
        
        scheduler.step()
        criterion.update_div_weight(epoch*n_samples + batch_idx, num_epochs*n_samples, args.div_decay_params)

        # save last model
        if batch_idx == num_batches - 1 :
            utils.log_string("saving model to file :{}".format('model_%d.pth' % (num_batches)),
                             log_file)
            torch.save(SINR.state_dict(),
                       os.path.join(model_outdir, 'model_%d.pth' % (num_batches)))
sanitychecks/train_basic_shape.py
sunyx523/StEik
# This file is borrowed from DiGS: https://github.com/Chumbyte/DiGS
import numpy as np
import json
import os
import trimesh
from scipy.spatial import cKDTree as KDTree
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
import surface_recon_args
# import utils.utils as utils

import torch
import recon_dataset as dataset
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import models.Net as Net
device = torch.device("cuda")
args = surface_recon_args.get_train_args()
Net = Net.Network(latent_size=0, in_dim=3, decoder_hidden_dim=args.decoder_hidden_dim, nl=args.nl, encoder_type='none',
                   decoder_n_hidden_layers=args.decoder_n_hidden_layers, neuron_type=args.neuron_type, init_type='mfgi')


dataset_path = args.dataset_path
raw_dataset_path = args.raw_dataset_path
mesh_path = args.logdir

print(dataset_path)
print(raw_dataset_path)
print(mesh_path)

# Print metrics to file in logdir as well
out_path = os.path.join(args.logdir, 'metric_summary.txt')
import builtins as __builtin__
def print(*args, **kwargs):
    # Override print function to also print to file
    __builtin__.print(*args, **kwargs)
    with open(out_path, 'a') as fp:
        __builtin__.print(*args, file=fp, **kwargs)

print("Metrics on ShapeNet")

shape_class_name_dict = \
        {
            "04256520": "sofa",
            "02691156": "airplane",
            "03636649": "lamp",
            "04401088": "telephone",
            "04530566": "watercraft",
            "03691459": "loudspeaker",
            "03001627": "chair",
            "02933112": "cabinet",
            "04379243": "table",
            "03211117": "display",
            "02958343": "car",
            "02828884": "bench",
            "04090263": "rifle",
        }

shape_class_name2id = {v:k for k,v in shape_class_name_dict.items()}

eval_type = "DeepSDF"
# eval_type = "Default"

def compute_dists(recon_points, gt_points, eval_type="Default"):
    recon_kd_tree = KDTree(recon_points)
    gt_kd_tree = KDTree(gt_points)
    re2gt_distances, re2gt_vertex_ids = recon_kd_tree.query(gt_points, workers=4)
    gt2re_distances, gt2re_vertex_ids = gt_kd_tree.query(recon_points, workers=4)
    if eval_type == 'DeepSDF':
        cd_re2gt = np.mean(re2gt_distances**2)
        cd_gt2re = np.mean(gt2re_distances**2)
        hd_re2gt = np.max(re2gt_distances)
        hd_gt2re = np.max(gt2re_distances)
        chamfer_dist = cd_re2gt + cd_gt2re
        hausdorff_distance = np.max((hd_re2gt, hd_gt2re))
    else:
        cd_re2gt = np.mean(re2gt_distances)
        cd_gt2re = np.mean(gt2re_distances)
        hd_re2gt = np.max(re2gt_distances)
        hd_gt2re = np.max(gt2re_distances)
        chamfer_dist = 0.5* (cd_re2gt + cd_gt2re)
        hausdorff_distance = np.max((hd_re2gt, hd_gt2re))
    return chamfer_dist, hausdorff_distance, cd_re2gt, cd_gt2re, hd_re2gt, hd_gt2re

order = ["car","chair","airplane","display","table","rifle","cabinet","loudspeaker","telephone","bench","sofa","watercraft","lamp",]
order = sorted(order)

chamfers = {}
IoUs = {}
# for shape_class in os.listdir(mesh_path):
for shape_class in order:
    if shape_class not in os.listdir(mesh_path):
        continue
    shape_class_id = shape_class_name2id[shape_class]
    gt_shape_class_path = os.path.join(dataset_path, shape_class)
    gt_raw_shape_class_path = os.path.join(raw_dataset_path, shape_class_id)

    result_meshes_path = os.path.join(mesh_path, shape_class, 'result_meshes')
    saved_weights_class_path = os.path.join(mesh_path, shape_class)
    shape_files = [f for f in os.listdir(result_meshes_path) if '_iter_' not in f]
    print("Found {} files for {}".format(len(shape_files), shape_class))

    chamfers[shape_class] = []
    IoUs[shape_class] = []
    for shape_file in shape_files:
        shape = shape_file.replace(".ply","")
        recon_shape_path = os.path.join(result_meshes_path, shape_file)
        recon_mesh = trimesh.load(recon_shape_path)
        
        gt_shape_path = os.path.join(gt_shape_class_path, shape_file)
        gt_shape_weights_path = os.path.join(saved_weights_class_path, shape_file.replace('.ply',''), 'trained_models', 'model.pth')
        gt_pc = trimesh.load(gt_shape_path)
        gt_raw_shape_path = os.path.join(gt_raw_shape_class_path, shape)
        points = np.load(os.path.join(gt_raw_shape_path, 'points.npz')) # [('points', (100000, 3)), ('occupancies', (12500,)), ('loc', (3,)), ('scale', ())]
        pointcloud = np.load(os.path.join(gt_raw_shape_path, 'pointcloud.npz')) # [('points', (100000, 3)), ('normals', (100000, 3)), ('loc', (3,)), ('scale', ())]
        gen_points = points['points'] # (100000,3)
        occupancies = np.unpackbits(points['occupancies']) # (100000)

        n_points=15000; n_samples=10000; grid_res=512
        test_set = dataset.ReconDataset(gt_shape_path, n_points*n_samples, n_samples=1, res=grid_res, sample_type='grid',
                                requires_dist=False)
        cp, scale, bbox = test_set.cp, test_set.scale, test_set.bbox
        Net.load_state_dict(torch.load(gt_shape_weights_path, map_location=device))
        Net.to(device)

        eval_points = (gen_points - cp) / scale
        eval_points = torch.tensor(eval_points, device=device, dtype=torch.float32)
        res = Net.decoder(eval_points)

        pred_occupancies = (res.reshape(-1)<0).int().detach().cpu().numpy()
        iou = (occupancies & pred_occupancies).sum() / (occupancies | pred_occupancies).sum()
        IoUs[shape_class].append(iou)
        print("IoU, {:.4f}".format(iou))


        gt_points = gt_pc.vertices
        recon_points = trimesh.sample.sample_surface(recon_mesh, 30000)[0]
        chamfer_dist, hausdorff_distance,  cd_re2gt, cd_gt2re, hd_re2gt, hd_gt2re = compute_dists(recon_points, gt_points, eval_type="DeepSDF")
        chamfers[shape_class].append(chamfer_dist)
        print('\t{: <10} {:e} {:e} {:e} {:e}'.format(shape, chamfer_dist, hausdorff_distance, cd_re2gt, hd_re2gt))
    dists_np = np.array(chamfers[shape_class])
    print("{}: Mean: {:e}, Median: {:e}, Std: {:e}".format(shape_class,  dists_np.mean(), np.median(dists_np), dists_np.std()))
print('Final Chamfers')


# for key in chamfers:
for key in order:
    if key not in chamfers:
        continue
    dists_np = np.array(chamfers[key])
    print("{:<15} : Mean: {:e}, Median: {:e}, Std: {:e}".format(key, dists_np.mean(), np.median(dists_np), dists_np.std()))
all_chamfers = np.concatenate([chamfers[k] for k in chamfers]).reshape(-1)
print("Overall: Mean: {:e}, Median: {:e}, Std: {:e}".format(all_chamfers.mean(), np.median(all_chamfers), all_chamfers.std()))

print('Final IoUs')
for key in order:
    if key not in IoUs:
        continue
    IoUs_np = np.array(IoUs[key])
    print("{:<15} : Mean: {:.4f}, Median: {:.4f}, Std: {:.4f}".format(key, IoUs_np.mean(), np.median(IoUs_np), IoUs_np.std()))
all_IoUs = np.concatenate([IoUs[k] for k in IoUs]).reshape(-1)
print("Overall: Mean: {:.4f}, Median: {:.4f}, Std: {:.4f}".format(all_IoUs.mean(), np.median(all_IoUs), all_IoUs.std()))

surface_reconstruction/compute_metrics_shapenet.py
sunyx523/StEik
# This file is borrowed from DiGS: https://github.com/Chumbyte/DiGS
import numpy as np
import json
import os
import trimesh
from scipy.spatial import cKDTree as KDTree
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
import surface_recon_args

args = surface_recon_args.get_train_args()

scan_path = os.path.join(args.dataset_path, "scans")
gt_path = os.path.join(args.dataset_path, "ground_truth")
mesh_path = args.results_path

# Or set manually
# scan_path = '/home/chamin/DiGSGithub/DiBS/data/deep_geometric_prior_data/scans'
# gt_path = '/home/chamin/DiGSGithub/DiBS/data/deep_geometric_prior_data/ground_truth'
# mesh_path = '/home/chamin/DiGSGithub/DiBS/surface_reconstruction/log/surface_reconstruction2/DiGS_surf_recon_experiment/result_meshes'

# Print metrics to file in logdir as well
out_path = os.path.join(args.logdir, 'metric_summary.txt')
import builtins as __builtin__
def print(*args, **kwargs):
    # Override print function to also print to file
    __builtin__.print(*args, **kwargs)
    with open(out_path, 'a') as fp:
        __builtin__.print(*args, file=fp, **kwargs)

print("Metrics on SRB (DC, DH, DC-> DH->)")


# eval_type = "DeepSDF"
eval_type = "Default"

def compute_dists(recon_points, gt_points, eval_type="Default"):
    recon_kd_tree = KDTree(recon_points)
    gt_kd_tree = KDTree(gt_points)
    re2gt_distances, re2gt_vertex_ids = recon_kd_tree.query(gt_points, workers=4)
    gt2re_distances, gt2re_vertex_ids = gt_kd_tree.query(recon_points, workers=4)
    if eval_type == 'DeepSDF':
        cd_re2gt = np.mean(re2gt_distances**2)
        cd_gt2re = np.mean(gt2re_distances**2)
        hd_re2gt = np.max(re2gt_distances)
        hd_gt2re = np.max(gt2re_distances)
        chamfer_dist = cd_re2gt + cd_gt2re
        hausdorff_distance = np.max((hd_re2gt, hd_gt2re))
    else:
        cd_re2gt = np.mean(re2gt_distances)
        cd_gt2re = np.mean(gt2re_distances)
        hd_re2gt = np.max(re2gt_distances)
        hd_gt2re = np.max(gt2re_distances)
        chamfer_dist = 0.5* (cd_re2gt + cd_gt2re)
        hausdorff_distance = np.max((hd_re2gt, hd_gt2re))
    return chamfer_dist, hausdorff_distance, cd_re2gt, cd_gt2re, hd_re2gt, hd_gt2re

shapes = ['anchor', 'daratech', 'dc', 'gargoyle', 'lord_quas']
for shape in shapes:
    scan_shape_path = os.path.join(scan_path, '{}.ply'.format(shape))
    gt_shape_path = os.path.join(gt_path, '{}.xyz'.format(shape))
    recon_shape_path = os.path.join(mesh_path, '{}.ply'.format(shape))
    # recon_shape_path = os.path.join(mesh_path, '{}_iter_6000.ply'.format(shape))
    # print(recon_shape_path)
    if not os.path.exists(recon_shape_path):
        continue
    scan_pc = trimesh.load(scan_shape_path)
    gt_pc = trimesh.load(gt_shape_path)
    recon_mesh = trimesh.load(recon_shape_path)

    scan_points = scan_pc.vertices
    gt_points = gt_pc.vertices
    # recon_points = trimesh.sample.sample_surface(recon_mesh, 30000)[0]
    recon_points = trimesh.sample.sample_surface(recon_mesh, 1000000)[0]
    chamfer_dist, hausdorff_distance,  cd_re2gt, cd_gt2re, hd_re2gt, hd_gt2re = compute_dists(recon_points, gt_points)
    sc_chamfer_dist, sc_hausdorff_distance, sc_cd_re2gt, sc_cd_gt2re, sc_hd_re2gt, sc_hd_gt2re = compute_dists(recon_points, scan_points)
    print('{: <10} {:.3f} {:.3f} {:.3f} {:.3f}'.format(shape, chamfer_dist, hausdorff_distance, sc_cd_re2gt, sc_hd_re2gt))
print()
surface_reconstruction/compute_metrics_srb.py
sunyx523/StEik
# This file is partly based on DiGS: https://github.com/Chumbyte/DiGS
import os
import os.path
import torch.utils.data as data
import numpy as np
import scipy.spatial as spatial
from abc import ABC, abstractmethod
import torch
import open3d as o3d


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)


        # 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()

        # generate grid points and find distance to closest point on the line
        self.grid_range = grid_range  # larger than 1 to avoid boundary issues
        x, y, z = np.linspace(-grid_range, grid_range, self.grid_res).astype(np.float32), \
                  np.linspace(-grid_range, grid_range, self.grid_res).astype(np.float32),\
                  np.linspace(-grid_range, grid_range, self.grid_res).astype(np.float32)
        xx, yy, zz = np.meshgrid(x, y, z)
        xx, yy, zz = xx.ravel(), yy.ravel(), zz.ravel()
        self.grid_points = np.stack([xx, yy, zz], axis=1).astype('f')
        self.nonmnfld_points = self.get_nonmnfld_points()

        self.point_idxs = np.arange(self.points.shape[0], dtype=np.int32)
        self.nonmnfld_points_idxs = np.arange(self.nonmnfld_points.shape[0], dtype=np.int32)
        self.generate_batch_indices()


    def generate_batch_indices(self):
        mnfld_idx = []
        nonmnfld_idx = []
        for i in range(self.n_samples):
            mnfld_idx.append(np.random.choice(self.point_idxs, self.n_points))
            nonmnfld_idx.append(np.random.choice(self.nonmnfld_points_idxs, self.n_points))
        self.mnfld_idx = np.array(mnfld_idx)
        self.nonmnfld_idx = np.array(nonmnfld_idx)


    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)
        # center and scale point cloud
        self.cp = points.mean(axis=0)
        points = points - self.cp[None, :]
        # self.scale = np.linalg.norm(points, axis=-1).max(-1)
        self.scale = np.abs(points).max()
        points = points / self.scale

        return points, normals

    def get_nonmnfld_points(self):
        if self.sample_type == 'grid':
            nonmnfld_points = self.grid_points
        elif self.sample_type == 'uniform':
            nonmnfld_points = np.random.uniform(-self.grid_range, self.grid_range,
                                                size=(self.grid_res * self.grid_res * self.grid_res, 3)).astype(np.float32)
        elif self.sample_type == 'gaussian':
            nonmnfld_points = self.sample_gaussian_noise_around_shape()
            idx = np.random.choice(range(nonmnfld_points.shape[1]), self.grid_res * self.grid_res)
            sample_idx = np.random.choice(range(nonmnfld_points.shape[0]), self.grid_res * self.grid_res)
            nonmnfld_points = nonmnfld_points[sample_idx, idx]
        elif self.sample_type == 'combined':
            nonmnfld_points1 = self.sample_gaussian_noise_around_shape()
            nonmnfld_points2 = self.grid_points
            idx1 = np.random.choice(range(nonmnfld_points1.shape[1]), int(np.ceil(self.grid_res * self.grid_res / 2)))
            idx2 = np.random.choice(range(nonmnfld_points2.shape[0]), int(np.floor(self.grid_res * self.grid_res / 2)))
            sample_idx = np.random.choice(range(nonmnfld_points1.shape[0]), int(np.ceil(self.grid_res * self.grid_res / 2)))

            nonmnfld_points = np.concatenate([nonmnfld_points1[sample_idx, idx1], nonmnfld_points2[idx2]], axis=0)
        else:
            raise Warning("Unsupported non manfold sampling type")
        return nonmnfld_points

    def sample_gaussian_noise_around_shape(self):
        n_noisy_points = int(np.round(self.grid_res * self.grid_res / self.n_points))
        noise = np.random.multivariate_normal([0, 0, 0], [[self.sampling_std, 0, 0], [0, self.sampling_std, 0],
                                                          [0, 0, self.sampling_std]],
                                              size=(self.points.shape[0], n_noisy_points)).astype(np.float32)
        nonmnfld_points = np.tile(self.points[:, None, :], [1, n_noisy_points, 1]) + noise
        nonmnfld_points = nonmnfld_points.reshape([nonmnfld_points.shape[0], -1, nonmnfld_points.shape[-1]])
        return nonmnfld_points

    @abstractmethod
    def get_points_distances_and_normals(self, points):
        # implement a function that computes the distance and normal vectors of nonmanifold points.
        # default implementation finds the nearest neighbor and return its normal and the distance projected onto it.
        # which is a coarse approximation

        # compute distance and normal (general case - closest point)
        kdtree = spatial.cKDTree(self.points)
        distances, nn_idx = kdtree.query(points, k=1, n_jobs=-1)
        normals = self.mnfld_n[nn_idx]
        p1p2 = points - self.points[nn_idx]
        sings = np.sign(np.einsum('ij,ij->i', p1p2, self.mnfld_n[nn_idx]))
        distances = sings * distances
        return distances.astype('f'), normals.astype('f')


    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)

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

    def __len__(self):
        return self.n_samples





surface_reconstruction/recon_dataset.py
sunyx523/StEik
# This file is partly based on DiGS: https://github.com/Chumbyte/DiGS
import argparse
import torch
import os


def get_train_args():
    parser = argparse.ArgumentParser(description='Local implicit functions experiment')
    parser.add_argument('--gpu_idx', type=int, default=0, help='set < 0 to use CPU')
    parser.add_argument('--logdir', type=str, default='./log/debug', help='log directory')
    parser.add_argument('--seed', type=int, default=3627473, help='random seed')
    parser.add_argument('--dataset_path', type=str, default='../data/deep_geometric_prior_data', help='path to dataset folder')
    parser.add_argument('--raw_dataset_path', type=str, default='../data/deep_geometric_prior_data', help='path to dataset folder')
    parser.add_argument('--file_name', type=str, default='lord_quas.ply',
                        help='name of file to reconstruct (within the dataset path)')
    parser.add_argument('--n_iterations', type=int, default=5000,
                        help='number of interations in the generated train and test set')
    parser.add_argument('--model_dirpath', type=str,
                        default='./models',
                        help='path to model directory for backup')
    parser.add_argument('--parallel', type=int, default=False, help='use data parallel')
    parser.add_argument('--results_path', type=str, default='./log/surface_reconstruction/DiGS_surf_recon_experiment/result_meshes',
                        help='path to results directory')

    # training paramaeters
    parser.add_argument('--lr', type=float, default=1e-4, 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('--n_points', type=int, default=30000, help='number of points in each point cloud')

    # Network architecture and loss
    parser.add_argument('--decoder_hidden_dim', type=int, default=256, help='length of decoder hidden dim')
    parser.add_argument('--encoder_hidden_dim', type=int, default=128, help='length of encoder hidden dim')
    parser.add_argument('--decoder_n_hidden_layers', type=int, default=8, help='number of decoder hidden layers')
    parser.add_argument('--nl', type=str, default='softplus', help='type of non linearity sine | relu')
    parser.add_argument('--latent_size', type=int, default=0, help='number of elements in the latent vector,'
                                                                   ' use 0 for reconstruction')
    parser.add_argument('--sphere_init_params', nargs='+', type=float, default=[1.6, 1.0],
                    help='radius and scaling')
    parser.add_argument('--neuron_type', type=str, default='quadratic', help='type of neuron')


    parser.add_argument('--encoder_type', type=str, default='none', help='type of encoder none | pointnet')
    parser.add_argument('--loss_type', type=str, default='siren', help='loss type to use: siren | siren_wo_n | igr | igr_wo_n'
                                                                       'siren_wo_n_w_div | siren_wo_n_w_div_w_reg |...')
    parser.add_argument('--div_decay_params', nargs='+', type=float, default=[0.0, 0.5, 0.75],
                        help='epoch number to evaluate')
    parser.add_argument('--div_type', type=str, default='dir_l1', help='divergence term norm dir_l1 | dir_l2 | full_l1 | full_l2')
    parser.add_argument('--grid_res', type=int, default=128,
                        help='uniform grid resolution')
    parser.add_argument('--div_decay', type=str, default='linear',
                        help='divergence term importance decay none | step | linear')
    parser.add_argument('--nonmnfld_sample_type', type=str, default='uniform',
                        help='how to sample points off the manifold - grid | gaussian | combined')
    parser.add_argument('--init_type', type=str, default='mfgi',
                        help='initialization type siren | geometric_sine | geometric_relu | mfgi')
    parser.add_argument('--loss_weights', nargs='+', type=float, default=[3e3, 1e2, 1e2, 5e1, 1e1],
                        help='loss terms weights sdf | inter | normal | eikonal | div')
    args = parser.parse_args()
    return args


def get_test_args():
    parser = argparse.ArgumentParser(description='Local implicit functions test experiment')
    parser.add_argument('--gpu_idx', type=int, default=1, help='set < 0 to use CPU')
    parser.add_argument('--logdir', type=str, default='./logs/', help='log directory')
    parser.add_argument('--file_name', type=str, default='daratech.ply', help='trained model name')
    parser.add_argument('--n_points', type=int, default=0, help='number of points in each point cloud, '
                                                                  'if 0 use training options')
    parser.add_argument('--batch_size', type=int, default=0, help='number of samples in a minibatch, if 0 use training')
    parser.add_argument('--grid_res', type=int, default=512, help='grid resolution for reconstruction')
    parser.add_argument('--export_mesh', type=bool, default=True, help='indicator to export mesh as ply file')
    parser.add_argument('--dataset_path', type=str, default='', help='path to dataset folder')

    test_opt = parser.parse_args()
    test_opt.logdir = os.path.join(test_opt.logdir, test_opt.file_name.split('.')[0])
    param_filename = os.path.join(test_opt.logdir, 'trained_models/params.pth')
    train_opt = torch.load(param_filename)

    test_opt.nl, test_opt.latent_size, test_opt.encoder_type, test_opt.n_iterations, test_opt.seed, \
        test_opt.decoder_hidden_dim, test_opt.encoder_hidden_dim, test_opt.decoder_n_hidden_layers, test_opt.init_type, test_opt.neuron_type \
        = train_opt.nl, train_opt.latent_size, train_opt.encoder_type, train_opt.n_iterations, train_opt.seed, \
          train_opt.decoder_hidden_dim, train_opt.encoder_hidden_dim, train_opt.decoder_n_hidden_layers, train_opt.init_type, train_opt.neuron_type
    test_opt.n_point_total = train_opt.n_points

    if test_opt.n_points == 0:
        test_opt.n_points = train_opt.n_points
    if "parallel" in train_opt:
        test_opt.parallel = train_opt.parallel
    else:
        test_opt.parallel = False
    return test_opt

surface_reconstruction/surface_recon_args.py
sunyx523/StEik
# This file is partly based on DiGS: https://github.com/Chumbyte/DiGS
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import recon_dataset as dataset
import torch
import utils.visualizations as vis
import numpy as np
import models.Net as Net
import torch.nn.parallel
import utils.utils as utils
import surface_recon_args

# get training parameters
args = surface_recon_args.get_test_args()

file_path = os.path.join(args.dataset_path, args.file_name)
if args.export_mesh:
    outdir = os.path.join(os.path.dirname(args.logdir), 'result_meshes')
    os.makedirs(outdir, exist_ok=True)
    output_ply_filepath = os.path.join(outdir, args.file_name)
# get data loader
torch.manual_seed(args.seed)
np.random.seed(args.seed)
test_set = dataset.ReconDataset(file_path, args.n_points, n_samples=1, res=args.grid_res, sample_type='grid',
                                requires_dist=False)
test_dataloader = torch.utils.data.DataLoader(test_set, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)

# get model
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_idx)
device = torch.device("cuda")

Net = Net.Network(latent_size=args.latent_size, in_dim=3, decoder_hidden_dim=args.decoder_hidden_dim,
                           nl=args.nl, encoder_type=args.encoder_type,
                           decoder_n_hidden_layers=args.decoder_n_hidden_layers,
                            init_type=args.init_type, neuron_type=args.neuron_type)
if args.parallel:
    if (device.type == 'cuda'):
        Net = torch.nn.DataParallel(Net)

model_dir = os.path.join(args.logdir, 'trained_models')
trained_model_filename = os.path.join(model_dir, 'model.pth')
Net.load_state_dict(torch.load(trained_model_filename, map_location=device))
Net.to(device)
latent = None

print("Converting implicit to mesh for file {}".format(args.file_name))
cp, scale, bbox = test_set.cp, test_set.scale, test_set.bbox
test_set, test_dataloader, clean_points_gt, normals_gt,  nonmnfld_points, data = None, None, None, None, None, None  # free up memory
mesh_dict = utils.implicit2mesh(Net.decoder, latent, args.grid_res, translate=-cp, scale=1/scale,
                                get_mesh=True, device=device, bbox=bbox)
vis.plot_mesh(mesh_dict["mesh_trace"], mesh=mesh_dict["mesh_obj"], output_ply_path=output_ply_filepath, show_ax=False,
              title_txt=args.file_name.split('.')[0], show=False)

print("Conversion complete.")
surface_reconstruction/test_surface_reconstruction.py
sunyx523/StEik
# This file is partly based on DiGS: https://github.com/Chumbyte/DiGS
import os, sys, time
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import recon_dataset as dataset
import numpy as np
import models.Net as Net
from models.losses import Loss
import torch.nn.parallel
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
import utils.utils as utils
import surface_recon_args


# get training parameters
args = surface_recon_args.get_train_args()

file_path = os.path.join(args.dataset_path, args.file_name)
logdir = os.path.join(args.logdir, args.file_name.split('.')[0])

# 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' % ('../models/Net.py', logdir))  # backup the models files
os.system('cp %s %s' % ('../models/losses.py', logdir))  # backup the losses files

# get data loaders
torch.manual_seed(0)  #change random seed for training set (so it will be different from test set
np.random.seed(0)
train_set = dataset.ReconDataset(file_path, args.n_points, args.n_iterations, args.grid_res, args.nonmnfld_sample_type)
train_dataloader = torch.utils.data.DataLoader(train_set, batch_size=1, shuffle=True, num_workers=4,
                                               pin_memory=True)
torch.manual_seed(args.seed)
np.random.seed(args.seed)

# get model
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_idx)
device = torch.device("cuda")

Net = Net.Network(latent_size=args.latent_size, in_dim=3, decoder_hidden_dim=args.decoder_hidden_dim, nl=args.nl,
                  encoder_type=args.encoder_type, decoder_n_hidden_layers=args.decoder_n_hidden_layers,
                  init_type=args.init_type, neuron_type=args.neuron_type, sphere_init_params=args.sphere_init_params)
if args.parallel:
    if (device.type == 'cuda'):
        Net = torch.nn.DataParallel(Net)

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, betas=(0.9, 0.999))
scheduler = lr_scheduler.ExponentialLR(optimizer, gamma=1.0) # Does nothing
Net.to(device)
criterion = Loss(weights=args.loss_weights, loss_type=args.loss_type, div_decay=args.div_decay, div_type=args.div_type)
num_batches = len(train_dataloader)

# For each batch in the dataloader
for batch_idx, data in enumerate(train_dataloader):
    # if batch_idx in [0, 1, 5, 10, 50, 100] or batch_idx % 500 == 0:
    #     try:
    #         shapename = args.file_name.split('.')[0]
    #         output_dir = os.path.join(logdir, '..', 'result_meshes')
    #         os.makedirs(output_dir, exist_ok=True)
    #         output_ply_filepath = os.path.join(output_dir, shapename+'_iter_{}.ply'.format(batch_idx))
    #         print('Saving to ', output_ply_filepath)
    #         cp, scale, bbox = train_set.cp, train_set.scale, train_set.bbox
    #         mesh_dict = utils.implicit2mesh(DiGSNet.decoder, None, 128, translate=-cp, scale=1/scale,
    #                             get_mesh=True, device=device, bbox=bbox)
    #         mesh_dict["mesh_obj"].export(output_ply_filepath, vertex_normal=True)
    #     except:
    #         print(traceback.format_exc())
    #         print('Could not generate mesh')
    #         print()

    Net.zero_grad()
    Net.train()

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

    mnfld_points.requires_grad_()
    nonmnfld_points.requires_grad_()

    output_pred = Net(nonmnfld_points, mnfld_points)

    loss_dict, _ = criterion(output_pred, mnfld_points, nonmnfld_points, mnfld_n_gt)
    lr = torch.tensor(optimizer.param_groups[0]['lr'])
    loss_dict["lr"] = lr
    utils.log_losses(log_writer_train, batch_idx, num_batches, loss_dict)

    loss_dict["loss"].backward()

    torch.nn.utils.clip_grad_norm_(Net.parameters(), 10.)

    optimizer.step()

    #Output training stats
    if batch_idx % 100 == 0:

        weights = criterion.weights
        utils.log_string("Weights: {}, lr={:.3e}".format(weights, lr), log_file)
        utils.log_string('Iteration: {:4d}/{} ({:.0f}%) Loss: {:.5f} = L_Mnfld: {:.5f} + '
                'L_NonMnfld: {:.5f} + L_Nrml: {:.5f} + L_Eknl: {:.5f} + L_Div: {:.5f}'.format(
             batch_idx, 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()), log_file)
        utils.log_string('Iteration: {:4d}/{} ({:.0f}%) Unweighted L_s : L_Mnfld: {:.5f},  '
                'L_NonMnfld: {:.5f},  L_Nrml: {:.5f},  L_Eknl: {:.5f},  L_Div: {:.5f}'.format(
             batch_idx, 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["div_loss"].item()), log_file)
        utils.log_string('', log_file)

    criterion.update_div_weight(batch_idx, args.n_iterations,
                                args.div_decay_params)  # assumes batch size of 1
    scheduler.step()

    # save model
    # if batch_idx % 1000 == 0 :
        # utils.log_string("saving model to file :generator_model.pth", log_file)
        # torch.save(Net.state_dict(),
        #             os.path.join(model_outdir, 'model.pth' ))
# save model
utils.log_string("saving model to file :generator_model.pth", log_file)
torch.save(Net.state_dict(), os.path.join(model_outdir, 'model.pth' ))
surface_reconstruction/train_surface_reconstruction.py
sunyx523/StEik
# This file is borrowed from DiGS: https://github.com/Chumbyte/DiGS
import os
from tensorboardX import SummaryWriter
import io
from PIL import Image
import numpy as np
import torch
from torch.autograd import grad
import warnings
from plyfile import PlyData
from scipy.spatial import cKDTree as KDTree
from tqdm import tqdm
from skimage import measure
import plotly.graph_objs as go
import trimesh


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, batch_idx, num_batches, loss_dict):
    # helper function to log losses to tensorboardx writer
    iteration = batch_idx + 1
    for loss in loss_dict.keys():
        writer.add_scalar(loss, loss_dict[loss].item(), iteration)
    return iteration

def log_weight_hist(writer, batch_idx, num_batches, net_blocks):
    # helper function to log losses to tensorboardx writer
    iteration = batch_idx + 1
    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, 'params.pth')
        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 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])

    #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


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.float16)
    return {"grid_points": grid_points,
            "shortest_axis_length": length,
            "xyz": [x, y, z],
            "shortest_axis_index": shortest_axis}


def read_vnf_ply(filename):
    # read vertices and normal vectors from a ply file
    plydata = PlyData.read(filename)
    x = np.asarray(plydata.elements[0].data['x'])
    y = np.asarray(plydata.elements[0].data['y'])
    z = np.asarray(plydata.elements[0].data['z'])
    nx = np.asarray(plydata.elements[0].data['nx'])
    ny = np.asarray(plydata.elements[0].data['ny'])
    nz = np.asarray(plydata.elements[0].data['nz'])
    return np.stack([x, y, z], axis=1), np.stack([nx, ny, nz], axis=1), plydata['face'].data


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, latent, grid_res, translate=[0., 0., 0.], scale=1, get_mesh=True, device=None,
                  bbox=np.array([[-1, 1], [-1, 1], [-1, 1]])):
    # 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)):
    for point in tqdm(torch.split(pnts, 100000, dim=0)):
        # point: (100000, 3)
        if (not latent is None):
            point = torch.cat([point, latent.unsqueeze(0).repeat(point.shape[0], 1), ], dim=1)
        point = get_cuda_ifavailable(point, device=device)
        z.append(decoder(point.type(torch.float32)).detach().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)
    # import pdb; pdb.set_trace()
    print(z.min(), z.max())

    verts, faces, normals, values = measure.marching_cubes(volume=z, level=0.0,
                                                                   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)
    trace = None
    # I, J, K = ([triplet[c] for triplet in faces] for c in range(3))

    # trace = [go.Mesh3d(x=verts[:, 0], y=verts[:, 1], z=verts[:, 2],
    #                       i=I, j=J, k=K, name='',
    #                       color='orange', opacity=0.5)]

    return { "mesh_trace": trace,
            "mesh_obj": mesh }


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
sunyx523/StEik
# This file is borrowed from DiGS: https://github.com/Chumbyte/DiGS
import plotly.graph_objects as go
import plotly.figure_factory as ff
import numpy as np
import plotly.offline as offline
import utils.utils as utils
import torch
import matplotlib.pyplot as plt
from PIL import Image
import os

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,
                               # autocontour=True,
                               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,
                                          example_idx=0, n_gt=None, n_pred=None,
                                          show_ax=True, show_bar=True, title_text='', colorscale='Geyser',
                                          res=256):
    # 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 = 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,
                           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 = []

    # 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")
    # 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
                               ))  # 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)
    return dist_img, curl_img, eikonal_img, div_img, z_diff_img


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='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,
                            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

    # 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,
                            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
sunyx523/StEik