lghizoni
lghizoni

Reputation: 155

How retrieve correct orientation of part from Step file with OpenCASCADE

I'm trying to retrieve correct position (translation in xyz), orientation (rotation in xyz, Euler angles), and colors of parts from a Step file by using a python wrapper of OpenCASCADE.

I tried implementing a snip from the code of DataExchange.py in my own to retrieve the information I need. The code works and gives the translation and rotation of the parts. However, lots of the parts have the exact same locations, and with 'weird' rotation numbers.

When I open the file with FreeCAD all the parts are visualized in the correct positions, but I can also see in FreeCAD the data of each, and some of parts indeed present the same locations as others. In my research I found something about the vertices actually having the correct locations, but I just need the data for the full part in order to use it in another application.

I've been cracking my head for weeks now without finding a solution. Finally, I'm reaching here for help. Here's my function (again, adapted from the code I found) to retrieve the information I need:

from OCC.Core.gp import gp_XYZ
from OCC.Core.BRepTools import breptools
from OCC.Core.TopLoc import TopLoc_Location
from OCC.Core.TDocStd import TDocStd_Document
from OCC.Core.UnitsMethods import unitsmethods
from OCC.Core.IFSelect import IFSelect_RetDone
from OCC.Core.Message import Message_ProgressRange
from OCC.Core.TDF import TDF_LabelSequence, TDF_Label
from OCC.Core.TCollection import TCollection_AsciiString
from OCC.Core.STEPCAFControl import STEPCAFControl_Reader
from OCC.Core.gp import gp_EulerSequence, gp_Extrinsic_XYZ
from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform
from OCC.Core.Quantity import Quantity_Color, Quantity_TOC_RGB
from OCC.Core.TColStd import TColStd_IndexedDataMapOfStringString
from OCC.Core.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ColorTool

