steffen/point_to_primitive
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

main_pipeline fixed

Browse Source
This commit is contained in:
Si11ium 2020-07-03 14:45:08 +02:00
parent 5353220890
commit 95b1503f78
2 changed files with 9 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
main_inference.py
View File

@ -18,6 +18,7 @@ from ml_lib.utils.model_io import SavedLightningModels
from datasets.shapenet import ShapeNetPartSegDataset
from utils.project_config import ThisConfig
raise BrokenPipeError('There are Imports that need to be fixed first!!!!')
def prepare_dataloader(config_obj):
dataset = ShapeNetPartSegDataset(config_obj.data.root, mode=GlobalVar.data_split.test,

11
main_pipeline.py
View File

@ -19,7 +19,7 @@ from datasets.shapenet import ShapeNetPartSegDataset
from models import PointNet2
from utils.pointcloud import cluster_cubes, append_onehotencoded_type, label2color, \
write_clusters, cluster2Color, cluster_dbscan
from utils.project_settings import GlobalVar, DataClass
from utils.project_settings import dataSplit, DataClass
class DisplayMode(DataClass):
@ -27,6 +27,7 @@ class DisplayMode(DataClass):
Types = 1,
Nothing = 2
def restore_logger_and_model(log_dir):
model = SavedLightningModels.load_checkpoint(models_root_path=log_dir, model=PointNet2, n=-1)
model = model.restore()
@ -47,7 +48,7 @@ def predict_prim_type(input_pc, model):
batch_to_data = BatchToData()
data = batch_to_data(input_data)
y = loaded_model(data.to(device='cuda' if torch.cuda.is_available() else 'cpu'))
y = model(data.to(device='cuda' if torch.cuda.is_available() else 'cpu'))
y_primary = torch.argmax(y.main_out, dim=-1).cpu().numpy()
if input_pc.shape[1] > 6:
@ -69,7 +70,7 @@ if __name__ == '__main__':
loaded_model.eval()
transforms = Compose([NormalizeScale(), ])
test_dataset = ShapeNetPartSegDataset('data', mode=GlobalVar.data_split.predict, collate_per_segment=False,
test_dataset = ShapeNetPartSegDataset('data', mode=dataSplit.predict, collate_per_segment=False,
refresh=True, transform=transforms, cluster_type=None)
grid_clusters = cluster_cubes(test_dataset[0], grid_clusters, max_points_per_cluster=grid_cluster_max_pts)
@ -133,14 +134,14 @@ if __name__ == '__main__':
# ========================== Result visualization ==========================
if display_mode == DisplayMode.Types:
pc = ps.register_point_cloud("points_" + str(i), final_pc[:, :3], radius=0.01)
pc = ps.register_point_cloud("points_" + str(0), final_pc[:, :3], radius=0.01)
pc.add_color_quantity("prim types", label2color(final_pc[:, 6].astype(np.int64)), True)
elif display_mode == DisplayMode.Clusters:
for i, result_cluster in enumerate(result_clusters):
pc = ps.register_point_cloud("points_" + str(i), result_cluster[:, :3], radius=0.01)
pc.add_color_quantity("prim types", cluster2Color(result_cluster,i), True)
pc.add_color_quantity("prim types", cluster2Color(result_cluster, i), True)
ps.show()