ap-canvas-creation-module-old/03_blender/sd_blender/__init__.py

# ----------------------------------------------------------
# File __init__.py
# ----------------------------------------------------------

# Imports
# Check if this add-on is being reloaded
if "bpy" in locals():
    # reloading .py files
    import bpy
    import importlib

    # from . import zs_renderscene as zsrs
    # importlib.reload(zsrs)


# or if this is the first load of this add-on
else:
    import bpy
    import json
    from pathlib import Path
    from dataclasses import dataclass
    from mathutils import Euler
    import math
    import os

    # from . import zs_renderscene as zsrs  # noqa


# Addon info
bl_info = {
    "name": "ZS Stable Diffusion Connection V0.0.1",
    "author": "Sergiu <sergiu@zixelise.com>",
    "Version": (0, 0, 1),
    "blender": (4, 00, 0),
    "category": "Scene",
    "description": "Stable Diffusion Connection",
}


# -------------------------------------------------------------------
# Functions
# -------------------------------------------------------------------


@dataclass
class AssetData:
    path: str
    collection_name: str
    type: str


product_asset = AssetData(
    path="sample_scene/01_Products", collection_name="01_Products", type="product"
)
element_asset = AssetData(
    path="sample_scene/02_Elements", collection_name="02_Elements", type="asset"
)
basic_shapes_asset = AssetData(
    path="sample_scene/03_BasicShapes",
    collection_name="03_BasicShapes",
    type="shape",
)


def load_scene():
    print("Loading Scene")
    # load scene data
    scene_data = load_scene_data()
    # create parent collections
    create_parent_collections(product_asset.collection_name)
    create_parent_collections(element_asset.collection_name)
    create_parent_collections(basic_shapes_asset.collection_name)
    create_parent_collections("05_Cameras")
    # append products
    products_data = load_objects_data(scene_data, product_asset.type)
    for index, product in enumerate(products_data):
        append_objects(product, index, product_asset)
    # append elements
    elements_data = load_objects_data(scene_data, element_asset.type)
    for index, product in enumerate(elements_data):
        append_objects(product, index, element_asset)
    # append shapes
    basic_shapes_data = load_objects_data(scene_data, basic_shapes_asset.type)
    for index, product in enumerate(basic_shapes_data):
        append_objects(product, index, basic_shapes_asset)

    # set lighting
    set_environment(scene_data)

    # set camera
    create_cameras(scene_data)

    # rename outputs

    set_output_paths(
        "D://Git//ap-canvas-creation-module//03_blender//sd_blender//sample_scene//Renders",
        scene_data["project_id"],
    )

    # setup compositing
    set_cryptomatte_objects("01_Products", "mask_product")


def load_scene_data():
    print("Loading Scene Data")
    # load scene data

    # to be replaced with actual data
    # open scene_info.json
    script_path = Path(__file__).resolve()
    scene_data_path = script_path.parent / "sample_scene" / "scene_info.json"
    with scene_data_path.open() as file:
        scene_data = json.load(file)
        print(scene_data)
    return scene_data


def load_objects_data(scene_data, object_type: str):
    print("Loading Assets Data")
    # load assets data
    objects_data = []
    for object in scene_data["scene"]["objects"]:
        if object["group_type"] == object_type:
            # get additional object data by id and combine with object data
            object_data = get_object_data_by_id(object["id"])
            object.update(object_data)
            objects_data.append(object)
    return objects_data

    # to be replaced with actual data


def get_object_data_by_id(object_id: str):
    print("Getting Object Data")
    # open assets_database.json
    script_path = Path(__file__).resolve()
    assets_data_path = script_path.parent / "sample_scene" / "assets_database.json"

    with assets_data_path.open() as file:
        assets_data = json.load(file)
    # get object data by id
    for object in assets_data:
        if object["id"] == object_id:
            return object


