framework/cad_stability_check/main.py

# Алгоритм генерации графа с помощью оценки стабильности подсборок в физическом движке PyBullet
from typing import Any, TypeVar, Type, cast
import FreeCAD as App
import json
import importOBJ
import Draft
import os
import Part
import numpy as np
from typing import TypeAlias
import FreeCADGui as Gui

ISequencesUnion: TypeAlias = dict[str : dict[str : list[str]]]


def importObjAtPath(path: str, inc: int):
    try:
        importOBJ.insert("" + path, App.ActiveDocument.Label)

        mesh = App.ActiveDocument.Objects[inc]
        shape = Part.Shape()
        shape.makeShapeFromMesh(mesh.Mesh.Topology, 0.05)
        solid = Part.makeSolid(shape)

        Part.show(solid)

        App.ActiveDocument.Objects[inc + 1].Label = App.ActiveDocument.Objects[inc].Name
        App.ActiveDocument.removeObject(App.ActiveDocument.Objects[inc].Name)
        return App.ActiveDocument.Objects[inc]
    except:
        print("path")
        print(path)
        print("inc")
        print(inc)


T = TypeVar("T")


def from_float(x: Any) -> float:
    assert isinstance(x, (float, int)) and not isinstance(x, bool)
    return float(x)


def to_float(x: Any) -> float:
    assert isinstance(x, float)
    return x


def from_str(x: Any) -> str:
    assert isinstance(x, str)
    return x


def to_class(c: Type[T], x: Any) -> dict:
    assert isinstance(x, c)
    return cast(Any, x).to_dict()


def euler_to_quaternion(yaw, pitch, roll):
    qx = np.sin(roll / 2) * np.cos(pitch / 2) * np.cos(yaw / 2) - np.cos(
        roll / 2
    ) * np.sin(pitch / 2) * np.sin(yaw / 2)
    qy = np.cos(roll / 2) * np.sin(pitch / 2) * np.cos(yaw / 2) + np.sin(
        roll / 2
    ) * np.cos(pitch / 2) * np.sin(yaw / 2)
    qz = np.cos(roll / 2) * np.cos(pitch / 2) * np.sin(yaw / 2) - np.sin(
        roll / 2
    ) * np.sin(pitch / 2) * np.cos(yaw / 2)
    qw = np.cos(roll / 2) * np.cos(pitch / 2) * np.cos(yaw / 2) + np.sin(
        roll / 2
    ) * np.sin(pitch / 2) * np.sin(yaw / 2)

    return [qx, qy, qz, qw]


class Coords:
    x: float
    y: float
    z: float

    def __init__(self, x: float, y: float, z: float) -> None:
        self.x = x
        self.y = y
        self.z = z

    @staticmethod
    def from_dict(obj: Any) -> "Coords":
        assert isinstance(obj, dict)
        x = from_float(obj.get("x"))
        y = from_float(obj.get("y"))
        z = from_float(obj.get("z"))
        return Coords(x, y, z)

    def to_dict(self) -> dict:
        result: dict = {}
        result["x"] = to_float(self.x)
        result["y"] = to_float(self.y)
        result["z"] = to_float(self.z)
        return result

    def vector(self):
        return App.Vector(self.x, self.y, self.z)


class MotionResultModel:
    id: str
    quaternion: Coords
    position: Coords

    def __init__(self, id: str, quaternion: Coords, position: Coords) -> None:
        self.id = id
        self.quaternion = quaternion
        self.position = position

    @staticmethod
    def from_dict(obj: Any) -> "MotionResultModel":
        id = from_str(obj.get("id"))
        quaternion = Coords.from_dict(obj.get("quaternion"))
        position = Coords.from_dict(obj.get("position"))
        return MotionResultModel(id, quaternion, position)

    def to_dict(self) -> dict:
        result: dict = {}
        result["id"] = from_str(self.id)
        result["quaternion"] = to_class(Coords, self.quaternion)
        result["position"] = to_class(Coords, self.position)
        return result


