Lets try the best solution we have currently.

script.py

@@ -12,6 +12,11 @@ import gc
 from utils import empty_solution
 from predict import predict_wireframe
 
+from fast_pointnet import load_pointnet_model
+from fast_voxel import load_3dcnn_model
+from fast_pointnet_class import load_pointnet_model as load_pointnet_class_model
+import torch
+
 if __name__ == "__main__":
     print ("------------ Loading dataset------------ ")
     param_path = Path('params.json')
@@ -72,12 +77,22 @@ if __name__ == "__main__":
 
     print(dataset, flush=True)
     
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    pnet_model = load_pointnet_model(model_path="pnet.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.4, 'edge_threshold': 0.6, 'only_predicted_connections': False}
+
     print('------------ Now you can do your solution ---------------')
     solution = []
 
     def process_sample(sample, i):
         try:
-            pred_vertices, pred_edges = predict_wireframe(sample)
+            pred_vertices, pred_edges = predict_wireframe(sample, pnet_model, voxel_model, pnet_class_model, config)
         except:
             pred_vertices, pred_edges = empty_solution()
         if i %10 == 0: