Update my_solution.py

my_solution.py

@@ -1,5 +1,3 @@
-# Description: This file contains the handcrafted solution for the task of wireframe reconstruction 
-
 import io
 from PIL import Image as PImage
 import numpy as np
@@ -11,8 +9,8 @@ from scipy.spatial.distance import cdist
 from hoho.read_write_colmap import read_cameras_binary, read_images_binary, read_points3D_binary
 from hoho.color_mappings import gestalt_color_mapping, ade20k_color_mapping
 
-from helpers import my_empty_solution
-from geom_solver import GeomSolver
+from geom_solver import GeomSolver, my_empty_solution
+
 
 def convert_entry_to_human_readable(entry):
     out = {}
@@ -35,25 +33,31 @@ def convert_entry_to_human_readable(entry):
 
 
 def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
-    return entry['__key__'], np.zeros((17,3)), [(0, 1)]            
+    # return (entry['__key__'], *my_empty_solution())
+    vertices0, edges0 = my_empty_solution()
     try:
         solver = GeomSolver(entry)
         vertices, edges = solver.get_vertices()
     except:
-        # print("Error")
-        vertices, edges = my_empty_solution()
-    
-    if len(vertices) == 1:
-        # print("Added one more vertex")
-        vertices = np.concatenate((vertices, np.zeros((1,3))))
+        vertices, edges = vertices0, edges0
+
+    if vertices.shape[0] < vertices0.shape[0]:
+        verts_new = vertices0
+        verts_new[:vertices.shape[0]] = vertices
+        vertices = verts_new
+
+    # if len(vertices) == 1:
+    #     # print("Added one more vertex")
+    #     vertices = np.concatenate((vertices, np.zeros((1,3))))
+    #     edges = [(0,0)]
 
     if (len(edges) < 1) and (len(vertices) >= 2):
         # print("Added only edges")
-        _, edges = my_empty_solution()
+       edges = edges0
 
     if (len(vertices) < 2) or (len(edges) < 1):
         # print("Added empty solution")
-        vertices, edges = my_empty_solution()
+        vertices, edges = vertices0, edges0
 
     if visualize:
         from hoho.viz3d import plot_estimate_and_gt
@@ -62,5 +66,5 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
                                 entry['wf_vertices'],
                                 entry['wf_edges'])
     if vertices.shape[-1] != 3:
-        vertices, edges = my_empty_solution() 
+        vertices, edges = vertices0, edges0
     return entry['__key__'], vertices, edges