class SequencesEvaluation:
    sequences: ISequencesUnion
    assemblyDir: str
    result: dict = {}

    def __init__(self, sequences, assemblyDir) -> None:
        self.sequences = sequences
        self.assemblyDir = assemblyDir
        pass

    def assemblyComputed(self):
        isOk = True
        for sequenceNumber, v in self.sequences.items():
            for assemblyNumber, assemblySequenced in v.items():
                if isOk:
                    # sequenceNumber += 1
                    isOk = False
                    self.comptedAssembly(
                        assemblySequenced, sequenceNumber, assemblyNumber
                    )

        pass

    def comptedAssembly(
        self, assembly: list[str], sequenceNumber: int, assemblyNumber: int
    ):
        assemblyParts = []
        for counter in range(len(assembly)):
            importObjAtPath(
                self.assemblyDir + "generation/meshes/" + assembly[counter] + ".obj",
                counter,
            )
            assemblyParts.append(
                {"part": App.ActiveDocument.Objects[counter], "name": assembly[counter]}
            )

        motionResult = json.loads(
            (
                open(
                    self.assemblyDir
                    + "stability"
                    + "/"
                    + str(sequenceNumber + 1)
                    + "/"
                    + str(assemblyNumber)
                    + "/"
                    + "motion_result.json"
                )
            ).read()
        )

        simulatorMotionResults: list["MotionResultModel"] = []
        for _k, v in motionResult.items():
            print(v)
            simulatorMotionResults.append(MotionResultModel.from_dict(v))

        for el in simulatorMotionResults:
            for e in assemblyParts:
                # сопоставляем детали
                if el.id == e.get("name"):
                    # вычисляем центр детали для перемещения
                    center = e.get("part").Shape.CenterOfMass
                    # получаем центр деталей из симуляции
                    new_center = App.Vector(el.position.vector())

                    # вычисляем вектор смещения
                    offset = new_center - center
                    # перемещаем деталь на вектор смещения
                    e.get("part").Placement.Base += offset

        # # импортируем меш связанный с зоной обьекта
        zonePart = importObjAtPath(
            self.assemblyDir
            + "stability/zones/meshes/zone_sub_assembly"
            + str(assembly.__len__() - 1)
            + ".obj",
            len(App.ActiveDocument.Objects),
        )
        # # получаем координаты зоны относительно детали за которой закреплена зона

        coords = json.loads(
            (
                open(
                    self.assemblyDir
                    + "stability/zones/sub_assembly_coords_"
                    + str(assemblyNumber - 1)
                    + ".json"
                )
            ).read()
        )
        assemblyCounter = len(assemblyParts)
        detailWhichZoneBindings = assemblyParts[assemblyCounter - 2].get("part")

        detailStabilityComputed = assemblyParts[assemblyCounter - 1].get("part")
        relativeCoordinates = coords.get("relativeCoordinates")

        relativeEuler = coords.get("relativeEuler")

        specificEuler = {
            "yaw": detailWhichZoneBindings.Placement.Rotation.toEuler()[0]
            - relativeEuler.get("yaw"),
            "pitch": detailWhichZoneBindings.Placement.Rotation.toEuler()[1]
            - relativeEuler.get("pitch"),
            "roll": detailWhichZoneBindings.Placement.Rotation.toEuler()[2]
            - relativeEuler.get("roll"),
        }
        quaternion = euler_to_quaternion(
            specificEuler.get("yaw"),
            specificEuler.get("pitch"),
            specificEuler.get("roll"),
        )
        rotation = App.Rotation(
            quaternion[0], quaternion[1], quaternion[2], quaternion[3]
        )

        detailStabilityComputed.Placement.Rotation = rotation
        centerVector = detailWhichZoneBindings.Shape.CenterOfMass
        vector = App.Vector(
            relativeCoordinates.get("x"),
            relativeCoordinates.get("y"),
            relativeCoordinates.get("z"),
        )
        print(vector)
        # TODO
        current_center = zonePart.Shape.CenterOfMass
        move_vector = App.Vector(centerVector + vector) - current_center
        zonePart.Placement.move(move_vector)

        computedStabilityResult = computedStability(zonePart, detailStabilityComputed)

        if sequenceNumber not in self.result.keys():
            self.result[sequenceNumber] = []

        self.result[sequenceNumber].append(
            {str(assemblyNumber): assembly, "result": computedStabilityResult}
        )
        # for part in App.ActiveDocument.Objects:
        #     App.ActiveDocument.removeObject(part.Name)


