PathNet: Path-selective Point Cloud Denoising
# *_*coding:utf-8 *_*
import os
import warnings
import numpy as np
import h5py
from torch.utils.data import Dataset
from data_utils.Pointfilter_Utils import pca_alignment

warnings.filterwarnings('ignore')

def pca_normalize(pc,pc2):
    normalize_data = np.zeros(pc.shape, dtype=np.float32)
    normalize_data2 = np.zeros(pc2.shape, dtype=np.float32)

    centroid = pc[:,0,:]
    centroid = np.expand_dims(centroid, axis=1)
    pc = pc - centroid
    pc2 = pc2 - centroid

    m = np.max(np.sqrt(np.sum(pc**2, axis=2)),axis = 1, keepdims=True)
    pc = pc / np.expand_dims(m, axis=-1)
    pc2 = pc2 / np.expand_dims(m, axis=-1)

    for B in range(pc.shape[0]):
        x, pca_martix_inv = pca_alignment(pc[B,:,:])
        x2 = np.array(np.linalg.inv(pca_martix_inv) * np.matrix(pc2[B,:,:].T)).T
        normalize_data[B, ...] = x
        normalize_data2[B, ...] = x2

    return normalize_data, normalize_data2

class PatchDataset(Dataset):
    def __init__(self,root = './data/', npoints=128, split='train', class_choice=None, normal_channel=False):
        self.npoints = npoints
        self.root = root

        self.catfile = os.path.join(self.root, 'train _data.hdf5')

        f = h5py.File(self.catfile,'r')

        self.inputs = f["inputs"][:]  #B, N ,3
        self.target = f["target"][:]
        self.label = f["label"][:]

        self.inputs,self.target= pca_normalize(self.inputs,self.target)
 
        self.label = self.label //2.3

        idx = np.arange(0,self.inputs.shape[0])
        np.random.seed(1111)
        np.random.shuffle(idx)
        self.inputs = self.inputs[idx][:,:,:3]
        self.target = self.target[idx][:,:,:3]
        self.label = self.label[idx][:]

        sample_size = int(self.inputs.shape[0] * 0.8)
        if(split == 'train'):
            self.inputs = self.inputs[:sample_size]
            self.target = self.target[:sample_size]
            self.label = self.label[:sample_size]
        elif(split == 'test'):
            self.inputs = self.inputs[sample_size:]
            self.target = self.target[sample_size:]
            self.label = self.label[:sample_size]

        print('The size of %s inputs is %d'%(split,self.inputs.shape[0]))


        self.seg_classes = {'circle': [0,1]}

        self.cache = {}  
        self.cache_size = 1000


    def __getitem__(self, index):
        if index in self.cache:
            inputs, target, label = self.cache[index]
        else:
            inputs = self.inputs[index].astype(np.float32)  #N,3
            target = self.target[index].astype(np.float32)
            label = self.label[index].astype(np.float32)

            if len(self.cache) < self.cache_size:
                self.cache[index] = (inputs, target, label)

        return inputs, target, label

    def __len__(self):
        return self.inputs.shape[0]




data_utils/DataLoader.py
ZeyongWei/PathNet
import scipy.spatial as sp
import numpy as np
import torch

import os

from Customer_Module.chamfer_distance.dist_chamfer import chamferDist
from plyfile import PlyData, PlyElement
nnd = chamferDist()

import logging

def log_string(str):
    logger.info(str)
    print(str)

name_dirs = ['benchmark81_10000','benchmark81_20000','benchmark81_50000','kinect_fusion','kinect_v1','kinect_v2']

name_dir = 'benchmark81_20000'

model = ''

iter_num = 2

results_dir = '//' + model + '/' + name_dir + '/'

back_logs = ['_output_end','_output_end_output_end']

back_log = back_logs[iter_num-1]

if(not os.path.exists('./data/test_data/' + name_dir + '/eval/' + model +'/')):
   os.makedirs('./data/test_data/' + name_dir + '/eval/' + model +'/')

if(os.path.exists('./data/test_data/' + name_dir + '/eval/' + model +'/eval_log_' + name_dir +'_' + str(iter_num) + '.csv')):
    os.remove('./data/test_data/' + name_dir + '/eval/' + model +'/eval_log_' + name_dir +'_' + str(iter_num) + '.csv')

logger = logging.getLogger("Eval"+'_' + model + '_' + name_dir + '_' + str(iter_num))
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('./data/test_data/' + name_dir + '/eval/' + model +'/eval_log_' + name_dir +'_' + str(iter_num) + '.csv')
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)


def Eval_With_Charmfer_Distance():
    log_string('************Errors under Chamfer Distance************')
    for shape_id, shape_name in enumerate(shape_names):
        if(name_dir.split('_')[0]=='benchmark81'):
            gt_pts = np.loadtxt(os.path.join('./data/test_data/gt/'+ name_dir +'/', shape_name.split('_0.')[0] + '.xyz'))
        else:
            gt_pts = np.loadtxt(os.path.join('./data/test_data/gt/'+ name_dir +'/', shape_name.split('_noisy')[0] + '.xyz'))
        pred_pts = np.loadtxt(os.path.join('./data/results/'+ results_dir +'/', shape_name + back_log + '.xyz'))[:,:3]
        with torch.no_grad():
            gt_pts_cuda = torch.from_numpy(np.expand_dims(gt_pts, axis=0)).cuda().float()
            pred_pts_cuda = torch.from_numpy(np.expand_dims(pred_pts, axis=0)).cuda().float()
            dist1, dist2 = nnd(pred_pts_cuda, gt_pts_cuda)
            chamfer_errors = torch.mean(dist1, dim=1) + torch.mean(dist2, dim=1)

            log_string('%12s  %.3f' % (shape_names[shape_id], round(chamfer_errors.item() * 100000, 3)))

