framework/cg/pipeline/cg_pipeline.py

# -*- coding: utf-8 -*-
#!/usr/bin/env python
'''
DESCRIPTION.
Convert and setup scene from FreeCAD data.
Support Blender compiled as a Python Module only!
'''
__version__ = '0.7'
import collections
import json
import logging
import shutil
import os
from itertools import zip_longest

from blender.utils.cleanup_orphan_data import cleanup_orphan_data
from blender.utils.collection_tools import (copy_collections_recursive,
                                            remove_collections_with_objects)
from utils.cmd_proc import cmd_proc
from blender.import_cad.build_blender_scene import json_to_blend
from blender.processing.restruct_hierarchy_by_lcs import restruct_hierarchy
from blender.processing.highpoly_setup import setup_meshes
from blender.processing.midpoly_setup import hightpoly_collections_to_midpoly
from blender.processing.lowpoly_setup import parts_to_shells
from blender.processing.uv_setup import uv_unwrap
from blender.texturing.bake_submitter import bw_bake
from blender.texturing.composing import compose_baked_textures
from blender.texturing.shading import assign_pbr_material
from blender.export.dae import export_dae
from blender.export.stl import export_stl
from blender.export.fbx import export_fbx
import bpy

# TODO Path
freecad_to_json_script = 'freecad_to_json.py'

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

# TODO NIX
freecad_bin = 'freecadcmd'
# TODO NIX BakwWrangler/nodes/node_tree.py:659
blender_bin = 'blender'

# TODO WEBAPP
cg_config = {
    'lcs_col_name': 'LCS',
    'parts_col_name': 'Parts',
    'midpoly_col_name': 'Midpoly',
    'lowpoly_col_name': 'Lowpoly',
    'lcs_inlet': 'in',
    'lcs_outlet': 'out',
    'lcs_root': 'root',
    'hightpoly': 'hp',
    'midpoly': 'mp',
    'lowpoly': 'lp',
    'render': 'render',
}


