macrovision/scripts/blender/addons/macrovision/ops.py

import bpy

from mathutils import Vector, Euler, Color
import json
import pathlib
import os
from math import pi
import random
import bmesh

VIEW_DATA = {
    "FRONT": [0, 1, 2, "Front"],
    "ANGLED": [0.25, 1, 2, "Angled"],
    "CORNER": [0.5, 1, 2, "Corner"],
    "SIDE": [1, 1, 2, "Side"],
    "BACK_CORNER": [1.5, 1, 2, "Back Corner"],
    "BACK": [2, 1, 2, "Back"],
    "TOP": [0, 0, 1, "Top"],
    "BOTTOM": [0, 2, 1, "Bottom"],
    "BOTTOM_FLIPPED": [2, 2, 1, "Bottom Flipped"],
}

def read_sci(property):
    return property.base * pow(10, property.power)

def get_bounds(objects):
    xl = []
    yl = []
    zl = []
    
    for obj in objects:
        if not obj.hide_render:
            for bounds in obj.bound_box:
                v = obj.matrix_world @ Vector(bounds)
                xl += [v[0] for c in obj.bound_box]
                yl += [v[1] for c in obj.bound_box]
                zl += [v[2] for c in obj.bound_box]
            
    return (
        Vector([min(xl), min(yl), min(zl)]),
        Vector([min(xl), min(yl), max(zl)]),
        Vector([min(xl), max(yl), min(zl)]),
        Vector([min(xl), max(yl), max(zl)]),
        Vector([max(xl), min(yl), min(zl)]),
        Vector([max(xl), min(yl), max(zl)]),
        Vector([max(xl), max(yl), min(zl)]),
        Vector([max(xl), max(yl), max(zl)]),
    )

class MVConfigCollection(bpy.types.Operator):
    bl_idname = "mv.config_collection"
    bl_label = "Configure root collection"

    @classmethod
    def poll(cls, context: bpy.types.Context):
        return True
    
    def execute(self, context: bpy.types.Context):
        coll = context.scene.collection.children[0]
        coll.name = "Macrovision"

        mats = [
            ("light", 0, 0, 1),
            ("medium", 0, 1, 0),
            ("dark", 1, 0, 0)
        ]
        
        for name, red, green, blue in mats:
            if name not in bpy.data.materials:
                bpy.data.materials.new(name)
            mat = bpy.data.materials[name]

            mat.use_nodes = True
            mat.use_fake_user = True

            mat.node_tree.nodes.clear()

            rgb = mat.node_tree.nodes.new("ShaderNodeRGB")
            out = mat.node_tree.nodes.new("ShaderNodeOutputMaterial")

            input = out.inputs['Surface']
            output = rgb.outputs['Color']

            mat.node_tree.links.new(input=input, output=output)
            
            output.default_value = (red, green, blue, 1.0)


        if "cam" not in bpy.data.objects:
            new_cam = bpy.data.cameras.new("cam")
            new_obj = bpy.data.objects.new("cam", new_cam)
            context.scene.collection.objects.link(new_obj)
        if "cam2" not in bpy.data.objects:
            new_cam = bpy.data.cameras.new("cam2")
            new_obj = bpy.data.objects.new("cam2", new_cam)
            context.scene.collection.objects.link(new_obj)

        if "Lines" not in bpy.data.materials:
            mat = bpy.data.materials.new("Lines")
            bpy.data.materials.create_gpencil_data(mat)

        if "lineart" not in bpy.data.objects:
            new_lineart = bpy.data.grease_pencils.new("lineart")
            new_obj = bpy.data.objects.new("lineart", new_lineart)

            new_obj.show_in_front = True

            context.scene.collection.objects.link(new_obj)

            modifier = new_obj.grease_pencil_modifiers.new(name='Lineart', type='GP_LINEART')
            material = bpy.data.materials["Lines"]

            new_lineart.materials.append(material)

            modifier.target_layer = "Lines"
            modifier.target_material = material

            new_lineart.layers.new("Lines")

            # we have to clear the bake, for some reason

            bpy.ops.object.lineart_clear_all()

            new_obj.grease_pencil_modifiers[0].use_custom_camera = True
            new_obj.grease_pencil_modifiers[0].source_camera = bpy.data.objects["cam2"]

        return {'FINISHED'}