def get_hdri_data_by_id(object_id: str):
    print("Getting HDRI Data")
    # open assets_database.json
    script_path = Path(__file__).resolve()
    assets_data_path = script_path.parent / "sample_scene" / "lighting_database.json"

    with assets_data_path.open() as file:
        assets_data = json.load(file)
    # get object data by id
    for object in assets_data:
        if object["id"] == object_id:
            return object


def append_objects(asset_info, index, asset_data: AssetData):
    print("Appending Objects")

    blendFileName = asset_info["name"]

    # visibleLayersJSONName = productInfo["Version"]

    # activeSwitchMaterials = json.dumps(productInfo["ActiveMaterials"])

    collectionName = blendFileName + "_" + str(index)

    append_active_layers(collectionName, asset_info, asset_data)

    # replace_switch_materials(shotInfo, productInfo["ActiveMaterials"])

    link_collection_to_collection(
        asset_data.collection_name, bpy.data.collections[collectionName]
    )


def append_active_layers(newCollectionName, product_info, asset_data: AssetData):
    utility_collections = [
        "MaterialLibrary",
        "Animation_Controller",
        "Animation_Rig",
        "Animation_Target",
    ]

    # visible_layers = utility_collections + product_info["VisibleLayers"]
    visible_layers = utility_collections

    script_path = Path(__file__).resolve().parent
    filePath = str(
        script_path
        / asset_data.path
        / product_info["name"]
        / "BLEND"
        / (product_info["name"] + ".blend")
    )

    print(filePath)

    # delete all objects and collections from product collection

    create_parent_collections(newCollectionName)

    # append active collections
    with bpy.data.libraries.load(filePath) as (data_from, data_to):
        data_to.collections = []
        for name in data_from.collections:

            if name in visible_layers:
                data_to.collections.append(name)
            if "NonConfigurable" in name:
                data_to.collections.append(name)

    # link appended colections to newCollection
    for collection in data_to.collections:
        # try:
        #     bpy.context.scene.collection.children.link(collection)
        # except:
        #     print(collection)
        link_collection_to_collection(newCollectionName, collection)

        # hide utility collections
        for utilityCollectionName in utility_collections:
            if utilityCollectionName in collection.name:
                # rename utility collections
                collection.name = newCollectionName + "_" + utilityCollectionName
                hide_collection(collection)

            # need to redo this in the future
            if "Animation_Target" in collection.name:
                # print object name from collection

                collection.objects[0].location = product_info["properties"][
                    "transform"
                ]["position"]

                # collection.objects[0].rotation_euler = product_info["properties"][
                #     "transform"
                # ]["rotation"]

                rotation_in_degrees = product_info["properties"]["transform"][
                    "rotation"
                ]
                rotation_in_radians = [math.radians(deg) for deg in rotation_in_degrees]

                collection.objects[0].rotation_euler = rotation_in_radians

                collection.objects[0].scale = product_info["properties"]["transform"][
                    "scale"
                ]

    # make all objects in collection local
    for obj in bpy.data.collections[newCollectionName].all_objects:
        if obj.type == "MESH":
            obj.make_local()
            obj.data.make_local()

    # remove duplicated material slots
    mats = bpy.data.materials

    for obj in bpy.data.collections[newCollectionName].all_objects:
        for slot in obj.material_slots:
            part = slot.name.rpartition(".")
            if part[2].isnumeric() and part[0] in mats:
                slot.material = mats.get(part[0])


def create_parent_collections(group_name: str):

    if collection_exists(group_name):

        remove_collection_and_objects(group_name)
    else:
        create_collection(group_name)


def set_environment(scene_data):
    # Find the group named NG_Canvas_Background in the Blender world
    world = bpy.context.scene.world

    if world is not None:
        # Get the node group named NG_Canvas_Background
        node_group = world.node_tree.nodes.get("NG_Canvas_Background")

        if node_group is not None:
            # Set the group's properties from scene_data
            node_group.inputs[0].default_value = float(
                scene_data["scene"]["environment"]["lighting"]["rotation"]
            )
            node_group.inputs[1].default_value = float(
                scene_data["scene"]["environment"]["lighting"]["intensity"]
            )

            if scene_data["scene"]["environment"]["lighting"]["visible"] == True:
                node_group.inputs[2].default_value = 1.0
            else:
                node_group.inputs[2].default_value = 0.0

            node_group.inputs[3].default_value = (
                scene_data["scene"]["environment"]["background"]["color"]["r"],
                scene_data["scene"]["environment"]["background"]["color"]["g"],
                scene_data["scene"]["environment"]["background"]["color"]["b"],
                1.0,
            )

            hdri_data = get_hdri_data_by_id(
                scene_data["scene"]["environment"]["lighting"]["id"]
            )

            # go inside the node group, and find the image texture node, and set the image to the one from the scene data
            for node in node_group.node_tree.nodes:
                if node.name == "environment_map":
                    node.image = bpy.data.images.load(
                        str(
                            Path(__file__).resolve().parent
                            / "sample_scene"
                            / "05_Lighting"
                            / "HDRI"
                            / (hdri_data["name"] + ".exr")
                        )
                    )
        else:
            print("Group 'NG_Canvas_Background' not found")


def create_cameras(scene_data):
    # # Get the 05_Cameras collection, or create it if it doesn't exist
    collection_name = "05_Cameras"
    if collection_name in bpy.data.collections:
        collection = bpy.data.collections[collection_name]
    else:
        return

    for camera_data in scene_data["scene"]["cameras"]:
        # Create a new camera object
        bpy.ops.object.camera_add()

        # Get the newly created camera
        camera = bpy.context.object

        # Set the camera's name
        camera.name = camera_data["name"]

        # Set the camera's position
        position = camera_data["properties"]["transform"]["position"]
        camera.location.x = position[0]
        camera.location.y = position[1]
        camera.location.z = position[2]

        # Set the camera's rotation
        rotation = camera_data["properties"]["transform"]["rotation"]
        # Convert the rotation from degrees to radians
        rotation = [math.radians(r) for r in rotation]
        camera.rotation_euler = Euler(rotation, "XYZ")

        # Set the camera's lens properties
        lens = camera_data["properties"]["lens"]
        type_mapping = {
            "PERSPECTIVE": "PERSP",
            "ORTHOGRAPHIC": "ORTHO",
            "PANORAMIC": "PANO",
        }
        camera.data.type = type_mapping.get(lens["type"].upper(), "PERSP")
        camera.data.angle = math.radians(lens["fov"])
        camera.data.clip_start = lens["near"]
        camera.data.clip_end = lens["far"]

        # Add the camera to the 05_Cameras collection
        collection.objects.link(camera)
        bpy.context.scene.collection.objects.unlink(camera)

        # Set the camera as the active camera if "active" is true
        if camera_data["properties"]["active"]:
            bpy.context.scene.camera = camera


def set_output_paths(base_path, project_name):
    # Get the current scene
    scene = bpy.context.scene

    # Check if the scene has a compositor node tree
    if scene.node_tree:
        # Iterate over all nodes in the node tree
        for node in scene.node_tree.nodes:
            # Check if the node is an output node
            if node.type == "OUTPUT_FILE":
                # Set the base path of the output node
                node.base_path = base_path + "//" + project_name
                # Iterate over all file slots of the output node
                # for file_slot in node.file_slots:
                #     # Set the path of the file slot
                #     file_slot.path = project_name


def set_cryptomatte_objects(collection_name, node_name):
    # Get all objects in the specified collection
    objects = bpy.data.collections[collection_name].all_objects

    # Convert the objects to a list
    object_names = [obj.name for obj in objects]

    print(object_names)

    # Get the current scene
    scene = bpy.context.scene

    # Check if the scene has a compositor node tree
    if scene.node_tree:
        # Iterate over all nodes in the node tree
        for node in scene.node_tree.nodes:
            # Check if the node is a Cryptomatte node with the specified name
            if node.type == "CRYPTOMATTE_V2" and node.name == node_name:
                # Set the Matte ID objects of the node
                node.matte_id = ",".join(object_names)


# -------------------------------------------------------------------
# Utilities
# -------------------------------------------------------------------


def remove_collection_and_objects(collection_name):
    oldObjects = list(bpy.data.collections[collection_name].all_objects)
    for obj in oldObjects:
        bpy.data.objects.remove(obj, do_unlink=True)

    old_collection = bpy.data.collections[collection_name]
    if old_collection is not None:
        old_collection_names = get_subcollection_names(old_collection)
    else:
        print("Collection not found.")

    # print line break
    print("-----------------------------------------------------------------")
    print(old_collection_names)
    print("-----------------------------------------------------------------")
    for old_collection_name in old_collection_names:
        for collection in bpy.data.collections:
            if collection.name == old_collection_name:
                bpy.data.collections.remove(collection)

    bpy.ops.outliner.orphans_purge(
        do_local_ids=True, do_linked_ids=True, do_recursive=True
    )


def get_subcollection_names(collection):
    subcollection_names = []

    for child in collection.children:
        subcollection_names.append(child.name)
        subcollection_names.extend(get_subcollection_names(child))

    return subcollection_names


def select_objects_in_collection(collection):
    """Recursively select all objects in the given collection and its subcollections."""
    for obj in collection.objects:
        obj.select_set(True)
    for subcollection in collection.children:
        select_objects_in_collection(subcollection)


def link_collection_to_collection(parentCollectionName, childCollection):

    if bpy.context.scene.collection.children:
        parentCollection = bpy.context.scene.collection.children.get(
            parentCollectionName
        )
        # Add it to the main collection

        # childCollection = bpy.context.scene.collection.children.get(childCollection)

        parentCollection.children.link(childCollection)
        # if child collection is in scene collection unlink it
        if bpy.context.scene.collection.children.get(childCollection.name):
            bpy.context.scene.collection.children.unlink(childCollection)
        # bpy.context.scene.collection.children.unlink(childCollection)


# link collection to collection
def link_collection_to_collection_old(parentCollectionName, childCollection):

    if bpy.context.scene.collection.children:
        parentCollection = bpy.context.scene.collection.children.get(
            parentCollectionName
        )
        # Add it to the main collection

        try:
            childCollection = bpy.data.collections[childCollection]
        except:
            print("Collection not found.")
            return

        parentCollection.children.link(childCollection)
        bpy.context.scene.collection.children.unlink(childCollection)


# function that checks if a collection exists
def collection_exists(collection_name):
    return collection_name in bpy.data.collections


# function that creates a new collection and adds it to the scene
def create_collection(collection_name):
    new_collection = bpy.data.collections.new(collection_name)
    bpy.context.scene.collection.children.link(new_collection)


def hide_collection(collection):
    collection.hide_render = True
    collection.hide_viewport = True


def check_if_selected_objects_have_parent(self):
    for obj in bpy.context.selected_objects:
        if obj.parent is None:
            message = f"Object {obj.name} has no parent"
            self.report({"ERROR"}, message)
            return False
    return True


def add_rig_controller_to_selection(self):
    # add "Rig_Controller_Main" to the selection
    has_controller_object = False
    for obj in bpy.data.objects:
        if obj.name == "Rig_Controller_Main":
            if obj.hide_viewport:
                self.report({"ERROR"}, "Rig_Controller_Main is hidden")
                return has_controller_object
            obj.select_set(True)
            # if object is not visible, make it visible

            has_controller_object = True
            return has_controller_object

    if not has_controller_object:
        message = f"Rig_Controller_Main not found"
        self.report({"ERROR"}, message)
        print("Rig_Controller_Main not found")

    return has_controller_object


def select_objects_in_collection(collection):
    """Recursively select all objects in the given collection and its subcollections."""
    for obj in collection.objects:
        obj.select_set(True)
    for subcollection in collection.children:
        select_objects_in_collection(subcollection)


def check_if_object_has_principled_material(obj):
    for slot in obj.material_slots:

        # if more than one slot is principled, return
        for node in slot.material.node_tree.nodes:
            # print(node.type)
            if node.type == "BSDF_PRINCIPLED":

                return True
            else:
                print("Object has no principled material", obj.name)
                return False

    return True