def Eval_With_Mean_Square_Error():
    log_string('************Errors under Mean Square Error************')
    for shape_id, shape_name in enumerate(shape_names):
        if(name_dir.split('_')[0]=='benchmark81'):
            gt_pts = np.loadtxt(os.path.join('./data/test_data/gt/' + name_dir +'/', shape_name.split('_0.')[0] + '.xyz'))
        else:
            gt_pts = np.loadtxt(os.path.join('./data/test_data/gt/' + name_dir +'/', shape_name.split('_noisy')[0] + '.xyz'))
        gt_pts_tree = sp.cKDTree(gt_pts)
        pred_pts = np.loadtxt(os.path.join('./data/results/'+ results_dir +'/', shape_name + back_log +'.xyz'))[:,:3]
        pred_dist, _ = gt_pts_tree.query(pred_pts, 10)

        log_string('%12s  %.3f' % (shape_names[shape_id], round(pred_dist.mean() * 1000, 3)))


if __name__ == '__main__':

    with open(os.path.join('./data/test_data/gt/'+ name_dir +'/', 'test.txt'), 'r') as f:
        shape_names = f.readlines()
    shape_names = [x.strip() for x in shape_names]
    shape_names = list(filter(None, shape_names))

    Eval_With_Charmfer_Distance()
    Eval_With_Mean_Square_Error()


eval_ours.py
ZeyongWei/PathNet
import torch.nn as nn
import torch.utils.data
import torch
import torch.nn.functional as F
import numpy as np
from einops import rearrange, repeat

class Analyser_Block(nn.Module):
    def __init__(self, path_num = 3):
        super(Analyser_Block, self).__init__()
        self.path_num = path_num

        #encoder
        self.mlp_convs1 = nn.Conv1d(512,512,1)
        self.mlp_convs2 = nn.Conv1d(512,512,1)

        self.mlp_bns1 = nn.BatchNorm1d(512)
        self.mlp_bns2 = nn.BatchNorm1d(512)
        
        # decoder
        self.fc1 = nn.Linear(512, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, self.path_num)

        self.bn1 = nn.BatchNorm1d(256)
        self.bn2 = nn.BatchNorm1d(128)
        self.drop1 = nn.Dropout(0.4)
        self.drop2 = nn.Dropout(0.4)

    def forward(self, f1):
        #encoder
        f1 = F.relu(self.mlp_bns1(self.mlp_convs1(f1)))
        f1 = F.relu(self.mlp_bns2(self.mlp_convs2(f1)))

        #decoder
        f1_max = torch.max(f1, axis = 2)[0]

        path_f1 = self.drop1(F.relu(self.bn1(self.fc1(f1_max))))
        path_f1 = self.drop2(F.relu(self.bn2(self.fc2(path_f1))))
        path_f1 = self.fc3(path_f1)  #B,3
        path_f1 = F.softmax(path_f1,dim=1)

        return path_f1

class get_analyser(nn.Module):
    def __init__(self, block_num = 3, path_num = 3):
        super(get_analyser, self).__init__()
        self.block_num = block_num
        self.path_num = path_num
        #Analyser_Block
        self.analyser_blocks = nn.ModuleList()
        for i in range(self.block_num):
            self.analyser_blocks.append(Analyser_Block(self.path_num))

    def forward(self, f1, ab_i = 0):

        path_f1 = self.analyser_blocks[ab_i](f1)

        return path_f1

class Path_Block(nn.Module):
    def __init__(self, path_num = 3):
        super(Path_Block, self).__init__()
        self.path_num = path_num

        #path-head
        self.mlp_convs_ph1 = nn.Conv1d(512,512,1)
        self.mlp_convs_ph2 = nn.Conv1d(512,512,1)

        self.mlp_bns_ph1 = nn.BatchNorm1d(512)
        self.mlp_bns_ph2 = nn.BatchNorm1d(512)

        #path0
        #pass

        #path1
        self.mlp_convs_p11 = nn.Conv1d(512,256,1)
        self.mlp_convs_p12 = nn.Conv1d(512,512,1)

        self.mlp_bns_p11 = nn.BatchNorm1d(256)
        self.mlp_bns_p12 = nn.BatchNorm1d(512)

        #path2
        self.mlp_convs_p21 = nn.Conv1d(512,256,1)
        self.mlp_convs_p22 = nn.Conv1d(512,512,1)

        self.mlp_bns_p21 = nn.BatchNorm1d(256)
        self.mlp_bns_p22 = nn.BatchNorm1d(512)

        self.mlp_convs_p23 = nn.Conv1d(512,256,1)
        self.mlp_convs_p24 = nn.Conv1d(512,512,1)

        self.mlp_bns_p23 = nn.BatchNorm1d(256)
        self.mlp_bns_p24 = nn.BatchNorm1d(512)

        #decoder
        self.fc1 = nn.Linear(512, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 3)

        self.bn1 = nn.BatchNorm1d(256)
        self.bn2 = nn.BatchNorm1d(128)
        self.drop1 = nn.Dropout(0.4)
        self.drop2 = nn.Dropout(0.4)

    def forward(self, f1, path_f1):
        #path_f1.shape #B

        #path_head
        f1_temp = F.relu(self.mlp_bns_ph1(self.mlp_convs_ph1(f1)))  #B, d, N
        f1 = f1 + F.relu(self.mlp_bns_ph2(self.mlp_convs_ph2(f1_temp)))  #B, d, N

        #path_chose

        if(self.path_num == 2):

            f1_out = f1
            idx = torch.where(path_f1==1)
            if(idx[0].shape[0]>0):
                f1_out[idx] = self.path1(f1_out[idx])
        
        #decoder
        t1 = self.path_denoise(f1_out)

        return f1_out, t1

    def path0(self,f1):

        #pass
        return f1

    def path1(self,f1):

        f1_temp = F.relu(self.mlp_bns_p11(self.mlp_convs_p11(f1)))
        f1_temp_max = torch.max(f1_temp, axis = 2)[0]
        f1_temp_max = f1_temp_max.unsqueeze(-1).repeat(1,1,f1_temp.shape[-1])
        f1_cat = torch.cat((f1_temp,f1_temp_max),1)
        f1 = f1 + F.relu(self.mlp_bns_p12(self.mlp_convs_p12(f1_cat)))

        return f1

    def path2(self,f1):
        f1_temp = F.relu(self.mlp_bns_p21(self.mlp_convs_p21(f1)))
        f1_temp_max = torch.max(f1_temp, axis = 2)[0]
        f1_temp_max = f1_temp_max.unsqueeze(-1).repeat(1,1,f1_temp.shape[-1])
        f1_temp = torch.cat((f1_temp,f1_temp_max),1)
        f2 = f1 + F.relu(self.mlp_bns_p22(self.mlp_convs_p22(f1_temp)))

        f1_temp = F.relu(self.mlp_bns_p23(self.mlp_convs_p23(f2)))
        f1_temp_max = torch.max(f1_temp, axis = 2)[0]
        f1_temp_max = f1_temp_max.unsqueeze(-1).repeat(1,1,f1_temp.shape[-1])
        f1_temp = torch.cat((f1_temp,f1_temp_max),1)
        f1 = f1 + F.relu(self.mlp_bns_p24(self.mlp_convs_p24(f1_temp)))

        return f1

    def path_denoise(self, f1):

        f1_max = torch.max(f1, axis = 2)[0]

        t1 = self.drop1(F.relu(self.bn1(self.fc1(f1_max))))
        t1 = self.drop2(F.relu(self.bn2(self.fc2(t1))))
        t1 = self.fc3(t1) #B,3

        return t1