class MVAssignMaterials(bpy.types.Operator):
    bl_idname = "mv.assign_materials"
    bl_label = "Assign Materials"
    
    def execute(self, context: bpy.types.Context):
        mv = bpy.data.collections["Macrovision"]
        collections = mv.children

        for coll in collections:
            object: bpy.types.Object
            for object in coll.objects:
                if object.type != "MESH":
                    continue

                if context.scene.mv_material_mode == 'RANDOM':
                    for index in range(len(object.material_slots)):
                        choices = [
                            bpy.data.materials['light'],
                            bpy.data.materials['medium'],
                            bpy.data.materials['dark']
                        ]

                        object.material_slots[index].material = random.choice(choices)
                elif context.scene.mv_material_mode == 'NAMES':
                    for index in range(len(object.material_slots)):
                        if object.material_slots[index].material.name in ('light', 'medium', 'dark'):
                            continue

                        material = object.material_slots[index].material
                        light_choices = context.scene.mv_material_names_light.split(",")
                        medium_choices = context.scene.mv_material_names_medium.split(",")

                        chosen = None

                        for choice in light_choices:
                            if choice in material.name or choice in object.name:
                                chosen = bpy.data.materials['light']
                                break

                        for choice in medium_choices:
                            if choice in material.name or choice in object.name:
                                chosen = bpy.data.materials['medium']
                                break

                        if chosen is None:
                            chosen = bpy.data.materials['dark']

                        object.material_slots[index].material = chosen
                elif context.scene.mv_material_mode == 'FACE_MAPS':
                    while len(object.material_slots) > 0:
                        bpy.ops.object.material_slot_remove({'object': object})
                    
                    light_choices = context.scene.mv_material_names_light.split(",")
                    medium_choices = context.scene.mv_material_names_medium.split(",")

                    for map in object.face_maps:
                        
                        chosen = None

                        if map.name in light_choices:
                            chosen = bpy.data.materials['light']
                        if map.name in medium_choices:
                            chosen = bpy.data.materials['medium']

                        if chosen is None:
                            chosen = bpy.data.materials['dark']

                        bpy.ops.object.material_slot_add({'object': object})
                        index = len(object.material_slots) - 1
                        object.material_slots[index].material = chosen
                    
                    bpy.ops.object.material_slot_add({'object': object})
                    missing_idx = len(object.material_slots) - 1
                    object.material_slots[missing_idx].material = bpy.data.materials['dark']
                    bm: bpy.types.BMesh

                    bpy.ops.object.select_all(action='DESELECT')
                    object.select_set(True)
                    bpy.ops.object.mode_set(mode = 'EDIT')
                    bm = bmesh.from_edit_mesh(object.data)

                    fm = bm.faces.layers.face_map.verify()

                    for face in bm.faces:
                        idx = face[fm]
                        print(idx)
                        if idx >= 0:
                            face.material_index = idx
                        else:
                            face.material_index = missing_idx
                    
                    object.data.update()
                    bm.free()
                    print("OK")
                    bpy.ops.object.mode_set(mode = 'OBJECT')
                elif context.scene.mv_material_mode == 'OBJECT_NAMES':
                    while len(object.material_slots) > 0:
                        bpy.ops.object.material_slot_remove({'object': object})
                    
                    light_choices = context.scene.mv_material_names_light.split(",")
                    medium_choices = context.scene.mv_material_names_medium.split(",")
                    
                    bpy.ops.object.material_slot_add({'object': object})
                    missing_idx = len(object.material_slots) - 1

                    chosen = 'dark'

                    if any([choice in object.name for choice in medium_choices]):
                        chosen = 'medium'
                    if any([choice in object.name for choice in light_choices]):
                        chosen = 'light'
                    
                    object.material_slots[0].material = bpy.data.materials[chosen]


                    


        return {'FINISHED'}

