From 1f7bfe77657a4876a0d14c55f8a614506677317e Mon Sep 17 00:00:00 2001 From: Markus Friedrich Date: Tue, 9 Jun 2020 10:14:30 +0200 Subject: [PATCH] added type detection pipeline --- main_pipeline.py | 94 +++++++++++++++ requirements.txt | 91 ++++++++++++++ utils/pointcloud.py | 282 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 467 insertions(+) create mode 100644 main_pipeline.py create mode 100644 requirements.txt create mode 100644 utils/pointcloud.py diff --git a/main_pipeline.py b/main_pipeline.py new file mode 100644 index 0000000..89c3cbf --- /dev/null +++ b/main_pipeline.py @@ -0,0 +1,94 @@ +from pathlib import Path + +import torch +from torch_geometric.data import Data +from tqdm import tqdm + +import polyscope as ps +import numpy as np + +from torch.utils.data import DataLoader + +# Dataset and Dataloaders +# ============================================================================= + +# Transforms +from ml_lib.point_toolset.point_io import BatchToData +from ml_lib.utils.model_io import SavedLightningModels + + +# Datasets +from datasets.full_pointclouds import FullCloudsDataset +from utils.pointcloud import read_pointcloud, normalize_pointcloud, cluster_cubes, append_onehotencoded_type, \ + label2color, polytopes_to_planes +from utils.project_config import GlobalVar, ThisConfig + + +def prepare_dataloader(config_obj): + dataset = FullCloudsDataset(config_obj.data.root, split=GlobalVar.data_split.test, + setting=GlobalVar.settings[config_obj.model.type]) + # noinspection PyTypeChecker + return DataLoader(dataset, batch_size=config_obj.train.batch_size, + num_workers=config_obj.data.worker, shuffle=False) + + +def restore_logger_and_model(log_dir): + model = SavedLightningModels.load_checkpoint(models_root_path=log_dir, n=-1) + model = model.restore() + if torch.cuda.is_available(): + model.cuda() + else: + model.cpu() + return model + +def predict_prim_type(input_pc, model): + + input_data = ( + torch.tensor(np.array([input_pc[:, 3:6]], np.float)), + torch.tensor(input_pc[:, 0:3]), + np.zeros(input_pc.shape[0]) + ) + + batch_to_data = BatchToData() + + data = batch_to_data(input_data[0], input_data[1], input_data[2]) + y = loaded_model(data.to(device='cuda' if torch.cuda.is_available() else 'cpu')) + y_primary = torch.argmax(y.main_out, dim=-1).squeeze().cpu().numpy() + + return np.concatenate((input_pc, y_primary.reshape(-1,1)), axis=1) + +if __name__ == '__main__': + + input_pc_path = 'data/pc/pc.txt' + + model_path = Path('trained_models/version_1') + config_filename = 'config.ini' + config = ThisConfig() + config.read_file((Path(model_path) / config_filename).open('r')) + loaded_model = restore_logger_and_model(model_path) + loaded_model.eval() + + input_pc = read_pointcloud(input_pc_path, ' ', False) + + input_pc = normalize_pointcloud(input_pc) + + grid_clusters = cluster_cubes(input_pc, [1,1,1], 2048) + + ps.init() + + for i,grid_cluster_pc in enumerate(grid_clusters): + + print("Cluster pointcloud size: {}".format(grid_cluster_pc.shape[0])) + + pc_with_prim_type = predict_prim_type(grid_cluster_pc, loaded_model) + + #pc_with_prim_type = polytopes_to_planes(pc_with_prim_type) + + pc_with_prim_type = append_onehotencoded_type(pc_with_prim_type) + + pc = ps.register_point_cloud("points_" + str(i), pc_with_prim_type[:, :3], radius=0.01) + pc.add_color_quantity("prim types", label2color(pc_with_prim_type[:, 6].astype(np.int64)), True) + + ps.show() + + print('Done') diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..0cc82f9 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,91 @@ +absl-py==0.9.0 +attrs==19.3.0 +bravado==10.6.2 +bravado-core==5.17.0 +cachetools==4.1.0 +certifi==2020.4.5.2 +chardet==3.0.4 +click==7.1.2 +combo==0.1.0 +cycler==0.10.0 +decorator==4.4.2 +future==0.18.2 +gitdb==4.0.5 +GitPython==3.1.3 +google-auth==1.16.1 +google-auth-oauthlib==0.4.1 +googledrivedownloader==0.4 +grpcio==1.29.0 +h5py==2.10.0 +idna==2.9 +imageio==2.8.0 +importlib-metadata==1.6.1 +isodate==0.6.0 +joblib==0.15.1 +jsonpointer==2.0 +jsonref==0.2 +jsonschema==3.2.0 +kiwisolver==1.2.0 +llvmlite==0.32.1 +Markdown==3.2.2 +matplotlib==3.2.1 +monotonic==1.5 +msgpack==1.0.0 +msgpack-python==0.5.6 +natsort==7.0.1 +neptune-client==0.4.109 +networkx==2.4 +numba==0.49.1 +numpy==1.18.5 +oauthlib==3.1.0 +pandas==1.0.4 +Pillow==7.1.2 +plyfile==0.7.2 +polyscope==0.1.2 +protobuf==3.12.2 +py3nvml==0.2.6 +pyasn1==0.4.8 +pyasn1-modules==0.2.8 +PyJWT==1.7.1 +pyod==0.8.0 +pyparsing==2.4.7 +pyrsistent==0.16.0 +python-dateutil==2.8.1 +pytorch-lightning==0.7.6 +pytz==2020.1 +PyWavelets==1.1.1 +PyYAML==5.3.1 +rdflib==5.0.0 +requests==2.23.0 +requests-oauthlib==1.3.0 +rfc3987==1.3.8 +rsa==4.0 +scikit-image==0.17.2 +scikit-learn==0.23.1 +scipy==1.4.1 +simplejson==3.17.0 +six==1.15.0 +smmap==3.0.4 +strict-rfc3339==0.7 +suod==0.0.4 +swagger-spec-validator==2.7.0 +tensorboard==2.2.2 +tensorboard-plugin-wit==1.6.0.post3 +test-tube==0.7.5 +threadpoolctl==2.1.0 +tifffile==2020.6.3 +torch==1.4.0+cpu +torch-cluster==1.5.4 +torch-geometric==1.4.3 +torch-scatter==2.0.4 +torch-sparse==0.6.1 +torchcontrib==0.0.2 +torchvision==0.5.0 +tqdm==4.45.0 +typing-extensions==3.7.4.2 +urllib3==1.25.9 +webcolors==1.11.1 +websocket-client==0.57.0 +Werkzeug==1.0.1 +xmltodict==0.12.0 +zipp==3.1.0 diff --git a/utils/pointcloud.py b/utils/pointcloud.py new file mode 100644 index 0000000..0c91799 --- /dev/null +++ b/utils/pointcloud.py @@ -0,0 +1,282 @@ +import numpy as np +from sklearn.cluster import DBSCAN + +from pyod.models.knn import KNN +from pyod.models.sod import SOD +from pyod.models.abod import ABOD +from pyod.models.sos import SOS +from pyod.models.pca import PCA +from pyod.models.ocsvm import OCSVM +from pyod.models.mcd import MCD +from pyod.models.lof import LOF +from pyod.models.cof import COF +from pyod.models.cblof import CBLOF +from pyod.models.loci import LOCI +from pyod.models.hbos import HBOS +from pyod.models.lscp import LSCP +from pyod.models.feature_bagging import FeatureBagging + +from utils.project_config import Classes + + +def polytopes_to_planes(pc): + pc[(pc[:, 6] == float(Classes.Box)) | (pc[:, 6] == float(Classes.Polytope)), 6] = float(Classes.Plane); + return pc + + +def mini_color_table(index, norm=True): + colors = [ + [0.,0.,0.], + [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 cluster2Color(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 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 read_pointcloud(path, delimiter=' ', hasHeader=True): + with open(path, 'r') as f: + if hasHeader: + # Get rid of the Header + _ = f.readline() + # This iterates over all lines, splits them and converts values to floats. Will fail on wrong values. + pc = [[float(x) for x in line.rstrip().split(delimiter)] for line in f if line != ''] + + return np.asarray(pc)[:, :6] + + +def write_pointcloud(file, pc, numCols=6): + np.savetxt(file, pc[:, :numCols], header=str(len(pc)) + ' ' + str(numCols), comments='') + + +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 cluster_per_column(pc, column): + clusters = [] + for i in range(0, int(np.max(pc[:, column]))): + cluster_pc = pc[pc[:, column] == i, :] + clusters.append(cluster_pc) + + return clusters + + +def cluster_cubes(data, cluster_dims, max_points_per_cluster=-1, min_points_per_cluster=-1): + if cluster_dims[0] is 1 and cluster_dims[1] is 1 and cluster_dims[2] is 1: + print("no need to cluster.") + return [farthest_point_sampling(data, 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 + final_clusters = [] + for key, cluster in clusters.items(): + c = np.vstack(cluster) + if c.shape[0] < min_points_per_cluster and -1 is not min_points_per_cluster: + continue + + if max_points_per_cluster is not -1: + final_clusters.append(farthest_point_sampling(c, max_points_per_cluster)) + else: + final_clusters.append(c) + + return final_clusters + + +def cluster_dbscan(data, selected_indices, eps, min_samples=5, metric='euclidean', algo='auto'): + # print('Clustering. Min Samples: ' + str(min_samples) + ' EPS: ' + str(eps) + "Selected Indices: " + str(selected_indices)) + + # 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, min_samples=min_samples, algorithm=algo).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 = [] + + for cluster_idx, cluster in enumerate(clusters): + cloud = o3d.PointCloud() + cloud.points = o3d.Vector3dVector(cluster[:, :3]) + cloud.colors = o3d.Vector3dVector(cluster2Color(cluster, cluster_idx)) + clouds.append(cloud) + + o3d.draw_geometries(clouds) + + +def write_clusters(path, clusters, type_column=6): + file = open(path, "w") + file.write(str(len(clusters)) + "\n") + + for cluster in clusters: + # print("Types: ", cluster[:, type_column]) + + types = np.unique(cluster[:, type_column], axis=0).astype(int) + + def type_mapping(t): + if t == 0: + return 2 + elif t == 1: + return 1 + elif t == 3: + return 4 + return t + + types = np.array([type_mapping(t) for t in types]) + print("Types: {}, Points: {}".format(types, cluster.shape[0])) + # draw_sample_data(cluster) + + np.savetxt(file, types.reshape(1, types.shape[0]), delimiter=';', header='', comments='', fmt='%i') + np.savetxt(file, cluster[:, :6], header=str(len(cluster)) + ' ' + str(6), comments='') + + +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 normalize_pointcloud(pc, factor=1.0): + 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 * factor) + pc[:, 3:6] /= (np.linalg.norm(pc[:, 3:6], ord=2, axis=1, keepdims=True)) + + return pc + + +def hierarchical_clustering(data, selected_indices_0, selected_indices_1, eps, min_samples=5, metric='euclidean', + algo='auto'): + total_clusters = [] + + clusters = cluster_dbscan(data, selected_indices_0, eps, min_samples, metric=metric, algo=algo) + + for cluster in clusters: + # cluster = normalize_pointcloud(cluster) + sub_clusters = cluster_dbscan(cluster, selected_indices_1, eps, min_samples, metric=metric, algo=algo) + total_clusters.extend(sub_clusters) + + return total_clusters + + +def filter_clusters(clusters, filter): + filtered_clusters = [] + + for c in clusters: + if filter(c): + filtered_clusters.append(c) + + return filtered_clusters + + +def split_outliers(pc, columns): + clf = LOF() # FeatureBagging() # detector_list=[LOF(), KNN()] + clf.fit(pc[:, columns]) + + # LOF, kNN + + return pc[clf.labels_ == 0], pc[clf.labels_ == 1] + + +def append_onehotencoded_type(data, factor=1.0): + types = data[:, 6].astype(int) + res = np.zeros((len(types), 8)) + res[np.arange(len(types)), types] = factor + + return np.column_stack((data, res))