tree_generator_script.py

import bpy
import bmesh
import random
from collections import namedtuple
import blenderbim.tool as tool
from math import sin, cos, pi
from random import uniform
from mathutils import Matrix, Vector

# TODO:
# - generate simple shapes
# - generate palm trees


def blender_simple_enum_items(*items):
    return [(i, i, i) for i in items]


SimpleTreeParams = namedtuple("SimpleTreeParams", "plant_height crown_diameter trunk_diameter")
LowPolyTreeParams = namedtuple(
    "LowPolyTreeParams", "plant_height crown_diameter crown_max_loc crown_taper trunk_height trunk_diameter random_seed"
)
PalmTreeParams = namedtuple("PalmTreeParams", "plant_height crown_diameter trunk_diameter")
TreePresetData = namedtuple("TreePresetData", "preset_class generator")


GeneratedGeometry = namedtuple("GeneratedGeometry", "verts_2D edges_2D faces_2D verts_3D faces_3D")


def generate_low_poly_tree(
    crown_taper=0.5,
    crown_max_loc=0.3,
    plant_height=12.0,
    trunk_height=2.0,
    crown_diameter=4,
    trunk_diameter=0.3,
    random_seed=10,
) -> GeneratedGeometry:
    random.seed(random_seed)
    res = 6
    crown_height = plant_height - trunk_height

    def circle(center, r, res):
        points = []
        for i in range(res):
            # fmt: off
            vert = (
                cos(i * 2 * pi / res) * r + center[0],
                sin(i * 2 * pi / res) * r + center[1],
                0.0 + center[2]
            )
            # fmt: on
            points.append(vert)
        return points

    def merge_lists(lists):
        merged_list = []
        for i in range(len(lists)):
            merged_list = merged_list + lists[i]
        return merged_list

    def XY_scale(vert, scale):
        vert_scaled = (vert[0] * scale, vert[1] * scale, vert[2])
        return vert_scaled

    def rand_add(vert, amp):
        # fmt: off
        vert_new = (
            vert[0] + uniform(-amp, amp),
            vert[1] + uniform(-amp, amp),
            vert[2] + uniform(-amp, amp)
        )
        # fmt: on
        return vert_new

    height_segments_trunk = 2
    height_segments_crown = 7
    height_segments_total = height_segments_trunk + height_segments_crown

    verts = []

    # Trunk
    for i in range(height_segments_trunk):
        points = circle((0.0, 0.0, i * trunk_height), trunk_diameter / 2, res)
        verts.append(points)

    # Crown
    h_crown_max_loc = int((height_segments_crown - 1) * crown_max_loc)
    crown_indices = []
    for i in range(0, h_crown_max_loc):
        crown_indices.append(i)
    for i in range(0, height_segments_crown - h_crown_max_loc):
        crown_indices.append(h_crown_max_loc - i)

    my_min_val = min(crown_indices)
    my_max_val = max(crown_indices)

    n_crown_taper = crown_taper * 0.8

    crown_diameters = []
    for x in crown_indices:
        crown_diameters.append(((x - my_min_val) / (my_max_val - my_min_val)) * n_crown_taper + (1 - n_crown_taper))

    randomize_amplitude = plant_height / 50

    crown_segment_height = crown_height / height_segments_crown
    for i in range(0, height_segments_crown - 1):
        points = circle(
            (0.0, 0.0, trunk_height + (i + 1) * crown_segment_height),
            crown_diameter * crown_diameters[i] / 2,
            res,
        )
        for i in range(0, len(points)):
            points[i] = rand_add(points[i], randomize_amplitude)
        verts.append(points)

    # Top
    verts.append(circle((0.0, 0.0, plant_height), 0.05, res))

    edge_loops = []
    for i in range(height_segments_total):
        edge_loops.append(range(i * res, (i + 1) * res))

    verts = merge_lists(verts)

    faces = []
    bottom = [*edge_loops[0]]
    faces.append(bottom)
    for i in range(height_segments_total - 1):
        for j in range(res):
            face = (edge_loops[i][j - 1], edge_loops[i][j], edge_loops[i + 1][j], edge_loops[i + 1][j - 1])
            faces.append(face)
    top = [*edge_loops[-1]]
    faces.append(top)

    verts_3D = [verts]
    faces_3D = [faces]

    verts_2D = [
        [
            (-0.5865642428398132, 0.8647078275680542, 0.0),
            (-0.7583824992179871, 0.6572927236557007, 0.0),
            (-0.7234422564506531, 0.5865941047668457, 0.0),
            (-0.8714070916175842, 0.46631893515586853, 0.0),
            (-0.9203394055366516, 0.15792329609394073, 0.0),
            (-0.6487269997596741, 0.11522211134433746, 0.0),
            (-0.9273074269294739, 0.0903107076883316, 0.0),
            (-0.9606701731681824, -0.2852219343185425, 0.0),
            (-0.8060722351074219, -0.35055163502693176, 0.0),
            (-0.8271104693412781, -0.4677186608314514, 0.0),
            (-0.6899875998497009, -0.6695915460586548, 0.0),
            (-0.6215555667877197, -0.6767659187316895, 0.0),
            (-0.4315463900566101, -0.8824808597564697, 0.0),
            (-0.30316293239593506, -0.9023936986923218, 0.0),
            (-0.32218849658966064, -0.8037619590759277, 0.0),
            (-0.21792533993721008, -0.9530860185623169, 0.0),
            (0.36917445063591003, -0.877334475517273, 0.0),
            (0.30831706523895264, -0.685136079788208, 0.0),
            (0.4778200685977936, -0.8094909191131592, 0.0),
            (0.9163193106651306, -0.368840754032135, 0.0),
            (0.827597439289093, -0.29377281665802, 0.0),
            (0.9478662610054016, -0.245104119181633, 0.0),
            (0.9667968153953552, 0.08458180725574493, 0.0),
            (0.7529001832008362, 0.05212751030921936, 0.0),
            (0.9102436900138855, 0.19098728895187378, 0.0),
            (0.8279229998588562, 0.5522171258926392, 0.0),
            (0.6635335087776184, 0.6461355686187744, 0.0),
            (0.5753684639930725, 0.8170149326324463, 0.0),
            (0.3976244330406189, 0.8779193162918091, 0.0),
            (0.18733686208724976, 0.7513588666915894, 0.0),
            (0.338344544172287, 0.910933256149292, 0.0),
            (-0.013130240142345428, 0.9798095226287842, 0.0),
            (-0.18649885058403015, 0.9414206743240356, 0.0),
            (-0.3193224370479584, 0.8835403919219971, 0.0),
        ]
    ]
    edges_2D = []
    faces_2D = []

    verts_2D_new = []
    for vert in verts_2D[0]:
        verts_2D_new.append(XY_scale(vert, crown_diameter / 2))

    range_of_verts = range(0, len(verts_2D[0]))
    for i in range_of_verts:
        edges_2D.append([range_of_verts[i - 1], i])
        faces_2D.append(i)

    verts_2D[0] = verts_2D_new
    edges_2D = [edges_2D]
    faces_2D = [[faces_2D]]
    return GeneratedGeometry(verts_2D, edges_2D, faces_2D, verts_3D, faces_3D)


