framework/cg/pipeline/cg_pipeline.py

# -*- coding: utf-8 -*-
#!/usr/bin/env python
'''
DESCRIPTION.
CG asset (models) creation pipelines from FreeCAD data.
Support Blender compiled as a Python Module only!
'''
__version__ = '0.8'
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 import (hierarchy_assembly,
                                                   hierarchy_separated_parts,
                                                   hierarchy_mono_part)
from blender.processing.render_assets import setup_meshes
from blender.processing.midpoly_setup import hightpoly_collections_to_midpoly
from blender.processing.visual_assets 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
from blender.export.ply import export_ply
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 = {}
"""
    'parts_col_name': 'Parts',
    'midpoly_col_name': 'Midpoly',
    'visual_col_name': 'Lowpoly',
    'lcs_inlet': 'in',
    'lcs_outlet': 'out',
    'lcs_root': 'root',
    'render': 'hp',
    'midpoly': 'mp',
    'visual': 'lp',
    'render': 'render',
}
"""

defined_pipeline_list = ['cad', 'render', 'visual', 'collision', 'baking', 'export']


def cg_pipeline(**kwargs):
    ''' CG assets (models) creation pipeline '''
    # set defaults
    # ______________________________
    assembly_name = (
        kwargs['fcstd_path'].replace('\\', '/').rpartition('/')[2].rpartition('.')[0])
    # clear other paths from kwargs (freecad issue)
    parts_sequence_path = kwargs.pop('parts_sequence_path', None)
    blend_path = kwargs.pop('blend_path', None)
    export_path = kwargs.pop('export_path', None)

    if kwargs['pipeline_type'] not in defined_pipeline_list:
        return logger.error('Pipeline type %s is not defined!', kwargs['pipeline_type'])

    # 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, 'models').replace('\\', '/')
    os.makedirs(blend_path, exist_ok=True)
    #blend_file = os.path.join(blend_path, f'{assembly_name}.blend').replace('\\', '/')
    remove_collections_with_objects()
    cleanup_orphan_data()

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

    # Save original cad setup as blender scene
    if kwargs['pipeline_type'] == 'cad':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        return blend_file

    # 2) prepare render (depend of cad_objects['objs_lcs'] and parts_sequence)
    # ______________________________
    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)
    else:
        parts_sequence = None
    lcs_inlets = [
        lcs_name for lcs_name in cad_objects['objs_lcs']
        if (bpy.data.objects[lcs_name].get('Robossembler_SocketFlow') == 'inlet'
            and bpy.data.objects[lcs_name].get('Robossembler_DefaultOrigin'))]

    # input cases
    # 1 case
    if parts_sequence and len(lcs_inlets) > 1:
        logger.info('Parts assembling sequence and LCS points found! '
                    'Launch "hierarchy_assembly" restructuring pipeline.')
        part_names = hierarchy_assembly(
            cad_objects['objs_lcs'], parts_sequence)
    # 2 case
    elif parts_sequence and len(lcs_inlets) < 2:
        return logger.error('Assembly do not have enough LCS points!')
    # 3 case
    elif not parts_sequence and lcs_inlets:
        logger.info('Parts assembling sequence not found! '
                    'Launch "hierarchy_separated_parts" restructuring pipeline.')
        part_names = hierarchy_separated_parts(
            cad_objects['objs_lcs'])
    # 4 case
    elif not parts_sequence and not lcs_inlets:
        logger.info('Parts assembling sequence and LCS points not found! '
                    'Launch "hierarchy_mono_part" restructuring pipeline.')
        part_names = hierarchy_mono_part()

    if not part_names:
        return logger.error('Can not generate parts!')

    # setup renders with materials only
    if cad_objects['objs_foreground']:
        setup_meshes(cad_objects['objs_foreground'],
                    sharpness=True, shading=True)
    # setup all renders
    else:
        setup_meshes(cad_objects['objs_background'],
                    sharpness=True, shading=True)

    # Save render setup as blender scene
    if kwargs['pipeline_type'] == 'render':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        logger.info('%s original hightpoly collections ready!', len(part_names))
        return blend_file

    # 3) prepare midpoly
    # ______________________________
    tmp_col_name = copy_collections_recursive(
        bpy.data.collections[cg_config['parts_col_name']],
        suffix=cg_config['midpoly']
    )
    midpoly_obj_names = hightpoly_collections_to_midpoly(
        tmp_col_name, part_names, **cg_config)

    # Save midpoly setup as blender scene
    if kwargs['pipeline_type'] == 'midpoly':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        logger.info('%s midpoly objects ready!', len(midpoly_obj_names))
        return blend_file

    # 4) prepare visual
    # ______________________________
    visual_obj_names = parts_to_shells(part_names, **cg_config)
    uv_unwrap(visual_obj_names)

    # Save visual setup as blender scene
    if kwargs['pipeline_type'] == 'visual':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        logger.info('%s visual objects ready!', len(visual_obj_names))
        return blend_file

    # 5) bake textures
    # ______________________________
    if kwargs['textures_resolution'] == 0:
        logger.info('Baking pipeline has been canceled!')
    else:
        textures_path = os.path.join(blend_path, 'textures').replace('\\', '/')
        bake_paths = bw_bake(visual_obj_names,
                            textures_path,
                            kwargs['textures_resolution'],
                            **cg_config)
        # Save baking result
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, 'baking')
        ).replace('\\', '/')
        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(visual_obj_names, textures_path)
        bpy.ops.file.make_paths_relative()

    # Save baking setup as blender scene
    if kwargs['pipeline_type'] == 'baking':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        logger.info('%s visual objects baked!', len(visual_obj_names))
        return blend_file

    # 6 export object meshes
    # ______________________________
    # TODO asset manager
    # Save blender scene
    if kwargs['pipeline_type'] == 'export':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)

    for part_name in part_names:
        export_fbx(
            obj_name=f'{part_name}_{cg_config["visual"]}',
            path=os.path.join(export_path, part_name, 'meshes').replace('\\', '/'))
        export_ply(
            obj_name=f'{part_name}_{cg_config["midpoly"]}',
            path=os.path.join(export_path, part_name, 'meshes').replace('\\', '/'))
        export_dae(
            obj_name=f'{part_name}_{cg_config["midpoly"]}',
            path=os.path.join(export_path, part_name, 'meshes').replace('\\', '/'))
        export_stl(
            obj_name=f'{part_name}_{cg_config["visual"]}',
            path=os.path.join(export_path, part_name, 'meshes').replace('\\', '/'))

    logger.info('%s parts exported!', len(part_names))
    return blend_file


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.01,
        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(
        '--pipeline_type',
        type=str,
        help='Set pipeline type: "cad", "render","visual", "collision", "baking", "export"',
        default='export',
        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)