x

predict.py

@@ -10,7 +10,7 @@ from PIL import Image as PImage
 import cv2
 #import open3d as o3d
 #from visu import plot_reconstruction_local, plot_wireframe_local, plot_bpo_cameras_from_entry_local
-import pyvista as pv
+#import pyvista as pv
 #from fast_pointnet import save_patches_dataset, predict_vertex_from_patch
 from fast_pointnet_v2 import save_patches_dataset, predict_vertex_from_patch
 #from fast_voxel import predict_vertex_from_patch_voxel

script.py

@@ -9,7 +9,7 @@ import gc
 from utils import empty_solution
 from predict import predict_wireframe
 
-from fast_pointnet import load_pointnet_model
+from fast_pointnet_v2 import load_pointnet_model
 from fast_pointnet_class import load_pointnet_model as load_pointnet_class_model
 import torch
 
@@ -74,13 +74,13 @@ if __name__ == "__main__":
     
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    pnet_model = load_pointnet_model(model_path="pnet.pth", device=device, predict_score=True)
+    pnet_model = load_pointnet_model(model_path="initial_epoch_60_v2.pth", device=device, predict_score=True)
 
     pnet_class_model = load_pointnet_class_model(model_path="pnet_class.pth", device=device)
 
     voxel_model = None
 
-    config = {'vertex_threshold': 0.32, 'edge_threshold': 0.65, 'only_predicted_connections': True}
+    config = {'vertex_threshold': 0.59, 'edge_threshold': 0.65, 'only_predicted_connections': True}
 
     print('------------ Now you can do your solution ---------------')
     solution = []