class get_model(nn.Module):
    def __init__(self, block_num = 3, path_num = 3):
        super(get_model, self).__init__()
        channel = 3
        self.block_num = block_num
        self.path_num = path_num

        #encoders
        self.mlp = np.array([64,128,256,512])

        self.mlp_convs = nn.ModuleList()

        self.mlp_bns = nn.ModuleList()

        last_channel = channel
        for i in range(self.mlp.shape[0]):
            self.mlp_convs.append(nn.Conv1d(last_channel, self.mlp[i], 1))
            self.mlp_bns.append(nn.BatchNorm1d(self.mlp[i]))

            last_channel = self.mlp[i]

        #path_blocks
        self.pbs = nn.ModuleList()
        for i in range(self.block_num):
            self.pbs.append(Path_Block(self.path_num))


    def forward(self, x, analyser, use_random_path = 0):
        B, _, N = x.shape

        #encoder
        x = x - x[:,:,0:1]  #B,3,N

        f0 = x
        for i in range(self.mlp.shape[0]):
            f0 = F.relu(self.mlp_bns[i](self.mlp_convs[i](f0)))

        feature_m = []
        path_maxprob_m = []
        path_m = []
        trans_m = []

        f1 = f0
        for pb_i in range(self.block_num):
            #path_analyser
            if(use_random_path == 1):
                path_prob_f1 = torch.rand(B,self.path_num)
                path_prob_f1 = F.softmax(path_prob_f1,-1)

            elif(pb_i < self.block_num and use_random_path == 0): 
                path_prob_f1 = analyser(f1, pb_i)

            else:
                print("error")
                return 0,0,0

            path_maxprob_f1, path_f1 = torch.max(path_prob_f1,axis = -1)#B
 
            #path_block
            f1, t1 = self.pbs[pb_i](f1,path_f1)

            #recode
            feature_m.append(f1)
            path_maxprob_m.append(path_maxprob_f1)
            path_m.append(path_f1)
            trans_m.append(t1) 

        return trans_m, path_m, path_maxprob_m


models/model.py
ZeyongWei/PathNet
import torch
import torch.nn as nn
import torch.nn.functional as F
from time import time
import numpy as np

def timeit(tag, t):
    print("{}: {}s".format(tag, time() - t))
    return time()

def pc_normalize(pc):
    l = pc.shape[0]
    centroid = np.mean(pc, axis=0)
    pc = pc - centroid
    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
    pc = pc / m
    return pc

def square_distance(src, dst):
    """
    Calculate Euclid distance between each two points.

    src^T * dst = xn * xm + yn * ym + zn * zm；
    sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
    sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
         = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst

    Input:
        src: source points, [B, N, C]
        dst: target points, [B, M, C]
    Output:
        dist: per-point square distance, [B, N, M]
    """
    B, N, _ = src.shape
    _, M, _ = dst.shape
    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
    dist += torch.sum(src ** 2, -1).view(B, N, 1)
    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
    return dist


def index_points(points, idx):
    """

    Input:
        points: input points data, [B, N, C]
        idx: sample index data, [B, S]
    Return:
        new_points:, indexed points data, [B, S, C]
    """
    device = points.device
    B = points.shape[0]
    view_shape = list(idx.shape)
    view_shape[1:] = [1] * (len(view_shape) - 1)
    repeat_shape = list(idx.shape)
    repeat_shape[0] = 1
    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
    new_points = points[batch_indices, idx, :]
    return new_points


def farthest_point_sample(xyz, npoint):
    """
    Input:
        xyz: pointcloud data, [B, N, 3]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [B, npoint]
    """
    device = xyz.device
    B, N, C = xyz.shape
    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10
    farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
    batch_indices = torch.arange(B, dtype=torch.long).to(device)
    for i in range(npoint):
        centroids[:, i] = farthest
        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
        dist = torch.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = torch.max(distance, -1)[1]
    return centroids


def query_ball_point(radius, nsample, xyz, new_xyz):
    """
    Input:
        radius: local region radius
        nsample: max sample number in local region
        xyz: all points, [B, N, 3]
        new_xyz: query points, [B, S, 3]
    Return:
        group_idx: grouped points index, [B, S, nsample]
    """
    device = xyz.device
    B, N, C = xyz.shape
    _, S, _ = new_xyz.shape
    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
    sqrdists = square_distance(new_xyz, xyz)
    group_idx[sqrdists > radius ** 2] = N
    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
    group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
    mask = group_idx == N
    group_idx[mask] = group_first[mask]
    return group_idx


