import time

from vtk import vtkDataSetReader, vtkDataSetWriter, vtkSurfaceNets3D

# read the input dataset - a vtkImageData containing integer labels from a medical image segemntation
# labels are: 0 for background, 1 to 17 different organs

rd = vtkDataSetReader()
rd.SetFileName("abdominal_organ_segmentation_reduced.vtk")
rd.Update()

input_label_image = rd.GetOutput()

# extract the organ surfaces 
snets = vtkSurfaceNets3D()
snets.SetInputData(input_label_image)
snets.SmoothingOn()
snets.SetNumberOfIterations(5)
snets.SetRelaxationFactor(0.1)
snets.SetConstraintDistance(5)

snets.Update()

extracted_surfaces = snets.GetOutput()

wr = vtkDataSetWriter()
wr.SetInputData(extracted_surfaces)
wr.SetFileName("extracted_surfaces.vtk")
wr.Write()