def unhide_all_objects():
    # show all objects using operator

    bpy.ops.object.hide_view_clear()


# -------------------------------------------------------------------
# Scene optimization
# -------------------------------------------------------------------
def export_non_configurable_to_fbx():
    # Get the current .blend file path

    blend_filepath = bpy.data.filepath
    if not blend_filepath:
        print("Save the .blend file first.")
        return

    # Get the parent directory of the .blend file
    blend_dir = os.path.dirname(blend_filepath)
    parent_dir = os.path.dirname(blend_dir)
    blend_filename = os.path.splitext(os.path.basename(blend_filepath))[0]

    # Create the FBX export path
    fbx_export_path = os.path.join(parent_dir, "FBX", f"{blend_filename}.fbx")

    # Ensure the FBX directory exists
    os.makedirs(os.path.dirname(fbx_export_path), exist_ok=True)

    # Deselect all objects
    bpy.ops.object.select_all(action="DESELECT")

    # Select all objects in the NonConfigurable collection
    collection_name = "NonConfigurable"
    if collection_name in bpy.data.collections:
        collection = bpy.data.collections[collection_name]
        for obj in collection.objects:
            obj.select_set(True)
    else:
        print(f"Collection '{collection_name}' not found.")
        return


def export_scene_to_fbx(self):
    # Ensure the .blend file is saved
    if not bpy.data.is_saved:
        print("Save the .blend file first.")
        self.report({"ERROR"}, "Save the .blend file first.")
        return

    # Get the current .blend file path
    blend_filepath = bpy.data.filepath
    if not blend_filepath:
        print("Unable to get the .blend file path.")
        return

    # Get the parent directory of the .blend file
    blend_dir = os.path.dirname(blend_filepath)
    parent_dir = os.path.dirname(blend_dir)
    blend_filename = os.path.splitext(os.path.basename(blend_filepath))[0]

    # Create the FBX export path
    fbx_export_path = os.path.join(parent_dir, "FBX", f"{blend_filename}.fbx")

    # Ensure the FBX directory exists
    os.makedirs(os.path.dirname(fbx_export_path), exist_ok=True)

    # unhide all objects
    unhide_all_objects()
    # Deselect all objects
    bpy.ops.object.select_all(action="DESELECT")

    # Select all objects in the NonConfigurable collection and its subcollections
    collection_name = "NonConfigurable"
    if collection_name in bpy.data.collections:
        collection = bpy.data.collections[collection_name]
        select_objects_in_collection(collection)
    else:
        print(f"Collection '{collection_name}' not found.")
        return

        # check if all objects selected have a parent, if not return

    if not check_if_selected_objects_have_parent(self):
        return

    if not add_rig_controller_to_selection(self):
        return

    # Export selected objects to FBX
    bpy.ops.export_scene.fbx(
        filepath=fbx_export_path,
        use_selection=True,
        global_scale=1.0,
        apply_unit_scale=True,
        bake_space_transform=False,
        object_types={"MESH", "ARMATURE", "EMPTY"},
        use_mesh_modifiers=True,
        mesh_smooth_type="FACE",
        use_custom_props=True,
        # bake_anim=False,
    )

    print(f"Exported to {fbx_export_path}")


