BUG: The trained model outputs 4 labels now

Due to an error, the trained model output 3 labels but now it has 3. Some small changes to the pipeline doesn't crash

Author: juanprietob

src/py/post_process.py

@@ -367,6 +367,11 @@ def ErodeLabel(vtkdata, labels, label):
 		else:
 			break
 
+def ReLabel(surf, labels, label, relabel):
+	for pid in range(labels.GetNumberOfTuples()):
+		if labels.GetTuple(pid)[0] == label:
+			labels.SetTuple(pid, (relabel,))
+
 def Threshold(vtkdata, labels, threshold_min, threshold_max):
 
 	threshold = vtk.vtkThreshold()

src/py/predict.py

@@ -130,7 +130,7 @@ def CreatePlane(Origin,Point1,Point2,Resolution):
 tree.SetDataSet(surf) 
 tree.BuildLocator()
 
-label_array = np.zeros([surf.GetNumberOfPoints(), 3])
+label_array = np.zeros([surf.GetNumberOfPoints(), 4])
 
 icosahedron = CreateIcosahedron(sphereRadius, numberOfSubdivisions)
 
@@ -249,12 +249,15 @@ def CreatePlane(Origin,Point1,Point2,Resolution):
 	for label in range(int(labels_range[0]), int(labels_range[1]) + 1):
 		print("Removing islands:", label)
 		post_process.RemoveIslands(surf, real_labels, label, 500)
+
+	print("Relabel...")
+	post_process.ReLabel(surf, real_labels, 3, -2)
 
 	print("Connectivity...")
 	post_process.ConnectivityLabeling(surf, real_labels, 2, 2)
 
 	print("Eroding...")
-	post_process.ErodeLabel(surf, real_labels, 0)
+	post_process.ErodeLabel(surf, real_labels, -2)
 
 	print("Writting:", outfilename)
 	polydatawriter = vtk.vtkPolyDataWriter()