main.py

import occ, mesh, utils
verts, faces, grps = mesh.from_stp('models/cylinder.stp')

from pytorch3d import ops, structures, loss as loss3d
from torch import Tensor, optim
V, F = Tensor(verts).to('cuda'), Tensor(faces).long().to('cuda')
G = { 'Cylinder': [], 'Plane': [] }
for typ, face in grps:
    if not typ in G:
        G[typ] = []
    G[typ].append(Tensor(face).long().to('cuda'))

AX = [utils.get_cylinder_axis(V, face) for face in G['Cylinder']]
F0 = G['Plane'][0].squeeze().unique()
B0 = V[F0, :]
B0[:, 1] += 1

F1 = G['Plane'][1].squeeze().unique()
B1 = V[F1, :]
B1[:, 1] += 1

X = V.clone().requires_grad_(True)
optimizer = optim.Adagrad([X], lr=0.01)
losses = []
for i in range(2000):
    loss = 0
    for face, ax in zip(G['Cylinder'], AX):
        loss += utils.get_cylinder_axis_loss(X, face, ax)
    for face in G['Plane']:
        loss += utils.get_plane_loss(X, face)
    loss += (B0 - X[F0, :]).norm(dim=1).mean()
    loss += (B1 - X[F1, :]).norm(dim=1).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    val = loss.cpu().item()
    losses.append(val)
    if i % 100 == 0:
        print(i, 'loss', val)

# %%
from matplotlib import pyplot as plt
import numpy as np
fig = plt.figure()
ax1 = fig.add_subplot(1, 2, 1, projection='3d')
ax1.scatter(verts[:, 0], verts[:, 1], verts[:, 2], c='g')
shape = X.detach().cpu().numpy()
for _, faces in grps:
    ax1.plot_trisurf(shape[:, 0], shape[:, 1], shape[:, 2], triangles=faces)
ax2 = fig.add_subplot(1, 2, 2)
ax2.plot(losses)
ax2.set_yscale('log')
plt.show()

mesh.py

import gmsh
import numpy as np
from gmsh import model

def from_stp(file: str, out = None):
    gmsh.initialize()

    model.occ.import_shapes(file)
    model.occ.synchronize()
    model.mesh.generate(2)

    node_tags, node_coords, _ = model.mesh.get_nodes_by_element_type(2)

    vert_coords, vert_map = { }, { }
    for i, j in enumerate(node_tags):
        vert_coords[j] = node_coords[i * 3:i * 3 + 3]
    for i, j in enumerate(vert_coords.keys()):
        vert_map[j] = i
    verts = np.array(list(vert_coords.values()))

    _, _, (node_tags, ) = model.mesh.get_elements(2)
    faces = []
    for i in range(0, len(node_tags), 3):
        a, b, c = node_tags[i:i + 3]
        faces.append([vert_map[a], vert_map[b], vert_map[c]])

    grps = []
    for _, surf_tag in model.get_entities(2):
        grp = []
        _, _, (node_tags, ) = model.mesh.get_elements(2, surf_tag)
        surf_type = model.get_type(2, surf_tag)
        for i in range(0, len(node_tags), 3):
            a, b, c = node_tags[i:i + 3]
            grp.append([vert_map[a], vert_map[b], vert_map[c]])
        if len(grp):
            grps.append([surf_type, grp])

    if out:
        gmsh.write(out)
    gmsh.finalize()
    return verts, faces, grps

if __name__ == '__main__':
    from_stp('cylinder2.stp')

occ.py

import numpy as np
from OCC.Core import BRepMesh
from OCC.Core.STEPControl import STEPControl_Reader
from OCC.Core.TopLoc import TopLoc_Location
from OCC.Core.BRep import BRep_Tool
from OCC.Core.BRepAdaptor import BRepAdaptor_Surface
from OCC.Extend.TopologyUtils import TopologyExplorer

def from_stp(file: str):
    reader = STEPControl_Reader()
    reader.ReadFile(file)
    reader.TransferRoot()
    shape = reader.Shape()

    mesher = BRepMesh.BRepMesh_IncrementalMesh()
    mesher.SetShape(shape)
    mesher.Perform()

    SURFACE_TYPES = [
        item.strip() for item in
        'plane, cylinder, cone, sphere, torus, beziersurface, bsplinesurface, surfaceofrevolution, surfaceofextrusion, othersurface'.split(',')
    ]
    verts, faces, grps = [], [], []
    exp = TopologyExplorer(shape)
    for g, face in enumerate(exp.faces()):
        loc = TopLoc_Location()
        mesh = BRep_Tool.Triangulation(face, loc)
        s = len(verts)
        grp = []
        for i in range(mesh.NbNodes()):
            p = mesh.Node(i + 1)
            verts.append([p.X(), p.Y(), p.Z()])
        for i in range(mesh.NbTriangles()):
            a, b, c = mesh.Triangle(i + 1).Get()
            f = [a - 1 + s, b - 1 + s, c - 1 + s]
            faces.append(f)
            grp.append(f)
        adaptor = BRepAdaptor_Surface(face)
        type = SURFACE_TYPES[adaptor.GetType()]
        grps.append([type, grp])

    # remove duplicated verts
    verts = np.floor(np.array(verts) / 1e-9) * 1e-9
    print('verts', verts.shape)
    verts, vert_map = np.unique(verts, return_inverse=True, axis=0)
    print('verts', verts.shape)
    faces = [[vert_map[v] for v in f] for f in faces]
    grps = [[[vert_map[v] for v in f] for f in faces] for faces in grps]
    
    return verts, faces, grps

utils.py

import torch
from torch import Tensor

def get_face_normals(V: Tensor, F: Tensor):
    a, b, c = [V[F[:, i], :] for i in range(3)]
    n = (b - a).cross(c - a, dim=1)
    n = n / n.norm(dim=1).reshape([-1, 1])
    return n

def get_plane_loss(V: Tensor, F: Tensor):
    n = get_face_normals(V, F)
    return n.std(dim=0).sum()

def get_cylinder_axis(V: Tensor, F: Tensor):
    n = get_face_normals(V, F)
    a = n[1:, :].cross(n[:-1, :], dim=1)
    a = a / a.norm(dim=1).reshape([-1, 1])
    a = a * torch.sign(a[:, 2]).reshape([-1, 1])
    a = a.mean(dim=0)
    return a / a.norm()

def get_cylinder_axis_loss(V: Tensor, F: Tensor, ax: Tensor):
    n = get_face_normals(V, F)
    return n.matmul(ax).square().mean()

def get_cylinder_face_loss(V: Tensor, F: Tensor, ax: Tensor):
    a, b, c = [V[F[:, i], :] for i in range(3)]

    # project on cylinder axis
    v = ax.reshape([1, -1])
    a, b, c = [x - v * x.matmul(ax).reshape([-1, 1]) for x in [a, b, c]]
    a, b, c = [x.norm(dim=1) for x in [a - b, b - c, c - a]]

    # get external circle radius
    p = (a + b + c) / 2
    s = p * (p - a) * (p - b) * (p - c)
    r = a * a * b * b * c * c / s
    r = r[r < 1e3]
    return r.std()