def export_scene_to_glb(self):
    # Ensure the .blend file is saved
    if not bpy.data.is_saved:
        print("Save the .blend file first.")
        self.report({"ERROR"}, "Save the .blend file first.")
        return

    # Get the current .blend file path
    blend_filepath = bpy.data.filepath
    if not blend_filepath:
        print("Unable to get the .blend file path.")
        return

    # Get the parent directory of the .blend file
    blend_dir = os.path.dirname(blend_filepath)
    parent_dir = os.path.dirname(blend_dir)
    blend_filename = os.path.splitext(os.path.basename(blend_filepath))[0]

    # Create the GLB export path
    glb_export_path = os.path.join(parent_dir, "WEB", f"{blend_filename}.glb")

    # Ensure the GLB directory exists
    os.makedirs(os.path.dirname(glb_export_path), exist_ok=True)

    # unhide all objects
    unhide_all_objects()
    # Deselect all objects
    bpy.ops.object.select_all(action="DESELECT")

    # Select all objects in the NonConfigurable collection and its subcollections
    collection_name = "WebGL"
    if collection_name in bpy.data.collections:
        collection = bpy.data.collections[collection_name]
        select_objects_in_collection(collection)
    else:
        print(f"Collection '{collection_name}' not found.")
        return

    if not check_if_selected_objects_have_parent(self):
        return
    # check if all objects selected have a parent, if not return
    if not add_rig_controller_to_selection(self):
        return

    # # for each selected objects, check if the the material is principled, if not return
    # for obj in bpy.context.selected_objects:
    #     if not check_if_object_has_principled_material(obj):
    #         message = f"Object {obj.name} has no principled material"
    #         self.report({"ERROR"}, message)
    #         return

    # Export selected objects to GLB
    bpy.ops.export_scene.gltf(
        filepath=glb_export_path,
        export_format="GLB",
        use_selection=True,
        export_apply=True,
        export_animations=False,
        export_yup=True,
        export_cameras=False,
        export_lights=False,
        export_materials="EXPORT",
        export_normals=True,
        export_tangents=True,
        export_morph=False,
        export_skins=False,
        export_draco_mesh_compression_enable=False,
        export_draco_mesh_compression_level=6,
        export_draco_position_quantization=14,
        export_draco_normal_quantization=10,
        export_draco_texcoord_quantization=12,
        export_draco_color_quantization=10,
        export_draco_generic_quantization=12,
        export_keep_originals=False,
        export_texture_dir="",
    )

    print(f"Exported to {glb_export_path}")


# -------------------------------------------------------------------
# Operators
# -------------------------------------------------------------------
# load scene operator
class ZSSD_OT_LoadScene(bpy.types.Operator):
    bl_idname = "zs_sd_loader.load_scene"
    bl_label = "Load Scene"
    bl_description = "Load Scene"

    def execute(self, context):
        load_scene()
        return {"FINISHED"}


# canvas exporter operator
class ZSSD_OT_ExportAssets(bpy.types.Operator):
    bl_idname = "zs_canvas.export_assets"
    bl_label = "Export Assets"
    bl_description = "Export Scene Assets to FBX and GLB"

    def execute(self, context):
        export_scene_to_fbx(self)
        export_scene_to_glb(self)
        return {"FINISHED"}


# parent class for panels
class ZSSDPanel:
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "ZS SD Loader"


# -------------------------------------------------------------------
#   Draw
# -------------------------------------------------------------------

# Panels


class ZSSD_PT_Main(ZSSDPanel, bpy.types.Panel):
    bl_label = "SD Loader"

    def draw(self, context):
        layout = self.layout

        scene = context.scene

        col = layout.column()

        self.is_connected = False

        col.label(text="Stable Diffusion Connection")

        # load scene button
        col.operator("zs_sd_loader.load_scene", text="Load Scene")

        # export assets button
        col.operator("zs_canvas.export_assets", text="Export Assets")


# modify after making products
blender_classes = [
    ZSSD_PT_Main,
    ZSSD_OT_LoadScene,
    ZSSD_OT_ExportAssets,
]


def register():

    # register classes
    for blender_class in blender_classes:
        bpy.utils.register_class(blender_class)

    # Has to be afqter class registering to correctly register property

    # register global properties

    # register list

    # list data

    # bpy.types.Scene.zs_product_list = bpy.props.CollectionProperty(
    #    type=ZS_Product_ListItem)

    # current item in list


def unregister():

    # unregister classes
    for blender_class in blender_classes:
        bpy.utils.unregister_class(blender_class)
    # unregister global properties

    # unregister list items
    # del bpy.types.Scene.my_list
    # del bpy.types.Scene.product_product_list_index