338 lines
9.6 KiB
Python
338 lines
9.6 KiB
Python
import sys
|
|
import os
|
|
import shutil
|
|
import math
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../') # add project root directory
|
|
|
|
from dataset.shapenet import ShapeNetPartSegDataset
|
|
from model.pointnet2_part_seg import PointNet2PartSegmentNet
|
|
import torch_geometric.transforms as GT
|
|
import torch
|
|
import argparse
|
|
import numpy as np
|
|
from sklearn.cluster import DBSCAN
|
|
from sklearn.preprocessing import StandardScaler
|
|
import open3d as o3d
|
|
import pointcloud as pc
|
|
|
|
def eval_sample(net, sample):
|
|
'''
|
|
sample: { 'points': tensor(n, 3), 'labels': tensor(n,) }
|
|
return: (pred_label, gt_label) with labels shape (n,)
|
|
'''
|
|
net.eval()
|
|
with torch.no_grad():
|
|
# points: (n, 3)
|
|
points, gt_label = sample['points'], sample['labels']
|
|
n = points.shape[0]
|
|
|
|
points = points.view(1, n, 3) # make a batch
|
|
points = points.transpose(1, 2).contiguous()
|
|
points = points.to(device, dtype)
|
|
|
|
pred = net(points) # (batch_size, n, num_classes)
|
|
pred_label = pred.max(2)[1]
|
|
pred_label = pred_label.view(-1).cpu() # (n,)
|
|
|
|
assert pred_label.shape == gt_label.shape
|
|
return (pred_label, gt_label)
|
|
|
|
|
|
def mini_color_table(index, norm=True):
|
|
colors = [
|
|
[0.5000, 0.5400, 0.5300], [0.8900, 0.1500, 0.2100], [0.6400, 0.5800, 0.5000],
|
|
[1.0000, 0.3800, 0.0100], [1.0000, 0.6600, 0.1400], [0.4980, 1.0000, 0.0000],
|
|
[0.4980, 1.0000, 0.8314], [0.9412, 0.9725, 1.0000], [0.5412, 0.1686, 0.8863],
|
|
[0.5765, 0.4392, 0.8588], [0.3600, 0.1400, 0.4300], [0.5600, 0.3700, 0.6000],
|
|
]
|
|
|
|
color = colors[index % len(colors)]
|
|
|
|
if not norm:
|
|
color[0] *= 255
|
|
color[1] *= 255
|
|
color[2] *= 255
|
|
|
|
return color
|
|
|
|
|
|
def label2color(labels):
|
|
'''
|
|
labels: np.ndarray with shape (n, )
|
|
colors(return): np.ndarray with shape (n, 3)
|
|
'''
|
|
num = labels.shape[0]
|
|
colors = np.zeros((num, 3))
|
|
|
|
minl, maxl = np.min(labels), np.max(labels)
|
|
for l in range(minl, maxl + 1):
|
|
colors[labels == l, :] = mini_color_table(l)
|
|
|
|
return colors
|
|
|
|
|
|
def clusterToColor(cluster, cluster_idx):
|
|
|
|
colors = np.zeros(shape=(len(cluster), 3))
|
|
point_idx = 0
|
|
for point in cluster:
|
|
colors[point_idx, :] = mini_color_table(cluster_idx)
|
|
point_idx += 1
|
|
|
|
return colors
|
|
|
|
|
|
def normalize_pointcloud(pc):
|
|
|
|
max = pc.max(axis=0)
|
|
min = pc.min(axis=0)
|
|
|
|
f = np.max([abs(max[0] - min[0]), abs(max[1] - min[1]), abs(max[2] - min[2])])
|
|
|
|
pc[:, 0:3] /= f
|
|
pc[:, 3:6] /= (np.linalg.norm(pc[:, 3:6], ord=2, axis=1, keepdims=True))
|
|
|
|
return pc
|
|
|
|
|
|
def farthest_point_sampling(pts, K):
|
|
|
|
if pts.shape[0] < K:
|
|
return pts
|
|
|
|
def calc_distances(p0, points):
|
|
return ((p0[:3] - points[:, :3]) ** 2).sum(axis=1)
|
|
|
|
farthest_pts = np.zeros((K, pts.shape[1]))
|
|
farthest_pts[0] = pts[np.random.randint(len(pts))]
|
|
distances = calc_distances(farthest_pts[0], pts)
|
|
for i in range(1, K):
|
|
farthest_pts[i] = pts[np.argmax(distances)]
|
|
distances = np.minimum(distances, calc_distances(farthest_pts[i], pts))
|
|
|
|
return farthest_pts
|
|
|
|
|
|
def append_onehotencoded_type(data, factor = 1.0):
|
|
|
|
types = data[:, 6].astype(int)
|
|
res = np.zeros((len(types), 4))
|
|
res[np.arange(len(types)), types] = factor
|
|
|
|
return np.column_stack((data, res))
|
|
|
|
|
|
def append_normal_angles(data):
|
|
|
|
def func(x):
|
|
theta = math.acos(x[2]) / math.pi
|
|
phi = (math.atan2(x[1], x[0]) + math.pi) / (2.0 * math.pi)
|
|
return (theta, phi)
|
|
|
|
res = np.array([func(xi) for xi in data[:, 3:6]])
|
|
|
|
print(res)
|
|
|
|
return np.column_stack((data, res))
|
|
|
|
|
|
def extract_cube_clusters(data, cluster_dims, max_points_per_cluster):
|
|
|
|
max = data[:,:3].max(axis=0)
|
|
max += max * 0.01
|
|
|
|
min = data[:,:3].min(axis=0)
|
|
min -= min * 0.01
|
|
|
|
size = (max - min)
|
|
|
|
clusters = {}
|
|
|
|
cluster_size = size / np.array(cluster_dims, dtype=np.float32)
|
|
|
|
print('Min: ' + str(min) + ' Max: ' + str(max))
|
|
print('Cluster Size: ' + str(cluster_size))
|
|
|
|
for row in data:
|
|
|
|
# print('Row: ' + str(row))
|
|
|
|
cluster_pos = ((row[:3] - min) / cluster_size).astype(int)
|
|
cluster_idx = cluster_dims[0] * cluster_dims[2] * cluster_pos[1] + cluster_dims[0] * cluster_pos[2] + cluster_pos[0]
|
|
clusters.setdefault(cluster_idx, []).append(row)
|
|
|
|
# Apply farthest point sampling to each cluster
|
|
for key, cluster in clusters.items():
|
|
c = np.vstack(cluster)
|
|
clusters[key] = farthest_point_sampling(c, max_points_per_cluster)
|
|
|
|
return clusters.values()
|
|
|
|
|
|
def extract_clusters(data, selected_indices, eps, min_samples, metric='euclidean', algo='auto'):
|
|
|
|
min_samples = min_samples * len(data);
|
|
|
|
print('Clustering. Min Samples: ' + str(min_samples) + ' EPS: ' + str(eps))
|
|
|
|
# 0,1,2 : pos
|
|
# 3,4,5 : normal
|
|
# 6: type index
|
|
# 7,8,9,10: type index one hot encoded
|
|
# 11,12: normal as angles
|
|
|
|
db_res = DBSCAN(eps=eps, metric=metric, n_jobs=-1, algorithm=algo, min_samples=min_samples).fit(data[:, selected_indices])
|
|
|
|
|
|
labels = db_res.labels_
|
|
n_clusters = len(set(labels)) - (1 if -1 in labels else 0)
|
|
n_noise = list(labels).count(-1)
|
|
print("Noise: " + str(n_noise) + " Clusters: " + str(n_clusters))
|
|
|
|
clusters = {}
|
|
for idx, l in enumerate(labels):
|
|
if l is -1:
|
|
continue
|
|
clusters.setdefault(str(l), []).append(data[idx, :])
|
|
|
|
|
|
npClusters = []
|
|
for cluster in clusters.values():
|
|
npClusters.append(np.array(cluster))
|
|
|
|
return npClusters
|
|
|
|
|
|
def draw_clusters(clusters):
|
|
|
|
clouds = []
|
|
|
|
cluster_idx = 0
|
|
for cluster in clusters:
|
|
|
|
cloud = o3d.PointCloud()
|
|
cloud.points = o3d.Vector3dVector(cluster[:,:3])
|
|
cloud.colors = o3d.Vector3dVector(clusterToColor(cluster, cluster_idx))
|
|
clouds.append(cloud)
|
|
cluster_idx += 1
|
|
|
|
o3d.draw_geometries(clouds)
|
|
|
|
|
|
def draw_sample_data(sample_data, colored_normals = False):
|
|
|
|
cloud = o3d.PointCloud()
|
|
cloud.points = o3d.Vector3dVector(sample_data[:,:3])
|
|
cloud.colors = \
|
|
o3d.Vector3dVector(label2color(sample_data[:, 6].astype(int)) if not colored_normals else sample_data[:, 3:6])
|
|
|
|
o3d.draw_geometries([cloud])
|
|
|
|
|
|
def recreate_folder(folder):
|
|
if os.path.exists(folder) and os.path.isdir(folder):
|
|
shutil.rmtree(folder)
|
|
os.mkdir(folder)
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../') # add project root directory
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--dataset', type=str, default='data', help='dataset path')
|
|
parser.add_argument('--npoints', type=int, default=2048, help='resample points number')
|
|
parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_246.pth', help='model path')
|
|
parser.add_argument('--sample_idx', type=int, default=0, help='select a sample to segment and view result')
|
|
opt = parser.parse_args()
|
|
print(opt)
|
|
|
|
if __name__ == '__main__':
|
|
|
|
# Create dataset
|
|
print('Create data set ..')
|
|
|
|
dataset_folder = './data/raw/predict/'
|
|
pointcloud_file = './pointclouds/0_pc.xyz'
|
|
|
|
pc = pc.read_pointcloud(pointcloud_file)
|
|
pc = normalize_pointcloud(pc)
|
|
pc = append_normal_angles(pc)
|
|
|
|
# pc = StandardScaler().fit_transform(pc)
|
|
|
|
recreate_folder(dataset_folder)
|
|
|
|
# Add full point cloud to prediction folder.
|
|
recreate_folder(dataset_folder + '0_0' + '/')
|
|
pc_fps = farthest_point_sampling(pc, opt.npoints)
|
|
pc.write_pointcloud(dataset_folder + '0_0' + '/pc.xyz', pc_fps)
|
|
|
|
pc_clusters = extract_cube_clusters(pc, [4,4,4], 1024)
|
|
#pc_clusters = extract_clusters(pc, [0, 1, 2, 3, 4, 5], eps=0.1, min_samples=0.0001, metric='euclidean', algo='auto')
|
|
# Add cluster point clouds to prediction folder.
|
|
for idx, pcc in enumerate(pc_clusters):
|
|
|
|
pcc = farthest_point_sampling(pcc, opt.npoints)
|
|
recreate_folder(dataset_folder + str(idx) + '/')
|
|
pc.write_pointcloud(dataset_folder + str(idx) + '/pc.xyz', pcc)
|
|
#draw_sample_data(pcc, False)
|
|
|
|
draw_clusters(pc_clusters)
|
|
|
|
# Load dataset
|
|
print('load dataset ..')
|
|
test_transform = GT.Compose([GT.NormalizeScale(), ])
|
|
|
|
test_dataset = ShapeNetPartSegDataset(
|
|
mode='predict',
|
|
root_dir=opt.dataset,
|
|
transform=None,
|
|
npoints=opt.npoints,
|
|
refresh=False
|
|
)
|
|
num_classes = test_dataset.num_classes()
|
|
|
|
print('test dataset size: ', len(test_dataset))
|
|
|
|
# Load model
|
|
print('Construct model ..')
|
|
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
|
|
dtype = torch.float
|
|
|
|
# net = PointNetPartSegmentNet(num_classes)
|
|
net = PointNet2PartSegmentNet(num_classes)
|
|
|
|
net.load_state_dict(torch.load(opt.model, map_location=device.type))
|
|
net = net.to(device, dtype)
|
|
net.eval()
|
|
|
|
result_clusters = []
|
|
|
|
# Iterate over all the samples
|
|
for sample in test_dataset:
|
|
|
|
print('Eval test sample ..')
|
|
pred_label, gt_label = eval_sample(net, sample)
|
|
sample_data = np.column_stack((sample["points"].numpy(), sample["normals"].numpy(), pred_label.numpy()))
|
|
print('Eval done.')
|
|
|
|
sample_data = normalize_pointcloud(sample_data)
|
|
|
|
sample_data = append_onehotencoded_type(sample_data, 1.0)
|
|
sample_data = append_normal_angles(sample_data)
|
|
|
|
print('Clustering ..')
|
|
print('Shape: ' + str(sample_data.shape))
|
|
|
|
clusters = extract_clusters(sample_data, [0, 1, 2, 3, 4, 5, 7, 8, 9, 10], eps=0.1, min_samples=0.0001, metric='euclidean', algo='auto')
|
|
|
|
print('Clustering done. ' + str(len(clusters)) + " Clusters.")
|
|
print(sample_data[:, 6])
|
|
|
|
draw_sample_data(sample_data, False)
|
|
|
|
result_clusters.extend(clusters)
|
|
|
|
# result_clusters.append(sample_data)
|
|
|
|
#draw_clusters(result_clusters)
|