def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False):
    """
    Input:
        npoint:
        radius:
        nsample:
        xyz: input points position data, [B, N, 3]
        points: input points data, [B, N, D]
    Return:
        new_xyz: sampled points position data, [B, npoint, nsample, 3]
        new_points: sampled points data, [B, npoint, nsample, 3+D]
    """
    B, N, C = xyz.shape
    S = npoint
    fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C]
    torch.cuda.empty_cache()
    new_xyz = index_points(xyz, fps_idx)
    torch.cuda.empty_cache()
    idx = query_ball_point(radius, nsample, xyz, new_xyz)
    torch.cuda.empty_cache()
    grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
    torch.cuda.empty_cache()
    grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
    torch.cuda.empty_cache()

    if points is not None:
        grouped_points = index_points(points, idx)
        new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
    else:
        new_points = grouped_xyz_norm
    if returnfps:
        return new_xyz, new_points, grouped_xyz, fps_idx
    else:
        return new_xyz, new_points


def sample_and_group_all(xyz, points):
    """
    Input:
        xyz: input points position data, [B, N, 3]
        points: input points data, [B, N, D]
    Return:
        new_xyz: sampled points position data, [B, 1, 3]
        new_points: sampled points data, [B, 1, N, 3+D]
    """
    device = xyz.device
    B, N, C = xyz.shape
    new_xyz = torch.zeros(B, 1, C).to(device)
    grouped_xyz = xyz.view(B, 1, N, C)
    if points is not None:
        new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
    else:
        new_points = grouped_xyz
    return new_xyz, new_points


class PointNetSetAbstraction(nn.Module):
    def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
        super(PointNetSetAbstraction, self).__init__()
        self.npoint = npoint
        self.radius = radius
        self.nsample = nsample
        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList()
        last_channel = in_channel
        for out_channel in mlp:
            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
            last_channel = out_channel
        self.group_all = group_all

    def forward(self, xyz, points):
        """
        Input:
            xyz: input points position data, [B, C, N]
            points: input points data, [B, D, N]
        Return:
            new_xyz: sampled points position data, [B, C, S]
            new_points_concat: sample points feature data, [B, D', S]
        """
        xyz = xyz.permute(0, 2, 1)
        if points is not None:
            points = points.permute(0, 2, 1)

        if self.group_all:
            new_xyz, new_points = sample_and_group_all(xyz, points)
        else:
            new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points)
        # new_xyz: sampled points position data, [B, npoint, C]
        # new_points: sampled points data, [B, npoint, nsample, C+D]
        new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            new_points =  F.relu(bn(conv(new_points)))

        new_points = torch.max(new_points, 2)[0]
        new_xyz = new_xyz.permute(0, 2, 1)
        return new_xyz, new_points


class PointNetSetAbstractionMsg(nn.Module):
    def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list):
        super(PointNetSetAbstractionMsg, self).__init__()
        self.npoint = npoint
        self.radius_list = radius_list
        self.nsample_list = nsample_list
        self.conv_blocks = nn.ModuleList()
        self.bn_blocks = nn.ModuleList()
        for i in range(len(mlp_list)):
            convs = nn.ModuleList()
            bns = nn.ModuleList()
            last_channel = in_channel + 3
            for out_channel in mlp_list[i]:
                convs.append(nn.Conv2d(last_channel, out_channel, 1))
                bns.append(nn.BatchNorm2d(out_channel))
                last_channel = out_channel
            self.conv_blocks.append(convs)
            self.bn_blocks.append(bns)

    def forward(self, xyz, points, xyz2 ):
        """
        Input:
            xyz: input points position data, [B, C, N]
            points: input points data, [B, D, N]
        Return:
            new_xyz: sampled points position data, [B, C, S]
            new_points_concat: sample points feature data, [B, D', S]
        """
        xyz = xyz.permute(0, 2, 1)
        if points is not None:
            points = points.permute(0, 2, 1)
        if xyz2 is None:
            xyz2 = xyz
        else:
            xyz2 = xyz2.permute(0, 2, 1)
        B, N, C = xyz.shape
        S = self.npoint
        new_xyz = index_points(xyz, farthest_point_sample(xyz, S))
        new_points_list = []
        for i, radius in enumerate(self.radius_list):
            K = self.nsample_list[i]
            group_idx = query_ball_point(radius, K, xyz2, new_xyz)
            grouped_xyz = index_points(xyz2, group_idx)
            grouped_xyz -= new_xyz.view(B, S, 1, C)
            if points is not None:
                grouped_points = index_points(points, group_idx)
                grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1)
            else:
                grouped_points = grouped_xyz

            grouped_points = grouped_points.permute(0, 3, 2, 1)  # [B, D, K, S]
            for j in range(len(self.conv_blocks[i])):
                conv = self.conv_blocks[i][j]
                bn = self.bn_blocks[i][j]
                grouped_points =  F.relu(bn(conv(grouped_points)))
            new_points = torch.max(grouped_points, 2)[0]  # [B, D', S]
            new_points_list.append(new_points)

        new_xyz = new_xyz.permute(0, 2, 1)
        new_points_concat = torch.cat(new_points_list, dim=1)
        return new_xyz, new_points_concat