def retrieveStepData(filename):

    output_shapes = {}

    # Create a document handler
    doc = TDocStd_Document("aStepDoc")

    # Get tools to handle parts
    shape_tool = XCAFDoc_DocumentTool.ShapeTool(doc.Main())
    color_tool = XCAFDoc_DocumentTool.ColorTool(doc.Main())

    step_reader = STEPCAFControl_Reader()
    step_reader.SetColorMode(True)
    step_reader.SetNameMode(True)

    status = step_reader.ReadFile(filename)
    if status == IFSelect_RetDone:
        step_reader.Transfer(doc)

    locs = []

    def _getSubShapes(lab, loc, level=0):

        l_subss = TDF_LabelSequence()
        shape_tool.GetSubShapes(lab, l_subss)

        l_comps = TDF_LabelSequence()
        shape_tool.GetComponents(lab, l_comps)

        name = lab.GetLabelName()

        # Check if part is Assembly
        if shape_tool.IsAssembly(lab):

            # Get shape to be converted to OBJ format
            shape = shape_tool.GetShape(lab)

            l_c = TDF_LabelSequence()
            shape_tool.GetComponents(lab, l_c)

            label = l_c.Value(1)
            name = label.GetLabelName()

            print(f"{name} : ASSEMBLY")

            if shape_tool.IsReference(label):

                label_reference = TDF_Label()
                shape_tool.GetReferredShape(label, label_reference)
                loc = shape_tool.GetLocation(label)

                trans = loc.Transformation()
                rot = trans.GetRotation()

                print("         ROTATION     :")
                print("              X  : {}".format(rot.X()))
                print("              Y  : {}".format(rot.Y()))
                print("              Z  : {}".format(rot.Z()))
                print("              W  : {}".format(rot.W()))

                tran = trans.TranslationPart()
                print("         TRANSLATION  :")
                print("              X  : {}".format(tran.X()))
                print("              Y  : {}".format(tran.Y()))
                print("              Z  : {}".format(tran.Z()))

            l_c = TDF_LabelSequence()
            shape_tool.GetComponents(lab, l_c)
            for i in range(l_c.Length()):
                label = l_c.Value(i + 1)
                name = label.GetLabelName()

                if shape_tool.IsReference(label):

                    label_reference = TDF_Label()
                    shape_tool.GetReferredShape(label, label_reference)
                    loc = shape_tool.GetLocation(label)

                    locs.append(loc)
                    _getSubShapes(label_reference, loc, level+1)
                    locs.pop()

        elif shape_tool.IsSimpleShape(lab):

            # Get shape to be converted to OBJ format
            shape = shape_tool.GetShape(lab)

            # Get Location and Oriention in World
            loc = TopLoc_Location()
            for l in locs:
                loc = loc.Multiplied(l)

            print(f"{name} : PART")

            # Retrieve translation and rotation data
            trans = loc.Transformation()

            rot = trans.GetRotation()

            print("         ROTATION     :")
            print("              X  : {}".format(rot.X()))
            print("              Y  : {}".format(rot.Y()))
            print("              Z  : {}".format(rot.Z()))
            print("              W  : {}".format(rot.W()))

            tran = trans.TranslationPart()
            print("         TRANSLATION  :")
            print("              X  : {}".format(tran.X()))
            print("              Y  : {}".format(tran.Y()))
            print("              Z  : {}".format(tran.Z()))

            c = Quantity_Color(0.5, 0.5, 0.5, Quantity_TOC_RGB)  # default color
            color_set = False
            if (
                color_tool.GetInstanceColor(shape, 0, c)
                or color_tool.GetInstanceColor(shape, 1, c)
                or color_tool.GetInstanceColor(shape, 2, c)
            ):
                color_tool.SetInstanceColor(shape, 0, c)
                color_tool.SetInstanceColor(shape, 1, c)
                color_tool.SetInstanceColor(shape, 2, c)
                color_set = True
                n = c.Name(c.Red(), c.Green(), c.Blue())
                print(
                    "    instance ONE color Name & RGB: ",
                    c,
                    n,
                    c.Red(),
                    c.Green(),
                    c.Blue(),
                )

            if not color_set:
                if (
                    XCAFDoc_ColorTool.GetColor(lab, 0, c)
                    or XCAFDoc_ColorTool.GetColor(lab, 1, c)
                    or XCAFDoc_ColorTool.GetColor(lab, 2, c)
                ):
                    color_tool.SetInstanceColor(shape, 0, c)
                    color_tool.SetInstanceColor(shape, 1, c)
                    color_tool.SetInstanceColor(shape, 2, c)

                    n = c.Name(c.Red(), c.Green(), c.Blue())
                    print(
                        "    shape ONE color Name & RGB: ",
                        c,
                        n,
                        c.Red(),
                        c.Green(),
                        c.Blue(),
                    )

            shape_disp = BRepBuilderAPI_Transform(shape, loc.Transformation()).Shape()
            if shape_disp not in output_shapes:
                output_shapes[shape_disp] = [lab.GetLabelName(), c]

            if(l_subss.Length() > 0):

                lab_subs = l_subss.Value(1)
                shape_sub = shape_tool.GetShape(lab_subs)

                c = Quantity_Color(0.5, 0.5, 0.5, Quantity_TOC_RGB)  # default color
                color_set = False
                if (
                    color_tool.GetInstanceColor(shape_sub, 0, c)
                    or color_tool.GetInstanceColor(shape_sub, 1, c)
                    or color_tool.GetInstanceColor(shape_sub, 2, c)
                ):
                    color_tool.SetInstanceColor(shape_sub, 0, c)
                    color_tool.SetInstanceColor(shape_sub, 1, c)
                    color_tool.SetInstanceColor(shape_sub, 2, c)
                    color_set = True
                    n = c.Name(c.Red(), c.Green(), c.Blue())
                    print(
                        "    instance TWO color Name & RGB: ",
                        c,
                        n,
                        c.Red(),
                        c.Green(),
                        c.Blue(),
                    )

                if not color_set:
                    if (
                        XCAFDoc_ColorTool.GetColor(lab_subs, 0, c)
                        or XCAFDoc_ColorTool.GetColor(lab_subs, 1, c)
                        or XCAFDoc_ColorTool.GetColor(lab_subs, 2, c)
                    ):
                        color_tool.SetInstanceColor(shape, 0, c)
                        color_tool.SetInstanceColor(shape, 1, c)
                        color_tool.SetInstanceColor(shape, 2, c)

                        print("    COLORS:")
                        print("         RED     : {}".format(c.Red()))
                        print("         GREEN   : {}".format(c.Green()))
                        print("         BLUE    : {}".format(c.Blue()))

            else:
                print("    COLORS:")
                print("         RED     : 0.50000")
                print("         GREEN   : 0.50000")
                print("         BLUE    : 0.50000")

    def _getShapes():

        labels = TDF_LabelSequence()
        shape_tool.GetFreeShapes(labels)

        for i in range(labels.Length()):
            root_item = labels.Value(i + 1)
            _getSubShapes(root_item, None)

    _getShapes()
    return output_shapes

I would really appreciate any inputs to my problem.

Upvotes: 0

Views: 103

Answers (0)

Related Questions