def bm_extrude(bm, geom, extrusion_vector):
    extruded = bmesh.ops.extrude_face_region(bm, geom=geom)
    translate_verts = [v for v in extruded["geom"] if isinstance(v, bmesh.types.BMVert)]
    bmesh.ops.translate(bm, vec=extrusion_vector, verts=translate_verts)


def generate_simple_tree(plant_height=2, crown_diameter=2, trunk_diameter=0.1) -> GeneratedGeometry:
    # NOTE: with ifc it will be just IFC sphere
    bm = bmesh.new()
    crown_radius = crown_diameter / 2

    # trunk
    circle_geom = bmesh.ops.create_circle(bm, cap_ends=True, segments=6, radius=trunk_diameter / 2)
    extrusion_vector = Vector((0, 0, 1)) * (plant_height - crown_diameter)
    bm_extrude(bm, [circle_geom["verts"][0].link_faces[0]], extrusion_vector)

    # crown
    crown_offset_matrix = Matrix.Translation((0, 0, plant_height - crown_radius))
    bmesh.ops.create_uvsphere(bm, u_segments=12, v_segments=12, radius=crown_radius, matrix=crown_offset_matrix)

    verts_3D = [v.co.copy() for v in bm.verts]
    faces_3D = [[v.index for v in face.verts] for face in bm.faces]

    # 2d shape
    shape_2D = bmesh.ops.create_circle(bm, cap_ends=True, segments=12, radius=crown_radius)["verts"]
    verts_2D = [v.co.copy() for v in shape_2D]
    faces_2D = [list(range(len(verts_2D)))]

    bm.free()

    return GeneratedGeometry([verts_2D], None, [faces_2D], [verts_3D], [faces_3D])