class PointNetFeaturePropagation(nn.Module):
    def __init__(self, in_channel, mlp):
        super(PointNetFeaturePropagation, self).__init__()
        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList()
        last_channel = in_channel
        for out_channel in mlp:
            self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm1d(out_channel))
            last_channel = out_channel

    def forward(self, xyz1, xyz2, points1, points2):
        """
        Input:
            xyz1: input points position data, [B, C, N]
            xyz2: sampled input points position data, [B, C, S]
            points1: input points data, [B, D, N]
            points2: input points data, [B, D, S]
        Return:
            new_points: upsampled points data, [B, D', N]
        """
        xyz1 = xyz1.permute(0, 2, 1)
        xyz2 = xyz2.permute(0, 2, 1)

        points2 = points2.permute(0, 2, 1)
        B, N, C = xyz1.shape
        _, S, _ = xyz2.shape

        if S == 1:
            interpolated_points = points2.repeat(1, N, 1)
        else:
            dists = square_distance(xyz1, xyz2)
            dists, idx = dists.sort(dim=-1)
            dists, idx = dists[:, :, :3], idx[:, :, :3]  # [B, N, 3]

            dist_recip = 1.0 / (dists + 1e-8)
            norm = torch.sum(dist_recip, dim=2, keepdim=True)
            weight = dist_recip / norm
            interpolated_points = torch.sum(index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2)

        if points1 is not None:
            points1 = points1.permute(0, 2, 1)
            new_points = torch.cat([points1, interpolated_points], dim=-1)
        else:
            new_points = interpolated_points

        new_points = new_points.permute(0, 2, 1)
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            new_points = F.relu(bn(conv(new_points)))
        return new_points


models/pointnet_util.py
ZeyongWei/PathNet
import numpy as np

def normalize_data(batch_data):
    """ Normalize the batch data, use coordinates of the block centered at origin,
        Input:
            BxNxC array
        Output:
            BxNxC array
    """
    B, N, C = batch_data.shape
    normal_data = np.zeros((B, N, C))
    for b in range(B):
        pc = batch_data[b]
        centroid = np.mean(pc, axis=0)
        pc = pc - centroid
        m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
        pc = pc / m
        normal_data[b] = pc
    return normal_data


def shuffle_data(data, labels):
    """ Shuffle data and labels.
        Input:
          data: B,N,... numpy array
          label: B,... numpy array
        Return:
          shuffled data, label and shuffle indices
    """
    idx = np.arange(len(labels))
    np.random.shuffle(idx)
    return data[idx, ...], labels[idx], idx

def shuffle_points(batch_data):
    """ Shuffle orders of points in each point cloud -- changes FPS behavior.
        Use the same shuffling idx for the entire batch.
        Input:
            BxNxC array
        Output:
            BxNxC array
    """
    idx = np.arange(batch_data.shape[1])
    np.random.shuffle(idx)
    return batch_data[:,idx,:]

def rotate_point_cloud_y_z(batch_data,batch_data2,batch_data3,threshold = 1,threshold2 = 1):
    """ Randomly rotate the point clouds to augument the dataset
        rotation is per shape based along up direction
        Input:
          BxNx3 array, original batch of point clouds
        Return:
          BxNx3 array, rotated batch of point clouds
    """
    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
    rotated_data2 = np.zeros(batch_data2.shape, dtype=np.float32)
    if(batch_data3 != []):
        rotated_data3 = np.zeros(batch_data3.shape, dtype=np.float32)
    rotation_matrixs = np.zeros([batch_data3.shape[0],3,3], dtype=np.float32)
    rotation_matrixs_inv = np.zeros([batch_data3.shape[0],3,3], dtype=np.float32)
    for k in range(batch_data.shape[0]):
        rotation_angle = np.random.uniform() * 2 * np.pi * threshold
        cosval = np.cos(rotation_angle)
        sinval = np.sin(rotation_angle)
        rotation_matrix1 = np.array([[cosval, 0, sinval],
                                    [0, 1, 0],
                                    [-sinval, 0, cosval]])

        rotation_angle2 = np.random.uniform() * 2 * np.pi * threshold2
        cosval2 = np.cos(rotation_angle2)
        sinval2 = np.sin(rotation_angle2)
        rotation_matrix2 = np.array([[cosval2, sinval2, 0],
                                    [-sinval2, cosval2, 0],
                                    [0, 0, 1]])
        rotation_matrix = np.dot(rotation_matrix1, rotation_matrix2)
        rotation_matrix_inv = np.linalg.inv(rotation_matrix)

        shape_pc = batch_data[k, ...]
        shape_pc2 = batch_data2[k, ...]
        rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
        rotated_data2[k, ...] = np.dot(shape_pc2.reshape((-1, 3)), rotation_matrix)
        if(batch_data3 != []):
            shape_pc3 = batch_data3[k, ...]
            rotated_data3[k, ...] = np.dot(shape_pc3.reshape((-1, 3)), rotation_matrix)
        rotation_matrixs[k, ...] = rotation_matrix
        rotation_matrixs_inv[k, ...] = rotation_matrix_inv
    if(batch_data3 != []):
        return rotated_data, rotated_data2, rotated_data3, rotation_matrixs, rotation_matrixs_inv
    return rotated_data, rotated_data2, rotation_matrixs, rotation_matrixs_inv

def rotate_point_cloud(batch_data,batch_data2,batch_data3 = [],threshold = 1):
    """ Randomly rotate the point clouds to augument the dataset
        rotation is per shape based along up direction
        Input:
          BxNx3 array, original batch of point clouds
        Return:
          BxNx3 array, rotated batch of point clouds
    """
    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
    rotated_data2 = np.zeros(batch_data2.shape, dtype=np.float32)
    if(batch_data3 != []):
        rotated_data3 = np.zeros(batch_data3.shape, dtype=np.float32)
    for k in range(batch_data.shape[0]):
        rotation_angle = np.random.uniform() * 2 * np.pi * threshold
        cosval = np.cos(rotation_angle)
        sinval = np.sin(rotation_angle)
        rotation_matrix = np.array([[cosval, 0, sinval],
                                    [0, 1, 0],
                                    [-sinval, 0, cosval]])
        shape_pc = batch_data[k, ...]
        shape_pc2 = batch_data2[k, ...]
        rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
        rotated_data2[k, ...] = np.dot(shape_pc2.reshape((-1, 3)), rotation_matrix)
        if(batch_data3 != []):
            shape_pc3 = batch_data3[k, ...]
            rotated_data3[k, ...] = np.dot(shape_pc3.reshape((-1, 3)), rotation_matrix)
    if(batch_data3 != []):
        return rotated_data, rotated_data2, rotated_data3
    return rotated_data, rotated_data2

