arzhela/QuadOpt-RL

1import string2import json3import os4 5from mesh_model.mesh_struct.mesh import Mesh6 7def getMeshFiles(d):8    l = []9    for sdp, Lsd, Lnf in os.walk(d):10        for f in Lnf:11            if (f.endswith(".msh")):12                l.append(f)13    l.sort()14    return l15 16def getFiles(d):17    l = []18    for sdp, Lsd, Lnf in os.walk(d):19        for f in Lnf:20            l.append(f)21    l.sort()22    return l23 24def read_medit(filename: string) -> Mesh:25    """26    Create a mesh read from a Medit .mesh file.27    :param filename: name of the file28    :return: a mesh29    """30 31    f = open(filename, 'r')32 33    nodes = []34    faces = []35    while True:36        line = f.readline()37        if not line:38            break39        if "Vertices" == line.strip():40            line = f.readline()41            nb_vertices = line.strip()42            for _ in range(int(nb_vertices)):43                line = f.readline()44                ls = line.split()45                x = float(ls[0])46                y = float(ls[1])47                nodes.append([x, y])48        if "Triangles" == line.strip():49            line = f.readline()50            nb_faces = line.strip()51            for _ in range(int(nb_faces)):52                line = f.readline()53                ls = line.split()54                n0 = int(ls[0])55                n1 = int(ls[1])56                n2 = int(ls[2])57                faces.append([n0 - 1, n1 - 1, n2 - 1])58 59    f.close()60    mesh = Mesh(nodes, faces)61 62    return mesh63 64 65def read_gmsh(filename: string) -> Mesh:66    """67    Create a mesh read from a gmsh .msh file.68    :param filename: name of the file69    :return: a mesh70    """71 72    f = open(filename, 'r')73 74    nodes = []75    faces = []76    while True:77        line = f.readline()78        if not line:79            break80        if "$Nodes" == line.strip():81            line = f.readline()82            ls = line.split()83            nb_blocs = int(ls[0])84            for _ in range(nb_blocs):85                line = f.readline()86                ls = line.split()87                nb_nodes_b = int(ls[3])88                # skip the tags89                for _ in range(nb_nodes_b):90                    line = f.readline()91                for _ in range(nb_nodes_b):92                    line = f.readline()93                    ls = line.split()94                    x = float(ls[0])95                    y = float(ls[1])96                    z = float(ls[2])97                    nodes.append([x, y, z])98 99        if "$Elements" == line.strip():100            line = f.readline()101            ls = line.split()102            nb_blocs = int(ls[0])103            for _ in range(nb_blocs):104                line = f.readline()105                ls = line.split()106                elem_type = int(ls[2])107                nb_elems_b = int(ls[3])108                for _ in range(nb_elems_b):109                    line = f.readline()110                    ls = line.split()111                    if elem_type == 2:112                        n0 = int(ls[1])113                        n1 = int(ls[2])114                        n2 = int(ls[3])115                        faces.append([n0 - 1, n1 - 1, n2 - 1])116                    elif elem_type == 3:117                        n0 = int(ls[1])118                        n1 = int(ls[2])119                        n2 = int(ls[3])120                        n3 = int(ls[4])121                        faces.append([n0 - 1, n1 - 1, n2 - 1, n3 - 1])122                    elif elem_type == 1:  # skip 2-node line elements123                        continue124                    elif elem_type == 15: # skip 1-node point elements125                        continue126                    else:127                        print("element_type " + str(elem_type) + " not handled")128                        exit(1)129    f.close()130    mesh = Mesh(nodes, faces)131 132    return mesh133 134def read_json(filename) -> Mesh:135    with open(filename, 'r') as f:136        json_mesh = json.load(f)137    mesh = Mesh(json_mesh['nodes'], json_mesh['faces'])138    return mesh139 140def read_dataset(dataset_dir) -> list[Mesh]:141    mesh_dataset = []142    for f in getFiles(dataset_dir):143        if (f.endswith(".msh")):144            msh_f = dataset_dir + "/" + f145            cmap = read_gmsh(msh_f)146            mesh_dataset.append(cmap)147        elif (f.endswith(".json")):148            json_f = dataset_dir + "/" + f149            cmap = read_json(json_f)150            mesh_dataset.append(cmap)151    return mesh_dataset