blender_runner.py

from pathlib import Path
from typing import List

import bpy
import numpy as np
from PIL import Image

from .modifiers.mesh.mesh_utils import make_translate
from .modifiers.face_morph_creator import FaceMorphCreator
from .modifiers.material_creator import MaterialCreator
from .modifiers.mesh.armature_creator import ArmatureCreator
from .modifiers.mesh.mesh_decimator import MeshDecimator
from .modifiers.mesh.oral_cavity_creator import OralCavityMeshCreator
from .modifiers.position_corrector import PositionCorrector
from .utils.logger import log_info

file_path = str(Path(__file__).absolute())
current_dir_path = file_path[: file_path.rfind("/")]
assets_path = f"{current_dir_path}/assets"


class BlenderRunner:

    def __init__(self):
        eye_vertex_group_mask_path = f"{assets_path}/eye_vertex_group_mask.png"
        mouth_vertex_group_mask_path = f"{assets_path}/mouth_vertex_group_mask.png"
        eyes_shift_mask_path = f"{assets_path}/eyes_shift_mask.png"
        decimate_mask_path = f"{assets_path}/decimate_mask.png"
        eye_decimate_mask_path = f"{assets_path}/eye_decimate_mask.png"

        teeth_object_path = f"{assets_path}/teeth_and_tongue.fbx"
        armature_with_template_meshes_path = f"{assets_path}/armature_with_template_head.fbx"
        person_neck_path = f"{assets_path}/person_neck.fbx"

        self.eyes_shift_mask_image = Image.open(eyes_shift_mask_path).convert("L")
        self.decimate_mask_image = Image.open(decimate_mask_path)
        self.eye_decimate_mask_image = Image.open(eye_decimate_mask_path)

        self.eyes_shift_mask_array = np.array(self.eyes_shift_mask_image) / 255.0
        self.decimate_mask_array = np.array(self.decimate_mask_image)
        self.eye_decimate_mask_array = np.array(self.eye_decimate_mask_image)

        self.position_corrector = PositionCorrector()
        self.oral_cavity_creator = OralCavityMeshCreator(
            teeth_object_path=teeth_object_path,
        )
        self.morph_creator = FaceMorphCreator(
            eye_vertex_group_mask_path=eye_vertex_group_mask_path,
            mouth_vertex_group_mask_path=mouth_vertex_group_mask_path,
        )
        self.material_creator = MaterialCreator()
        self.mesh_decimator = MeshDecimator()
        self.armature_creator = ArmatureCreator(
            armature_with_template_meshes_path=armature_with_template_meshes_path,
            person_neck_path=person_neck_path,
        )

    def postprocess(
        self,
        is_generation_for_prod: bool,
        base_state_mesh_path: str,
        eyes_full_mesh_path: str,
        eyes_half_mesh_path: str,
        emotion_meshes_paths: List[str],
        texture_path: str,
        eyes_texture_path: str,
        debug_dir_path: str,
        output_file_path: str,
    ):
        log_info(f"Starting postprocessing mesh {base_state_mesh_path}")

        bpy.ops.wm.read_homefile(use_empty=True)

        meshes_files = [
            base_state_mesh_path,
            eyes_half_mesh_path,
            *emotion_meshes_paths,
        ]

        # import our files
        for filepath in meshes_files:
            bpy.ops.wm.obj_import(filepath=filepath)

        meshes_objects_names = list(map(
            lambda path: path[path.rfind("/") + 1 : path.rfind(".")][:63],
            meshes_files
        ))
        base_state_mesh_name = meshes_objects_names[0]
        eyes_mesh_name = meshes_objects_names[1]

        base_head_obj: bpy.types.Object = bpy.data.objects[base_state_mesh_name]
        eyes_obj: bpy.types.Object = bpy.data.objects[eyes_mesh_name]

        emotion_objects_names = meshes_objects_names[2:]
        emotion_objects: List[bpy.types.Object] = list(map(
            lambda name: bpy.data.objects[name],
            emotion_objects_names,
        ))

        base_obj_with_emotion_objects = [base_head_obj, *emotion_objects]

        self.position_corrector.correct_head_rotation(
            eyes_full_mesh_path=eyes_full_mesh_path,
            base_head_obj=base_head_obj,
            eyes_obj=eyes_obj,
            emotion_objects=emotion_objects,
        )

        # remove selection and select base state file again
        bpy.ops.object.select_all(action="DESELECT")
        bpy.context.view_layer.objects.active = base_head_obj

        # correct_eyes_position
        self.position_corrector.correct_eyes_position(
            eyes_obj=eyes_obj,
        )

        # add oral cavity with teeth and tongue to all head meshes
        self.oral_cavity_creator.add_oral_cavity(
            objects=base_obj_with_emotion_objects,
            debug_dir_path=debug_dir_path,
        )

        # create morphs
        self.morph_creator.create_morphs(
            base_obj=base_head_obj,
            morph_target_objects=emotion_objects,
        )

        # remove selection and select base state file again
        bpy.ops.object.select_all(action="DESELECT")
        bpy.context.view_layer.objects.active = base_head_obj

        base_head_obj.select_set(True)
        bpy.ops.export_scene.fbx(
            filepath=f"{debug_dir_path}/before_adjust_morph.fbx",
            use_selection=True,
            use_mesh_modifiers=True,
            object_types={"MESH"},
            path_mode="COPY",
            embed_textures=True,
            add_leaf_bones=False,
            bake_anim=False,
        )
        base_head_obj.select_set(False)

        # adjust morphs
        self.morph_creator.adjust_morphs(
            base_obj=base_head_obj,
        )

        # fix meshes' origins
        self.position_corrector.correct_origin(base_head_obj, eyes_obj)

        # add teeth and tongue
        teeth_obj = self.oral_cavity_creator.add_teeth_and_tongue(
            head_obj=base_head_obj,
            debug_file_path=f"{debug_dir_path}/after_adding_teeth.fbx",
        )

        # remove selection and select base state file again
        bpy.ops.object.select_all(action="DESELECT")
        bpy.context.view_layer.objects.active = base_head_obj

        # apply texture on head
        self.material_creator.apply_head_texture(base_head_obj, texture_path)

        # remove selection and select eyes object
        bpy.ops.object.select_all(action="DESELECT")
        bpy.context.view_layer.objects.active = eyes_obj

        # apply texture on eyes
        self.material_creator.apply_eyes_texture(eyes_obj, eyes_texture_path)

        base_head_obj.select_set(True)
        eyes_obj.select_set(True)
        teeth_obj.select_set(True)
        bpy.ops.export_scene.fbx(
            filepath=f"{debug_dir_path}/before_decimation.fbx",
            use_selection=True,
            use_mesh_modifiers=True,
            object_types={"MESH"},
            path_mode="COPY",
            embed_textures=True,
            add_leaf_bones=False,
            bake_anim=False,
        )
        base_head_obj.select_set(False)
        eyes_obj.select_set(False)
        teeth_obj.select_set(False)

        # apply modifiers on eyes
        if is_generation_for_prod:
            self.mesh_decimator.decimate_object_by_mask(
                obj=eyes_obj,
                mask_image=self.eye_decimate_mask_image,
                mask_array=self.eye_decimate_mask_array,
            )
            # self.__apply_eyes_modifiers(eyes_obj)

        # create vertex groups
        self.morph_creator.create_vertex_groups(base_head_obj)

        # set smooth shading
        self.material_creator.set_smooth_shading(
            head_obj=base_head_obj,
            eyes_obj=eyes_obj,
            teeth_obj=teeth_obj,
        )

        if is_generation_for_prod:
            # adjust scale to match boy head better
            self.armature_creator.adjust_scale(base_head_obj, eyes_obj, teeth_obj)

            # move head to the location of head bone
            self.armature_creator.change_location(base_head_obj, eyes_obj, teeth_obj)

            # change neck
            self.armature_creator.change_neck(base_head_obj)

        # optimize polygon number
        if is_generation_for_prod:
            self.mesh_decimator.decimate_object_by_mask(
                obj=base_head_obj,
                mask_image=self.decimate_mask_image,
                mask_array=self.decimate_mask_array,
            )

        if is_generation_for_prod:
            # attach head to armature and neck template
            armature_obj = self.armature_creator.attach_to_armature(
                base_head_obj,
                eyes_obj,
                teeth_obj,
            )
        else:
            base_head_obj.select_set(True)
            eyes_obj.select_set(True)
            teeth_obj.select_set(True)
            make_translate((0.0, 0.0, 0.2))
            bpy.ops.object.select_all(action="DESELECT")

        # export selection
        bpy.ops.object.select_all(action="DESELECT")
        base_head_obj.select_set(True)
        eyes_obj.select_set(True)
        teeth_obj.select_set(True)

        bpy.ops.object.transform_apply(
            location=True,
            rotation=True,
            scale=True,
        )

        if is_generation_for_prod:
            armature_obj.select_set(True)

        self.__rename_objects(
            base_head_obj,
            eyes_obj,
        )

        bpy.ops.export_scene.fbx(
            filepath=output_file_path,
            use_selection=True,
            use_mesh_modifiers=True,
            object_types={ "MESH", "ARMATURE" },
            path_mode="COPY",
            embed_textures=True,
            add_leaf_bones=False,
            bake_anim=False,
        )
        log_info(f"Finished postprocessing mesh {base_state_mesh_path}")

    def __apply_eyes_modifiers(
        self,
        eyes_obj: bpy.types.Object,
    ):
        log_info(f"Started applying eyes modifiers")

        decimate_modifier = eyes_obj.modifiers.new(name="Decimate", type="DECIMATE")
        decimate_modifier.decimate_type = "COLLAPSE"
        decimate_modifier.ratio = 0.4

        bpy.ops.object.modifier_apply(modifier=decimate_modifier.name)

        log_info("Finished applying eyes modifiers")

    def __rename_objects(
        self,
        head_obj: bpy.types.Object,
        eyes_obj: bpy.types.Object,
    ):
        head_name = "BoyHead"
        head_obj.name = head_name
        head_obj.data.name = head_name

        eyes_name = "Eyes"
        eyes_obj.name = eyes_name
        eyes_obj.data.name = eyes_name