def rotate_point_cloud_z(batch_data,batch_data2,batch_data3 = [],threshold = 1):
    """ Randomly rotate the point clouds to augument the dataset
        rotation is per shape based along up direction
        Input:
          BxNx3 array, original batch of point clouds
        Return:
          BxNx3 array, rotated batch of point clouds
    """
    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
    rotated_data2 = np.zeros(batch_data2.shape, dtype=np.float32)
    if(batch_data3 != []):
        rotated_data3 = np.zeros(batch_data3.shape, dtype=np.float32)
    for k in range(batch_data.shape[0]):
        rotation_angle = np.random.uniform() * 2 * np.pi * threshold
        cosval = np.cos(rotation_angle)
        sinval = np.sin(rotation_angle)
        rotation_matrix = np.array([[cosval, sinval, 0],
                                    [-sinval, cosval, 0],
                                    [0, 0, 1]])
        shape_pc = batch_data[k, ...]
        shape_pc2 = batch_data2[k, ...]

        rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
        rotated_data2[k, ...] = np.dot(shape_pc2.reshape((-1, 3)), rotation_matrix)
        if(batch_data3 != []):
            shape_pc3 = batch_data3[k, ...]
            rotated_data3[k, ...] = np.dot(shape_pc3.reshape((-1, 3)), rotation_matrix)
    if(batch_data3 != []):
        return rotated_data, rotated_data2, rotated_data3
    return rotated_data, rotated_data2

def rotate_point_cloud_with_normal(batch_xyz_normal):
    ''' Randomly rotate XYZ, normal point cloud.
        Input:
            batch_xyz_normal: B,N,6, first three channels are XYZ, last 3 all normal
        Output:
            B,N,6, rotated XYZ, normal point cloud
    '''
    for k in range(batch_xyz_normal.shape[0]):
        rotation_angle = np.random.uniform() * 2 * np.pi
        cosval = np.cos(rotation_angle)
        sinval = np.sin(rotation_angle)
        rotation_matrix = np.array([[cosval, 0, sinval],
                                    [0, 1, 0],
                                    [-sinval, 0, cosval]])
        shape_pc = batch_xyz_normal[k,:,0:3]
        shape_normal = batch_xyz_normal[k,:,3:6]
        batch_xyz_normal[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
        batch_xyz_normal[k,:,3:6] = np.dot(shape_normal.reshape((-1, 3)), rotation_matrix)
    return batch_xyz_normal

def rotate_perturbation_point_cloud_with_normal(batch_data, angle_sigma=0.06, angle_clip=0.18):
    """ Randomly perturb the point clouds by small rotations
        Input:
          BxNx6 array, original batch of point clouds and point normals
        Return:
          BxNx3 array, rotated batch of point clouds
    """
    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
    for k in range(batch_data.shape[0]):
        angles = np.clip(angle_sigma*np.random.randn(3), -angle_clip, angle_clip)
        Rx = np.array([[1,0,0],
                       [0,np.cos(angles[0]),-np.sin(angles[0])],
                       [0,np.sin(angles[0]),np.cos(angles[0])]])
        Ry = np.array([[np.cos(angles[1]),0,np.sin(angles[1])],
                       [0,1,0],
                       [-np.sin(angles[1]),0,np.cos(angles[1])]])
        Rz = np.array([[np.cos(angles[2]),-np.sin(angles[2]),0],
                       [np.sin(angles[2]),np.cos(angles[2]),0],
                       [0,0,1]])
        R = np.dot(Rz, np.dot(Ry,Rx))
        shape_pc = batch_data[k,:,0:3]
        shape_normal = batch_data[k,:,3:6]
        rotated_data[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), R)
        rotated_data[k,:,3:6] = np.dot(shape_normal.reshape((-1, 3)), R)
    return rotated_data


def rotate_point_cloud_by_angle(batch_data, rotation_angle):
    """ Rotate the point cloud along up direction with certain angle.
        Input:
          BxNx3 array, original batch of point clouds
        Return:
          BxNx3 array, rotated batch of point clouds
    """
    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
    for k in range(batch_data.shape[0]):
        #rotation_angle = np.random.uniform() * 2 * np.pi
        cosval = np.cos(rotation_angle)
        sinval = np.sin(rotation_angle)
        rotation_matrix = np.array([[cosval, 0, sinval],
                                    [0, 1, 0],
                                    [-sinval, 0, cosval]])
        shape_pc = batch_data[k,:,0:3]
        rotated_data[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
    return rotated_data

def rotate_point_cloud_by_angle_with_normal(batch_data, rotation_angle):
    """ Rotate the point cloud along up direction with certain angle.
        Input:
          BxNx6 array, original batch of point clouds with normal
          scalar, angle of rotation
        Return:
          BxNx6 array, rotated batch of point clouds iwth normal
    """
    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
    for k in range(batch_data.shape[0]):
        #rotation_angle = np.random.uniform() * 2 * np.pi
        cosval = np.cos(rotation_angle)
        sinval = np.sin(rotation_angle)
        rotation_matrix = np.array([[cosval, 0, sinval],
                                    [0, 1, 0],
                                    [-sinval, 0, cosval]])
        shape_pc = batch_data[k,:,0:3]
        shape_normal = batch_data[k,:,3:6]
        rotated_data[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
        rotated_data[k,:,3:6] = np.dot(shape_normal.reshape((-1,3)), rotation_matrix)
    return rotated_data



def rotate_perturbation_point_cloud(batch_data, angle_sigma=0.06, angle_clip=0.18):
    """ Randomly perturb the point clouds by small rotations
        Input:
          BxNx3 array, original batch of point clouds
        Return:
          BxNx3 array, rotated batch of point clouds
    """
    rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
    for k in range(batch_data.shape[0]):
        angles = np.clip(angle_sigma*np.random.randn(3), -angle_clip, angle_clip)
        Rx = np.array([[1,0,0],
                       [0,np.cos(angles[0]),-np.sin(angles[0])],
                       [0,np.sin(angles[0]),np.cos(angles[0])]])
        Ry = np.array([[np.cos(angles[1]),0,np.sin(angles[1])],
                       [0,1,0],
                       [-np.sin(angles[1]),0,np.cos(angles[1])]])
        Rz = np.array([[np.cos(angles[2]),-np.sin(angles[2]),0],
                       [np.sin(angles[2]),np.cos(angles[2]),0],
                       [0,0,1]])
        R = np.dot(Rz, np.dot(Ry,Rx))
        shape_pc = batch_data[k, ...]
        rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), R)
    return rotated_data


