|
- # import torch
- # import torch.distributions as dist
- import tensorflow as tf
- import os
- import shutil
- import argparse
- import time
- import trimesh
- import random
- from collections import defaultdict
- import fd.config
- import fn.config
- import fn.checkpoints
- import fd.checkpoints
- from generation import Generator3D6
- import numpy as np
-
- from tqdm import tqdm
-
- def farthest_point_sample(xyz, pointnumber):
- #device ='cuda'
- N, C = xyz.shape
- #torch.seed()
- np.random.seed(1234)
- #xyz=tf.convert_to_tensor(xyz)
- centroids = np.zeros(pointnumber, dtype=np.int64)
-
- distance = np.ones(N) * 1e32
- farthest = np.random.randint(0, N, (1,))
- farthest[0]=N/2
- #farthest = tf.convert_to_tensor(farthest, dtype=tf.64)
- for i in tqdm(range(pointnumber)):
- centroids[i] = farthest
- centroid = xyz[farthest, :]
-
- #print(xyz.shape)
- dist = np.sum((xyz - centroid) ** 2, axis=-1)
-
- mask = dist < distance
- distance[mask] = dist[mask]
-
- farthest = np.argmax(distance, axis=-1)
- return centroids.astype(np.int64)
-
- tpointnumber=8192
-
- datalist=['test/cow.xyz',
- 'test/coverrear_Lp.xyz',
- 'test/chair.xyz',
- 'test/camel.xyz',
- 'test/casting.xyz',
- 'test/duck.xyz',
- 'test/eight.xyz',
- 'test/elephant.xyz',
- 'test/elk.xyz',
- 'test/fandisk.xyz',
- 'test/genus3.xyz',
- 'test/horse.xyz',
- 'test/Icosahedron.xyz',
- 'test/kitten.xyz',
- 'test/moai.xyz',
- 'test/Octahedron.xyz',
- 'test/pig.xyz',
- 'test/quadric.xyz',
- 'test/sculpt.xyz',
- 'test/star.xyz']
-
- outlist=['test_out/cow.xyz',
- 'test_out/coverrear_Lp.xyz',
- 'test_out/chair.xyz',
- 'test_out/camel.xyz',
- 'test_out/casting.xyz',
- 'test_out/duck.xyz',
- 'test_out/eight.xyz',
- 'test_out/elephant.xyz',
- 'test_out/elk.xyz',
- 'test_out/fandisk.xyz',
- 'test_out/genus3.xyz',
- 'test_out/horse.xyz',
- 'test_out/Icosahedron.xyz',
- 'test_out/kitten.xyz',
- 'test_out/moai.xyz',
- 'test_out/Octahedron.xyz',
- 'test_out/pig.xyz',
- 'test_out/quadric.xyz',
- 'test_out/sculpt.xyz',
- 'test_out/star.xyz']
-
-
- cfg1 = fn.config.load_config('configs/fn.yaml')
- cfg2 = fd.config.load_config('configs/fd.yaml')
-
- model = fn.config.get_model(cfg1)
- model2 = fd.config.get_model(cfg2)
-
- #checkpoint_io1 = fn.checkpoints.CheckpointIO('out/fn/results', model=model)
- #load_dict =checkpoint_io1.load( 'model_best.h5')
- model.build(input_shape=(4, 16, 100, 3))
- model.load_weights('out/fn/model_best.h5')
-
- #checkpoint_io2 = fd.checkpoints.CheckpointIO('out/fd/results', model=model2)
- #load_dict =checkpoint_io2.load( 'out/fd/model_best.h5')
- model2.build(input_shape=(4, 16, 100, 3))
- model2.load_weights('out/fd/model_best.h5')
-
- generator=Generator3D6(model, model2)
-
-
- for k in range(len(datalist)):
- print("processing "+datalist[k])
- #normalization
- xyzname=datalist[k]
- cloud =np.loadtxt(xyzname)
- cloud=cloud[:,0:3]
- bbox=np.zeros((2,3))
- bbox[0][0]=np.min(cloud[:,0])
- bbox[0][1]=np.min(cloud[:,1])
- bbox[0][2]=np.min(cloud[:,2])
- bbox[1][0]=np.max(cloud[:,0])
- bbox[1][1]=np.max(cloud[:,1])
- bbox[1][2]=np.max(cloud[:,2])
- loc = (bbox[0] + bbox[1]) / 2
- scale = (bbox[1] - bbox[0]).max()
- scale1 = 1/scale
- for i in range(cloud.shape[0]):
- cloud[i]=cloud[i]-loc
- cloud[i]=cloud[i]*scale1
- np.savetxt("test.xyz",cloud)
- cloud=np.expand_dims(cloud,0)
-
- #upsampling
- pointcloud = np.array(generator.upsample(cloud))
-
- #farthest_point_sample
- print("farthest point sample")
- for i in range(pointcloud.shape[0]):
- pointcloud[i]=pointcloud[i]*scale
- pointcloud[i]=pointcloud[i]+loc
-
- centroids=farthest_point_sample(pointcloud, tpointnumber)
- np.savetxt(outlist[k],pointcloud[centroids])
- print("done")
|