def cg_pipeline(**kwargs):
    ''' CG asset creation pipeline '''
    assembly_name = kwargs['fcstd_path'].rpartition('/')[2].rpartition('.')[0]
    # freecad don't like other paths
    parts_sequence_path = kwargs.pop('parts_sequence_path', None)
    blend_path = kwargs.pop('blend_path', None)
    export_path = kwargs.pop('export_path', None)

    # for eatch sequence
    parts_sequence = None
    if parts_sequence_path and os.path.isfile(parts_sequence_path):
        with open(parts_sequence_path, 'r', encoding='utf-8') as sequence_file:
            parts_sequence = json.load(sequence_file)
    # output file management
    if not blend_path:
        blend_path = kwargs['fcstd_path'].replace('\\', '/').rpartition('/')[0]
    if not export_path:
        export_path = os.path.join(blend_path, 'assets').replace('\\', '/')
    os.makedirs(blend_path, exist_ok=True)

    # prepare blend file
    blend_file = os.path.join(blend_path, f'{assembly_name}.blend').replace('\\', '/')
    remove_collections_with_objects()
    cleanup_orphan_data()

    # 0 сonvert FreeCAD scene to Blender scene
    imported_objects = json_to_blend(
        json.loads(
            cmd_proc(freecad_bin,
                    freecad_to_json_script,
                    '--',
                    **kwargs
                    ).split('FreeCAD ')[0]
        ), **cg_config
    )

    # Save original freecad setup as blender scene
    bpy.ops.wm.save_as_mainfile(filepath=f'{blend_file.rpartition(".")[0]}_orig.blend')

    # 1 prepare highpoly
    part_names = None
    lcs_pipeline = True
    if imported_objects['objs_lcs']:
        part_names = restruct_hierarchy(
            imported_objects['objs_lcs'], parts_sequence, **cg_config)

    # non lcs pipeline
    if not part_names:
        lcs_pipeline = False
        part_names = [[obj for obj in bpy.data.objects if not obj.parent][0].name]

    if imported_objects['objs_foreground']:
        setup_meshes(imported_objects['objs_foreground'],
                    sharpness=True, shading=True)
    else:
        setup_meshes(imported_objects['objs_background'],
                    sharpness=True, shading=True)

    # 2 prepare midpoly
    copy_col_name = copy_collections_recursive(
        bpy.data.collections[cg_config['parts_col_name']],
        suffix=cg_config['midpoly']
    )
    midpoly_obj_names = hightpoly_collections_to_midpoly(
        copy_col_name, part_names, lcs_pipeline, **cg_config)

    # 3 prepare lowpoly
    lowpoly_obj_names = parts_to_shells(part_names, lcs_pipeline, **cg_config)
    uv_unwrap(lowpoly_obj_names)

    # Save before baking
    bpy.ops.wm.save_as_mainfile(filepath=blend_file)

    # 4 bake textures
    if kwargs['textures_resolution'] != 0:
        textures_path = os.path.join(blend_path, 'textures').replace('\\', '/')
        bake_paths = bw_bake(lowpoly_obj_names,
                            textures_path,
                            kwargs['textures_resolution'],
                            **cg_config)
        # Save baking result
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        # 5 prepare textures
        bpy.ops.wm.quit_blender()
        compose_baked_textures(textures_path, bake_paths, kwargs['textures_resolution'])
        for bake_path in bake_paths:
            shutil.rmtree(bake_path)
        bpy.ops.wm.open_mainfile(filepath=blend_file)
        assign_pbr_material(lowpoly_obj_names, textures_path)
        bpy.ops.file.make_paths_relative()
        # Save assigned lowpoly assets
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        for part_name in part_names:
            export_fbx(
                obj_name=f'{part_name}_{cg_config["lowpoly"]}',
                path=export_path,
                subdir='fbx')


    # 6 export object meshes and urdf
    to_urdf = collections.defaultdict(list)
    link = {}
    for part_name in part_names:
        link_prop = {}
        link_prop['visual'] = export_dae(
            obj_name=f'{part_name}_{cg_config["midpoly"]}',
            path=export_path,
            subdir='dae')
        link_prop['collision'] = export_stl(
            obj_name=f'{part_name}_{cg_config["lowpoly"]}',
            path=export_path,
            subdir='collision')
        link[part_name] = link_prop
    to_urdf['links'].append(link)

    # TODO export urdf
    config = kwargs.pop('config', None)
    # config = {'sequences': [['cube1', 'cube2', 'cube3', 'cube4'], ['cube2', 'cube1', 'cube4', 'cube3']]}
    if config:
        for sequence in config['sequences']:
            joint = {}
            # TODO collect pairs 0_1, 1_2, 2_3, 3_4, ...
            for pair in zip_longest(sequence[0::2], sequence[1::2]):
                joint_prop = {}
                if pair[1]:
                    joint_prop['type'] = 'fixed'
                    location = list(bpy.data.objects.get(pair[1]).location)
                    rotation = list(bpy.data.objects.get(pair[0]).rotation_euler)
                    # origin
                    # round location values
                    for idx, axis in enumerate(location):
                        location[idx] = round(axis, 5)
                    joint_prop['location'] = location
                    joint_prop['rotation'] = rotation
                    # parent
                    joint_prop['parent'] = pair[0]
                    # child
                    joint_prop['child'] = pair[1]

                    joint['_'.join(pair)] = joint_prop

            to_urdf['sequence'].append(joint)

    print(json.dumps(to_urdf, indent=4))

    logger.info('%s original hightpoly collections ready!', len(part_names))
    logger.info('%s midpoly objects ready!', len(midpoly_obj_names))
    logger.info('%s lowpoly objects ready!', len(lowpoly_obj_names))

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(
         description='Convert and setup FreeCAD solid objects to 3d assets mesh files.'
    )
    parser.add_argument(
        '--fcstd_path',
        type=str,
        help='Path to source FreeCAD scene',
        required=True
    )
    parser.add_argument(
        '--tesselation_method',
        type=str,
        help='Select tesselation method: Standard or FEM.',
        default='Standard',
        required=False
    )
    parser.add_argument(
        '--linear_deflection',
        type=float,
        help='Max linear distance error',
        default=0.1,
        required=False
    )
    parser.add_argument(
        '--angular_deflection',
        type=float,
        help='Max angular distance error',
        default=20.0,
        required=False
    )
    parser.add_argument(
        '--fem_size',
        type=float,
        help='For FEM method only! Finite element size in mm',
        default=50.0,
        required=False
    )
    parser.add_argument(
        '--skiphidden',
        type=bool,
        help='Skip processing for hidden FreeCAD objects',
        default=True,
        required=False
    )
    parser.add_argument(
        '--property_forse_nonsolid',
        type=str,
        help='FreeCAD property to enable processing for nonsolid objects',
        default='Robossembler_NonSolid',
        required=False
    )
    parser.add_argument(
        '--parts_sequence_path',
        type=str, help='Path to parts assembling sequence json file.',
        required=False
    )
    parser.add_argument(
        '--export_path',
        type=str, help='Path for export assets. If not, fcstd_path will be used instead.',
        required=False
    )
    parser.add_argument(
        '--blend_path',
        type=str,
        help='Path for blender scene. If not, fcstd_path will be used instead.',
        required=False
    )
    parser.add_argument(
        '--textures_resolution',
        type=int,
        help='Set baking texture resolution. Recomended - 4096 pix ',
        default=512,
        required=False
    )

    args = parser.parse_args()

    cg_kwargs = {key: getattr(args, key) for key in dir(args) if not key.startswith('_')}

    cg_pipeline(**cg_kwargs)

    logger.info('CG Pipeline Completed!')