def get_part_center(part):
    shape = None
    if not hasattr(part, "Shape"):
        shape = part.Mesh
    if hasattr(part, "Shape"):
        shape = part.Shape

    center = shape.BoundBox.Center
    return App.Vector(center[0], center[1], center[2])


def move_second_part_to_match_center(first_part, second_part):
    first_center = get_part_center(first_part)
    second_center = get_part_center(second_part)
    offset = first_center - second_center
    second_part.Placement.move(offset)


def create(part):
    clone = Draft.make_clone([part], forcedraft=True)

    clone.Scale = App.Vector(1.30, 1.30, 1.30)
    clone_corr = (
        App.Vector(0.4476673941774023, -2.109332894191716, -0.5918687740295264)
        - clone.Placement.Base
    ).scale(*App.Vector(-0.25, -0.25, -0.25))
    clone.Placement.move(clone_corr)
    App.ActiveDocument.recompute()

    return clone


def getFullPathObj(assemblyFolder: str, name: str):
    return assemblyFolder + "sdf/meshes/" + name + ".obj"


def computedStability(refElement, childElement):
    rootElement = childElement.Shape.BoundBox
    # Создание обьекта на котором делается операция пересечения
    App.activeDocument().addObject("Part::MultiCommon", "Common")
    App.activeDocument().Common.Shapes = [
        refElement,
        childElement,
    ]
    App.ActiveDocument.getObject("Common").ViewObject.ShapeColor = getattr(
        App.ActiveDocument.getObject(refElement.Name).getLinkedObject(True).ViewObject,
        "ShapeColor",
        App.ActiveDocument.getObject("Common").ViewObject.ShapeColor,
    )
    App.ActiveDocument.getObject("Common").ViewObject.DisplayMode = getattr(
        App.ActiveDocument.getObject(childElement.Name)
        .getLinkedObject(True)
        .ViewObject,
        "DisplayMode",
        App.ActiveDocument.getObject("Common").ViewObject.DisplayMode,
    )
    App.ActiveDocument.recompute()
    obj = App.ActiveDocument.getObjectsByLabel("Common")[0]

    shp = obj.Shape
    bbox = shp.BoundBox
    # Если после операции пересечения зона обьекта совпадает с зоной тестируемого обьекта то тест прошел успешно
    if (
        bbox.XLength == rootElement.XLength
        and bbox.YLength == rootElement.YLength
        and rootElement.ZLength == bbox.ZLength
    ):
        return True
    return False