def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05):
    """ Randomly jitter points. jittering is per point.
        Input:
          BxNx3 array, original batch of point clouds
        Return:
          BxNx3 array, jittered batch of point clouds
    """
    B, N, C = batch_data.shape
    assert(clip > 0)
    jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1*clip, clip)
    jittered_data += batch_data
    return jittered_data

def shift_point_cloud(batch_data, batch_data2, shift_range=0.1):
    """ Randomly shift point cloud. Shift is per point cloud.
        Input:
          BxNx3 array, original batch of point clouds
        Return:
          BxNx3 array, shifted batch of point clouds
    """
    B, N, C = batch_data.shape
    shifts = np.random.uniform(-shift_range, shift_range, (B,3))
    for batch_index in range(B):
        batch_data[batch_index,:,:] += shifts[batch_index,:]
        batch_data2[batch_index,:,:] += shifts[batch_index,:]
    return batch_data,batch_data2


def random_scale_point_cloud(batch_data, batch_data2, scale_low=0.8, scale_high=1.25):
    """ Randomly scale the point cloud. Scale is per point cloud.
        Input:
            BxNx3 array, original batch of point clouds
        Return:
            BxNx3 array, scaled batch of point clouds
    """
    B, N, C = batch_data.shape
    scales = np.random.uniform(scale_low, scale_high, B)
    for batch_index in range(B):
        batch_data[batch_index,:,:] *= scales[batch_index]
        batch_data2[batch_index,:,:] *= scales[batch_index]
    return batch_data,batch_data2

def random_point_dropout(batch_pc, max_dropout_ratio=0.5):
    ''' batch_pc: BxNx3 '''
    for b in range(batch_pc.shape[0]):
        dropout_ratio =  np.random.random()*max_dropout_ratio # 0~0.875
        drop_idx = np.where(np.random.random((batch_pc.shape[1]))<=dropout_ratio)[0]
        if len(drop_idx)>0:
            batch_pc[b,drop_idx,:] = batch_pc[b,0,:] # set to the first point
    return batch_pc

def random_point_dropout2(batch_pc, max_dropout_ratio=0.5):
    ''' batch_pc: BxNx3 '''
    for b in range(batch_pc.shape[0]):
        dropout_ratio =  np.random.random()*max_dropout_ratio # 0~0.875
        keep_idx = np.where(np.random.random((batch_pc.shape[1]))>=dropout_ratio)[0]
        if len(keep_idx)>4:
            batch_pc = batch_pc[b,keep_idx,:].reshape(1,-1,3)
    return batch_pc




provider.py
ZeyongWei/PathNet
"""
Author: Wei
Date: Mar 2022
"""
import argparse
import os
from data_utils.DataLoader import PatchDataset
import torch
import datetime
import logging
from pathlib import Path
import sys
import importlib
import shutil
from tqdm import tqdm
import provider
import numpy as np
from scipy import spatial
from glob import glob
from data_utils.Pointfilter_Utils import pca_alignment
import time
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = BASE_DIR
sys.path.append(os.path.join(ROOT_DIR, 'models'))


def parse_args():
    parser = argparse.ArgumentParser('Model')
    parser.add_argument('--model', type=str, default='model', help='model name [default: model]')
    parser.add_argument('--batch_size', type=int, default=100, help='Batch Size during training [default: 16]')
    parser.add_argument('--epoch',  default=1000, type=int, help='Epoch to run [default: 251]')
    parser.add_argument('--learning_rate', default=0.000001, type=float, help='Initial learning rate [default: 0.001]')
    parser.add_argument('--gpu', type=str, default='0', help='GPU to use [default: GPU 0]')
    parser.add_argument('--optimizer', type=str, default='Adam', help='Adam or SGD [default: Adam]')
    parser.add_argument('--optimizer2', type=str, default='Adam', help='Adam or SGD [default: Adam]')
    parser.add_argument('--log_dir', type=str, default='model', help='Log path [default: None]')
    parser.add_argument('--decay_rate', type=float, default=1e-4, help='weight decay [default: 1e-4]')
    parser.add_argument('--npoint', type=int,  default=128, help='Point Number [default: 2048]')
    parser.add_argument('--normal', action='store_true', default=False, help='Whether to use normal information [default: False]')
    parser.add_argument('--step_size', type=int,  default=20, help='Decay step for lr decay [default: every 20 epochs]')
    parser.add_argument('--lr_decay', type=float,  default=0.5, help='Decay rate for lr decay [default: 0.5]')
    parser.add_argument('--use_random_path', type=int, default=0, help='wether use random path 0 no ,1 yes, 2 all 0, 3 all 1 ')
    parser.add_argument('--block_num', type=int, default=6, help='num of denosier block')
    parser.add_argument('--path_num', type=int, default=2, help='path num of each denosier block')

    return parser.parse_args()