class MVExport(bpy.types.Operator):
    bl_idname = "mv.export"
    bl_label = "Export objects"

    def execute(self, context: bpy.context):
        path_info = pathlib.Path(bpy.data.filepath).parent.joinpath("macrovision-directory.txt")
        config_path = pathlib.Path(open(path_info).read().strip())

        json_path = config_path.joinpath("config.json")
        config = json.load(open(json_path.resolve(), encoding="utf-8"))

        parent_workdir = config["work-directory"]

        c = bpy.data.objects["cam"]
        c2 = bpy.data.objects["cam2"]
        context.scene.camera = c

        lineart = bpy.data.objects["lineart"]
        lineart.grease_pencil_modifiers['Lineart'].source_type = 'COLLECTION'
        #lineart.data.stroke_thickness_space = 'SCREENSPACE'
        
        c.data.type = "ORTHO"
        c2.data.type = "ORTHO"

        bpy.data.scenes["Scene"].render.resolution_x = 2000
        bpy.data.scenes["Scene"].render.resolution_y = 2000
        bpy.data.scenes["Scene"].render.film_transparent = True
        bpy.data.scenes["Scene"].view_settings.view_transform = "Raw"

        bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = 0

        mv = bpy.data.collections["Macrovision"]
        collections = mv.children

        all_data = {}

        all_data["name"] = context.scene.mv_name
        all_data["kind"] = context.scene.mv_kind
        all_data["trace_alpha"] = context.scene.mv_trace_mode
        all_data["forms"] = []

        default_views = []

        for view in context.scene.mv_views:
            key = view.view
            default_views.append([
                VIEW_DATA[key][0],
                VIEW_DATA[key][1],
                VIEW_DATA[key][2],
                view.name if view.name != "" else VIEW_DATA[key][3]
            ])
        print(default_views)

        workdir = pathlib.Path(parent_workdir).joinpath(all_data["name"])

        os.makedirs(workdir, exist_ok=True)

        coll: bpy.types.Collection

        for coll in collections:
            coll.hide_render = True

        for coll in collections:
            if bpy.context.scene.view_layers['ViewLayer'].layer_collection.children[0].children[coll.name].exclude:
                continue
            coll.hide_render = False
            
            bpy.ops.object.select_all(action='DESELECT')
            
            for obj in coll.objects:
                obj.select_set(True)

            data = {}
            
            data["name"] = coll.name
            data["views"] = []

            bounds = get_bounds(coll.objects)
            bounds_seq = []
            for bound in bounds:
                bounds_seq += [bound[0], bound[1], bound[2]]
            bound_min = bounds[0]
            bound_max = bounds[-1]
            dimensions = bound_max - bound_min
            size = max(dimensions)
            global_bbox_center = 0.5 * (bound_min + bound_max)
            
            view_list = []

            lineart.grease_pencil_modifiers['Lineart'].source_collection = coll
            
            if "Views" in coll:
                for view in coll["Views"].split(","):
                    view_list.append(VIEW_DATA[view])
            else:
                view_list = default_views

            for angles in view_list:
                c.location = global_bbox_center
                c.rotation_euler = Euler([angles[1] * pi / 2, 0, angles[0] * pi / 2])
                print(list(bound_min) + list(bound_max))
                _, c.data.ortho_scale = c.camera_fit_coords(bpy.context.evaluated_depsgraph_get(), bounds_seq)
                c.location = Vector([c.location[0], c.location[1], c.location[2]])
                c.data.ortho_scale *= 1.1
                rot = c.rotation_euler.to_matrix()
                rot.invert()
                c.location += Vector([0, 0, size * 2]) @ rot
                c.data.clip_start = size / 4
                c.data.clip_end = size * 8

                c2.rotation_euler = Euler([angles[1] * pi / 2 + 0.001, 0.001, angles[0] * pi / 2 + 0.001])
                c2.location = c.location
                c2.data.ortho_scale = c.data.ortho_scale
                c2.data.clip_start = c.data.clip_start
                c2.data.clip_end = c2.data.clip_end

                scale_factor = (10 ** context.scene.mv_scale_factor)

                height = dimensions[angles[2]] * scale_factor

                view_data = {
                    "name": angles[3],
                    "height": height
                }

                volume_mode = context.scene.mv_entity_volume_mode
                mass_mode = context.scene.mv_entity_mass_mode

                volume = None

                if volume_mode == "MANUAL":
                    volume = read_sci(coll.mv_entity_volume)
                else:
                    # this over-estimates, since overlapping objects will count the same space
                    volume = 0
                    for object in coll.objects:
                        if object.type == "MESH":
                            bm = bmesh.new()
                            bm.from_object(object, context.evaluated_depsgraph_get())
                            volume += bm.calc_volume()
                            bm.free()
                
                mass = None

                if mass_mode == "MANUAL":
                    mass = read_sci(coll.mv_entity_mass)
                elif mass_mode == "DENSITY":
                    mass = volume * read_sci(coll.mv_entity_density) * 1000

                if volume_mode != "OFF":
                    view_data["volume"] = volume * (scale_factor ** 3)
                if mass is not None:
                    view_data["mass"] = mass * (scale_factor ** 3)

                data["views"].append(view_data)

                lineart.hide_render = False
                filename = f"{coll.name}-{angles[3]}.png"
                bpy.context.scene.render.filepath = workdir.joinpath(filename).resolve().__str__()
                bpy.ops.render.render(write_still = True)
                lineart.hide_render = True
                filename = f"{coll.name}-{angles[3]}-noline.png"
                bpy.context.scene.render.filepath = workdir.joinpath(filename).resolve().__str__()
                bpy.ops.render.render(write_still = True)
                
            all_data["forms"].append(data)
            coll.hide_render = True

        with open(workdir.joinpath("data.json"), "w") as file:
            json.dump(all_data, file)

        return {'FINISHED'}

clses = [
    MVExport,
    MVAssignMaterials,
    MVConfigCollection
]