def autoStabilityZoneComputed(stepFilesPaths: list[str], directoryStableZonesPath: str):
    cadObjects = []
    for count in range(len(stepFilesPaths)):
        importObjAtPath(stepFilesPaths[count], count)
        cadObjects.append(App.ActiveDocument.Objects[count])

    assemblesBindings = []

    for increment in range(len(cadObjects)):
        if increment != 0:
            detailForEvaluationZ = cadObjects[increment]
            zoneBindingDetailZ = cadObjects[increment - 1]
            assemblesBindings.append(
                {
                    "zoneBindingDetail": detailForEvaluationZ,
                    "detailForEvaluation": zoneBindingDetailZ,
                    "relativeCoordinates": None,
                    "zonePart": None,
                }
            )
    for increment in range(len(assemblesBindings)):
        el = assemblesBindings[increment]

        zoneBindingDetail = el.get("zoneBindingDetail")
        zoneBindingDetailCenterVector = zoneBindingDetail.Shape.CenterOfMass
        zoneDetail = create(el.get("detailForEvaluation"))
        move_second_part_to_match_center(el.get("zoneBindingDetail"), zoneDetail)
        zoneDetail.Label = "zone_sub_assembly" + str(increment + 1)

        zoneDetail.ViewObject.ShapeColor = (0.40, 0.74, 0.71)
        zoneDetail.ViewObject.Transparency = 50
        zoneDetailCenterVector = el.get("detailForEvaluation").Shape.CenterOfMass

        el["relativeCoordinates"] = {
            "x": zoneBindingDetailCenterVector.x - zoneDetailCenterVector.x,
            "y": zoneBindingDetailCenterVector.y - zoneDetailCenterVector.y,
            "z": zoneBindingDetailCenterVector.z - zoneDetailCenterVector.z,
        }

        el["relativeEuler"] = {
            "yaw": zoneBindingDetail.Placement.Rotation.toEuler()[0]
            - el.get("detailForEvaluation").Placement.Rotation.toEuler()[0],
            "pitch": zoneBindingDetail.Placement.Rotation.toEuler()[1]
            - el.get("detailForEvaluation").Placement.Rotation.toEuler()[1],
            "roll": zoneBindingDetail.Placement.Rotation.toEuler()[2]
            - el.get("detailForEvaluation").Placement.Rotation.toEuler()[2],
        }
        el["zonePart"] = zoneDetail

    meshesPath = directoryStableZonesPath + "meshes/"
    if not os.path.exists(directoryStableZonesPath):
        os.makedirs(directoryStableZonesPath)
    if not os.path.exists(meshesPath):
        os.makedirs(meshesPath)
    zonesSaved = {}
    for counter in range(len(assemblesBindings)):
        zoneComputed = assemblesBindings[counter]
        mesh = zoneComputed.get("zonePart")
        zonesSavePath = meshesPath + mesh.Label + ".obj"
        importOBJ.export([mesh], zonesSavePath)
        zonesSaved[mesh.Label] = "meshes/" + mesh.Label + ".obj"
    for counter in range(len(assemblesBindings)):
        el = assemblesBindings[counter]
        savePath = zonesSaved[el.get("zonePart").Label]
        el["zonePart"] = savePath
        el["detailForEvaluation"] = el["detailForEvaluation"].Label
        el["zoneBindingDetail"] = el["zoneBindingDetail"].Label

        json_result = json.dumps(el)

        file_to_open = (
            directoryStableZonesPath + "sub_assembly_coords_" + str(counter) + ".json"
        )

        f = open(
            file_to_open,
            "w",
        )

        f.write(json_result)
        f.close()


def main():
    App.newDocument()
    env = json.loads((open("./env.json")).read())

    assemblyDir = env.get("aspPath")
    sequencesJSON = json.loads((open(assemblyDir + "sequences.json")).read())
    directoryStableZones = assemblyDir + "stability/zones/"

    sequences = sequencesJSON.get("sequences")
    stepStructure = json.loads((open(assemblyDir + "step-structure.json")).read())

    stepFilesPaths = []
    for step in stepStructure:
        stepFilesPaths.append(assemblyDir + "generation/meshes/" + step + ".obj")
    if not os.path.exists(directoryStableZones):
        print("Zones not found automatic calculation started")
        autoStabilityZoneComputed(stepFilesPaths, directoryStableZones)

    sequencesJoin = {}

    for arrayCounter in range(len(sequences)):
        for indexCounter in range(len(sequences[arrayCounter])):
            if indexCounter != 0:
                if sequencesJoin.get(arrayCounter) == None:
                    sequencesJoin[arrayCounter] = {
                        indexCounter: sequences[arrayCounter][0 : indexCounter + 1]
                    }
                else:
                    sequencesJoin[arrayCounter][indexCounter] = sequences[arrayCounter][
                        0 : indexCounter + 1
                    ]

    seqEvaluation = SequencesEvaluation(sequencesJoin, assemblyDir)
    seqEvaluation.assemblyComputed()


main()