def generate_palm_tree(plant_height=2, crown_diameter=5.8, trunk_diameter=0.1) -> GeneratedGeometry:
    palm_verts = [
        Vector((-3.6415634155273438, -3.8275668621063232, -1.0948246717453003)),
        Vector((-0.3698049783706665, -1.805966854095459, 0.493910551071167)),
        Vector((-1.6397790908813477, -0.5056067705154419, 0.04353363811969757)),
        Vector((0.0, 0.0, 0.0)),
        Vector((5.023162364959717, -1.6531240940093994, -1.0948246717453003)),
        Vector((1.8090602159500122, 0.4589407444000244, 0.493910551071167)),
        Vector((1.1666079759597778, -1.2413609027862549, 0.04353363811969757)),
        Vector((1.6942416429519653, 5.038582801818848, -1.0948246717453003)),
        Vector((-0.43215513229370117, 1.833944320678711, 0.493910551071167)),
        Vector((1.2652605772018433, 1.1839056015014648, 0.04353363811969757)),
        Vector((-4.577561378479004, 2.689336061477661, -1.0948246717453003)),
        Vector((-1.8646581172943115, -0.0367276668548584, 0.493910551071167)),
        Vector((-0.8873085975646973, 1.4957730770111084, 0.04353363811969757)),
    ]
    palm_faces = [(2, 1, 0), (2, 3, 1), (6, 5, 4), (6, 3, 5), (9, 8, 7), (9, 3, 8), (12, 11, 10), (12, 3, 11)]

    scale = Vector()
    scale.xyz = crown_diameter / 11.6
    leaves_height = 0.493910551071167 * scale.z
    plant_height -= leaves_height
    for v in palm_verts:
        v *= scale
        v.z += plant_height

    # 3d
    bm = bmesh.new()
    palm_bm_verts = [bm.verts.new(v) for v in palm_verts]
    [bm.faces.new([palm_bm_verts[i] for i in face]) for face in palm_faces]

    circle_geom = bmesh.ops.create_circle(bm, cap_ends=True, segments=6, radius=trunk_diameter / 2)
    extrusion_vector = Vector((0, 0, 1)) * plant_height
    bm_extrude(bm, [circle_geom["verts"][0].link_faces[0]], extrusion_vector)
    verts_3D = [v.co.copy() for v in bm.verts]
    faces_3D = [[v.index for v in face.verts] for face in bm.faces]

    bm.free()

    # 2d
    remove_z = Vector((1, 1, 0))
    verts_2D = [v * remove_z for v in palm_verts]
    faces_2D = palm_faces.copy()

    return GeneratedGeometry([verts_2D], None, [faces_2D], [verts_3D], [faces_3D])


tree_presets = {
    "simple": TreePresetData(SimpleTreeParams, generate_simple_tree),
    "low_poly": TreePresetData(LowPolyTreeParams, generate_low_poly_tree),
    "palm_tree": TreePresetData(PalmTreeParams, generate_palm_tree),
}


def update_tree_geometry_bmesh():
    obj = bpy.context.active_object
    props = obj.tree_generator_props

    get_params = lambda tree_class: {prop_name: getattr(props, prop_name) for prop_name in tree_class._fields}
    preset = tree_presets.get(props.preset)
    prop_data = get_params(preset.preset_class)
    tree_geometry = preset.generator(**prop_data)

    if props.mode == "3d":
        verts = tree_geometry.verts_3D[0]
        faces = tree_geometry.faces_3D[0]
    else:
        verts = tree_geometry.verts_2D[0]
        faces = tree_geometry.faces_2D[0]

    bm = tool.Blender.get_bmesh_for_mesh(obj.data, clean=True)
    bm_verts = [bm.verts.new(v) for v in verts]
    for face in faces:
        bm.faces.new((bm_verts[i] for i in face))
    tool.Blender.apply_bmesh(obj.data, bm, obj)


