framework/insertion_vector_predicate/main.py

from scipy.spatial.transform import Rotation
import shutil
from spatialmath import *
from spatialmath.base import *
from assembly.assets.process_mesh import process_mesh
from assembly.examples.run_joint_plan import get_planner
from assembly.baselines.run_joint_plan import PyPlanner
from assembly.assets.subdivide import subdivide_to_size
import numpy as np
import json
import sys
import os
os.environ['OMP_NUM_THREADS'] = '1'

project_base_dir = os.path.abspath(os.path.join(
    os.path.dirname(os.path.abspath(__file__)), './')) + '/assembly/'


sys.path.append(project_base_dir)
sys.path.append(project_base_dir + '/baselines/')
sys.path.append(project_base_dir + '/assets/')


class FS:
    def readJSON(path: str):
        return json.loads((open(path)).read())

    def writeFile(data, filePath, fileName):

        file_to_open = filePath + fileName

        f = open(file_to_open, 'w', )

        f.write(data)

    def readFile(path: str):
        return open(path).read()

    def readFilesTypeFolder(pathFolder: str, fileType='.json'):
        return os.listdir(pathFolder)

    def readFolder(pathFolder: str):
        return list(map(lambda el: pathFolder + '/' + el,  os.listdir(pathFolder)))


def listGetFirstValue(iterable, default=False, pred=None):
    return next(filter(pred, iterable), default)


def filterModels(filterModels, filterModelsDescription):
    models = []
    for el in filterModelsDescription:
        models.append(listGetFirstValue(
            filterModels, None, lambda x:  x.name == el))
    return models


def main():
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument('--asp-path', type=str, required=True)
    args = parser.parse_args()
    aspDir = args.asp_dir

    # Коректировка пути до папки с генерацией ASP
    if (aspDir == None):
        args.print_helper()
    if (aspDir[aspDir.__len__() - 1] != '/'):
        aspDir += '/'
    # Получение списка папок с .obj обьектами
    assembles = FS.readFolder(aspDir + 'assembly')
    assemblyDirNormalize = []
    for el in assembles:
        try:
            # Пост обработка .obj обьектов
            process_mesh(source_dir=el, target_dir=el +
                         '/process/', subdivide=el, verbose=True)
            assemblyDirNormalize.append(el + '/process/')
        except Exception as e:
            print('ERRROR:')
            print(e)

    for el in assemblyDirNormalize:
        asset_folder = os.path.join(project_base_dir, aspDir)
        assembly_dir = os.path.join(asset_folder, el)
        planner = get_planner('bfs')(assembly_dir, assembly_dir, 0, [
            1], False, 'sdf', 0.05, 0.01, 100, 100, True)

        # Планирование пути
        status, t_plan, path = planner.plan(
            120, seed=1, return_path=True, render=False, record_path=None
        )
        coords = []

        for k in path:
            seMatrix = SE3(k)
            euler = seMatrix.eul()
            coord = seMatrix.A[0:3, 3]
            rot = Rotation.from_euler('xyz', euler, degrees=True).as_quat()
            coords.append({'quadrelion': [rot[0], rot[1], rot[2], rot[3]], 'xyz': [
                          coord[0], coord[1], coord[2]], 'euler': [euler[0], euler[1], euler[2]]})
        # Запись пути в кортеж
        planingObject = {
            "time": t_plan,
            "insertion_path": coords,
            "status": status,
        }
        # Запись результата планирования
        FS.writeFile(json.dumps(planingObject),
                     el[0:el.__len__() - 8], 'insertion_path.json')

        try:
            planner = PyPlanner(assembly_dir, 'process', still_ids=[1],)
            status, t_plan, path = planner.plan(
                planner_name=args.planner,
                step_size=args.step_size,
                max_time=args.max_time,
                seed=args.seed,
                return_path=True,
                simplify=args.simplify,
                render=args.render
            )

            print(f'Status: {status}, planning time: {t_plan}')

            if args.save_dir is not None:
                planner.save_path(path, args.save_dir, args.n_save_state)
        except Exception as e:
            print(e)


main()