def pca_normalize(pc):
    normalize_data = np.zeros(pc.shape, dtype=np.float32)
    martix_inv = np.zeros((pc.shape[0],3,3),dtype=np.float32)

    centroid = pc[:,0,:]
    centroid = np.expand_dims(centroid, axis=1)
    pc = pc - centroid

    m = np.max(np.sqrt(np.sum(pc**2, axis=2)),axis = 1, keepdims=True)
    scale_inv = m #B
    pc = pc / np.expand_dims(m, axis=-1)

    for B in range(pc.shape[0]):
        x, pca_martix_inv = pca_alignment(pc[B,:,:])
        normalize_data[B, ...] = x
        martix_inv[B, ...] = pca_martix_inv

    return normalize_data, martix_inv, scale_inv

def main(args,test_data_dir):
    def log_string(str):
        logger.info(str)
        print(str)

    '''HYPER PARAMETER'''
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

    '''CREATE DIR'''
    experiment_dir = Path('./log/')
    experiment_dir = experiment_dir.joinpath('path_denoise/'+ args.model)
    checkpoints_dir = experiment_dir.joinpath('checkpoints/')

    '''LOG'''
    args = parse_args()
    name_dir = test_data_dir 
    root = 'data/test_data/'

    DATA_PATH = 'data/test_data/' + name_dir + '/' 
    samples = glob(DATA_PATH+"/*.xyz")  
    samples.sort()
    #print(samples)
    batch_size = args.batch_size

    block_num = args.block_num
    path_num = args.path_num

    '''MODEL LOADING'''
    MODEL = importlib.import_module(args.model)

    denoiser = MODEL.get_model(block_num, path_num).cuda()
    analyser = MODEL.get_analyser(block_num, path_num).cuda()
 
    checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth')
    denoiser.load_state_dict(checkpoint['denoiser_model_state_dict'])
    analyser.load_state_dict(checkpoint['analyser_model_state_dict'])


    '''save dir'''
    save_dir = './data/test_data/results/' + name_dir
    if(not os.path.exists(save_dir)):
       os.makedirs(save_dir)

    for i,item in tqdm(enumerate(samples)):
        print(item)
        start_time = time.time()
        data_name = item.split('/')[-1][:-4]
        input_data = np.loadtxt(DATA_PATH + data_name + '.xyz').astype(np.float32) #

        with torch.no_grad():
            denoiser = denoiser.eval()
            analyser = analyser.eval()
            source_data = input_data[:,:3]

            path_first = None

            for iter_time in range(2):

                inputs = input_data[:,:3]
                nbrs = spatial.cKDTree(inputs)  # kd tree

                batch_num = int(inputs.shape[0] / batch_size)
                add_batch = 0 if(inputs.shape[0] % batch_size ==0)else 1

                normalize_trans_all = []
                path_all = []
                
                dist,idxs = nbrs.query(inputs,k = int(128)) #s

                input_patchs = inputs[idxs,:] # B,128,3

                normalize_input_patchs, martix_inv, scale_inv = pca_normalize(input_patchs)  # B,K,3  #B,3,3  #B
                print('time1:', time.time()-start_time)
                for i in tqdm(range(batch_num+add_batch)):
                    points = normalize_input_patchs[i*batch_size:(i+1)*batch_size,:,:] # B,128,3
                
                    points = torch.Tensor(points)
                    points = points.float().cuda()
      
                    points = points.transpose(2, 1)

                    trans_m, path_m, _ = denoiser(points,analyser,args.use_random_path)

                    trans_m = torch.cat(trans_m,axis = 0).reshape(block_num,-1,3).transpose(1,0) #B,block_num,3

                    path_m = torch.cat(path_m,axis = 0).reshape(block_num,-1).transpose(1,0)  #B,block_num

                    normalize_trans_all.append(trans_m)
                    path_all.append(path_m)

                normalize_trans_all = torch.cat(normalize_trans_all, axis = 0)  # B,6,3
                path_all = torch.cat(path_all, axis = 0)

                normalize_trans_all = normalize_trans_all.cpu().numpy().astype(np.float32).reshape(-1,block_num,3)
                path_all = path_all.cpu().numpy().astype(np.float32).reshape(-1,block_num)

                trans_all = np.matmul(martix_inv, normalize_trans_all.transpose(0,2,1)).transpose(0,2,1) # B,6,3

                trans_all = trans_all * np.expand_dims(scale_inv,axis = -1)

                outputs_all = inputs[:,:3].reshape(-1,1,3) - trans_all
            
                #scale_recover
                if(iter_time ==0):
                    path_first = path_all

                #save
                inputs_start = np.concatenate((inputs[:,:3],path_all[:,0].reshape(-1,1)),axis = -1)
                np.savetxt(save_dir+"/"+ data_name +"_input_start.xyz", inputs_start.astype(np.float32), fmt = '%.6f')

                inputs_sum = np.concatenate((inputs[:,:3],np.sum(path_all, axis = 1).reshape(-1,1)),axis = -1)
                np.savetxt(save_dir+"/"+ data_name +"_input_sum.xyz",inputs_sum.astype(np.float32),fmt = '%.6f')

                output_end = np.concatenate((outputs_all[:,-1,:],np.sum(path_first, axis = 1).reshape(-1,1)),axis = -1)
                np.savetxt(save_dir+"/"+ data_name +"_output_end.xyz",output_end.astype(np.float32),fmt = '%.6f')

                for i in range(0,block_num-1):
                    path_i = path_all[:,i+1].reshape(-1,1) 
                    output_i = outputs_all[:,i,:]
                    output_i = np.concatenate((output_i,path_i),axis = -1)
                    np.savetxt(save_dir+"/"+ data_name +"_output_" + str(i) + ".xyz",output_i.astype(np.float32),fmt = '%.6f')

                input_data = outputs_all[:,-1,:] #iter
                data_name = data_name +'_output_end'
            print('time2:', time.time()-start_time)
if __name__ == '__main__':
    args = parse_args()
    test_data_dir =  'benchmark81_20000' # test dir ['benchmark81_10000','benchmark81_20000','benchmark81_50000','kinect_fusion','kinect_v1','kinect_v2']
    main(args,test_data_dir)


test.py
ZeyongWei/PathNet