IVUS-processing/generate_surface_mesh.py at master · msbazzi/IVUS-processing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
import vtk
import xml.etree.ElementTree as ET

def read_contour_from_ctgr(file_path):
    tree = ET.parse(file_path)
    root = tree.getroot()
    contours = []

    for timestep in root.findall('timestep'):
        for contour in timestep.findall('contour'):
            points = []
            for point in contour.find('control_points').findall('point'):
                x = float(point.get('x'))
                y = float(point.get('y'))
                z = float(point.get('z'))
                points.append([x, y, z])
            contours.append(points)

    return contours

def create_vtk_polydata(points):
    polydata = vtk.vtkPolyData()
    points_vtk = vtk.vtkPoints()
    lines = vtk.vtkCellArray()

    num_points = len(points)
    for i, point in enumerate(points):
        points_vtk.InsertNextPoint(point)
        if i > 0:
            lines.InsertNextCell(2)
            lines.InsertCellPoint(i - 1)
            lines.InsertCellPoint(i)

    # Close the contour by connecting the last point to the first
    if num_points > 2:
        lines.InsertNextCell(2)
        lines.InsertCellPoint(num_points - 1)
        lines.InsertCellPoint(0)

    polydata.SetPoints(points_vtk)
    polydata.SetLines(lines)

    return polydata

def create_surface_from_contours(contours):
    append_filter = vtk.vtkAppendPolyData()

    for i in range(len(contours) - 1):
        polydata1 = create_vtk_polydata(contours[i])
        polydata2 = create_vtk_polydata(contours[i + 1])

        loft_surface = vtk.vtkRuledSurfaceFilter()
        loft_surface.SetInputData(polydata1)
        loft_surface.SetInputData(polydata2)
        loft_surface.SetResolution(50, 5)
        loft_surface.SetRuledModeToResample()
        loft_surface.Update()

        append_filter.AddInputData(loft_surface.GetOutput())

    append_filter.Update()

    # Clean the polydata to remove any duplicate points
    clean_filter = vtk.vtkCleanPolyData()
    clean_filter.SetInputData(append_filter.GetOutput())
    clean_filter.Update()

    return clean_filter.GetOutput()

def write_stl(polydata, filename):
    if polydata is None or polydata.GetNumberOfPoints() == 0:
        print(f"No data to write for {filename}.")
        return

    stl_writer = vtk.vtkSTLWriter()
    stl_writer.SetFileName(filename)
    stl_writer.SetInputData(polydata)
    stl_writer.Write()

def main():
    inner_contours = read_contour_from_ctgr('/home/bazzi/TEVG/FSG/IVUS-processing/contours_ctgr/contours_inner.ctgr')
    outer_contours = read_contour_from_ctgr('/home/bazzi/TEVG/FSG/IVUS-processing/contours_ctgr/contours_outer.ctgr')

    print(f"Read {len(inner_contours)} inner contours.")
    print(f"Read {len(outer_contours)} outer contours.")

    inner_surface = create_surface_from_contours(inner_contours)
    outer_surface = create_surface_from_contours(outer_contours)

    # Write the resulting surface to STL files
    write_stl(inner_surface, 'inner_wall.stl')
    write_stl(outer_surface, 'outer_wall.stl')

    print('Inner and outer wall surfaces have been written to STL files.')

if __name__ == "__main__":
    main()