class SaveTreeGeneratorParams(bpy.types.Operator):
    bl_idname = "object.save_tree_generator_params"
    bl_label = "Save Tree Generator Params To Clipboard"

    def execute(self, context):
        obj = context.object
        props = obj.tree_generator_props
        preset = tree_presets[props.preset]
        props_data = "\t".join(str(round(getattr(props, prop_name), 3)) for prop_name in preset.preset_class._fields)
        props_data = f"{props.preset}\t{props_data}"
        bpy.context.window_manager.clipboard = props_data
        self.report({"INFO"}, f"Tree generator props saved to clipboard: {props_data}")
        return {"FINISHED"}


class LoadTreeGeneratorParams(bpy.types.Operator):
    bl_idname = "object.load_tree_generator_params"
    bl_label = "Load Tree Generator Params From Clipboard"

    def execute(self, context):
        obj = context.object
        props = obj.tree_generator_props
        props_data = bpy.context.window_manager.clipboard.split("\t")
        preset_name = props_data[0]
        props.preset = preset_name

        preset = tree_presets[preset_name]
        props_data = preset.preset_class(*props_data[1:])
        for prop_name in preset.preset_class._fields:
            value = getattr(props_data, prop_name)
            value = type(getattr(props, prop_name))(value)
            setattr(props, prop_name, value)
        self.report({"INFO"}, f"Tree generator props loaded from clipboard: {props_data}")
        return {"FINISHED"}


class TreeGeneratorPanel(bpy.types.Panel):
    bl_label = "Tree Generator Panel"
    bl_idname = "OBJECT_PT_tree_generator"
    bl_space_type = "PROPERTIES"
    bl_region_type = "WINDOW"
    bl_context = "modifier"

    def draw(self, context):
        layout = self.layout
        obj = context.active_object
        props = obj.tree_generator_props

        layout.prop(props, "is_editing")
        if props.is_editing:
            layout.operator("object.save_tree_generator_params")
            layout.operator("object.load_tree_generator_params")
            layout.prop(props, "preset")
            layout.prop(props, "mode")
            for prop_name in tree_presets.get(props.preset).preset_class._fields:
                layout.prop(props, prop_name)
            update_tree_geometry_bmesh()


class TreeGeneratorProps(bpy.types.PropertyGroup):
    is_editing: bpy.props.BoolProperty(name="Enable Tree Generator", default=False)
    preset: bpy.props.EnumProperty(
        name="Preset", items=blender_simple_enum_items(*tree_presets.keys()), default="low_poly"
    )
    mode: bpy.props.EnumProperty(name="Display Mode", items=blender_simple_enum_items("3d", "2d"))
    random_seed: bpy.props.IntProperty(default=10)

    # low poly
    plant_height: bpy.props.FloatProperty(default=12.0, min=0, subtype="DISTANCE")
    crown_diameter: bpy.props.FloatProperty(default=4.0, min=0, subtype="DISTANCE")
    crown_max_loc: bpy.props.FloatProperty(default=0.3, min=0, max=1)
    crown_taper: bpy.props.FloatProperty(default=0.5, min=0, max=1)
    trunk_height: bpy.props.FloatProperty(default=2.0, min=0, subtype="DISTANCE")
    trunk_diameter: bpy.props.FloatProperty(default=0.3, min=0, subtype="DISTANCE")

    # simple
    plant_height: bpy.props.FloatProperty(default=2, subtype="DISTANCE")


bpy.utils.register_class(TreeGeneratorPanel)
bpy.utils.register_class(TreeGeneratorProps)
bpy.utils.register_class(LoadTreeGeneratorParams)
bpy.utils.register_class(SaveTreeGeneratorParams)
bpy.types.Object.tree_generator_props = bpy.props.PointerProperty(type=TreeGeneratorProps)


print("Trees setup in current .blend project:")
trees_data = []
for obj in bpy.data.objects:
    props = obj.tree_generator_props
    preset_name = props.preset
    preset = tree_presets.get(preset_name).preset_class
    tree_data = [obj.name, preset_name]
    preset_data = []
    for prop_name in preset._fields:
        value = round(getattr(props, prop_name), 3)
        preset_data.append(value)
    tree_data.append(tuple(preset_data))
    print(f"{tuple(tree_data)},")