Quick low-poly tree addon for Blender

quick_tree.py

import bpy
import bmesh
from random import random, uniform, choice, randrange
from mathutils import Matrix, Vector
from bmesh.types import BMVert, BMFace, BMEdge
from bpy.props import (
    IntProperty,
    IntVectorProperty,
    FloatProperty,
    FloatVectorProperty,
)

bl_info = {
    "name": "Quick Tree for Fugue",
    "author": "Francis Tseng",
    "version": (1, 0),
    "blender": (2, 92, 0),
    "category": "Add Mesh",
    "location": "View3D > Add > Quick Tree",
    "description": "Adds a tree generator to the Add Mesh menu",
}

def verts_for_faces(faces):
    verts = set()
    for face in faces:
        for vert in face.verts:
            verts.add(vert)
    return list(verts)


def scale_rotate_faces(bm, faces, scale=None, rotation=None):
    bm.normal_update()

    # Set proper transformation orientation
    c = sum((f.calc_center_median() for f in faces), Vector())/len(faces)
    T = Matrix.Translation(-c)

    if scale is not None:
        bmesh.ops.scale(
            bm,
            space=T,
            vec=(scale, scale, scale),
            verts=verts_for_faces(faces)
        )

    if rotation is not None:
        rot_angle, rot_axis = rotation
        bmesh.ops.rotate(
            bm,
            space=T,
            matrix=Matrix.Rotation(rot_angle, 4, rot_axis),
            verts=verts_for_faces(faces)
        )


def extrude_faces(bm, faces, length, scale=None, rotation=None):
    prev_faces = [f for f in bm.faces]

    # Extrude
    norm = sum((f.normal for f in faces), Vector())/len(faces)
    vec = length * norm.normalized()
    extruded = bmesh.ops.extrude_face_region(bm, geom=faces)
    translate_verts = [v for v in extruded['geom'] if isinstance(v, BMVert)]
    bmesh.ops.translate(bm, vec=vec, verts=translate_verts)

    # Keep track of the new faces
    lead_faces = [v for v in extruded['geom'] if isinstance(v, BMFace)]
    side_faces = [f for f in bm.faces if f not in prev_faces and f not in lead_faces]

    # Clean up
    bmesh.ops.delete(bm, geom=faces, context='FACES')
    bmesh.ops.remove_doubles(bm, verts=bm.verts, dist=0.001)

    # Apply transformations to the lead faces
    scale_rotate_faces(bm, lead_faces, scale, rotation)

    bm.faces.ensure_lookup_table()
    bm.normal_update()
    return lead_faces, side_faces


def pinch_faces(bm, faces):
    if len(faces) > 1:
        face = bmesh.utils.face_join(faces)
    else:
        face = faces[0]
    scale_rotate_faces(bm, [face], 0.)
    bmesh.ops.remove_doubles(bm, verts=face.verts, dist=0.001)


class Tree:
    def __init__(self, params):
        self.params = params

    def grow(self):
        self.bm = bmesh.new()
        lead_faces = self.gen_base()
        for i, (lead_faces, side_faces, length) in enumerate(self.gen_trunk(lead_faces)):
            if i > self.params['MIN_BRANCH_HEIGHT'] and random() < self.params['BRANCH_PROB']:
                self.gen_branch(side_faces, length)

        # "Pinch" end of trunk
        lead_faces, _ = extrude_faces(self.bm, lead_faces, self.params['TRUNK_TIP_LENGTH'])
        pinch_faces(self.bm, lead_faces)
        return self.bm

    def get_faces(self, idxs):
        return [self.bm.faces[i] for i in idxs]

    def gen_base(self):
        bmesh.ops.create_cube(self.bm, size=1.0)

        # Prep the base
        self.bm.edges.ensure_lookup_table()
        self.bm.verts.ensure_lookup_table()

        # Make a hexagon
        edges = [self.bm.edges[i] for i in [2,0,7,8]]
        bmesh.ops.subdivide_edges(self.bm, edges=edges, cuts=1)
        bmesh.ops.scale(
            self.bm,
            vec=(1.2, 1, 1),
            verts=self.bm.verts[-4:]
        )

        self.bm.faces.ensure_lookup_table()
        bmesh.ops.scale(
            self.bm,
            vec=(0.626, 1, 1),
            verts=verts_for_faces(self.get_faces([1, 3]))
        )

        # Bottom base
        bottom_scale = uniform(*self.params['TRUNK_BASE_BOTTOM_SCALE_RANGE'])
        scale_rotate_faces(self.bm, self.get_faces([7, 4]), bottom_scale)

        # Top base
        lead_faces = self.get_faces([5, 6])
        top_scale = uniform(*self.params['TRUNK_BASE_TOP_SCALE_RANGE'])
        tilt = uniform(*self.params['TRUNK_TILT_RANGE'])
        axis = (random(), random(), random())
        scale_rotate_faces(self.bm, lead_faces, top_scale, (tilt, axis))
        return lead_faces

    def gen_trunk(self, lead_faces):
        segments = randrange(*self.params['TRUNK_SEGMENTS_RANGE'])
        for _ in range(segments):
            length = uniform(*self.params['TRUNK_SEGMENT_LENGTH_RANGE'])
            tilt = uniform(*self.params['TRUNK_TILT_RANGE'])
            scale = uniform(*self.params['TRUNK_GIRTH_SCALE_RANGE'])

            axis = (random(), random(), random())
            lead_faces, side_faces = extrude_faces(self.bm, lead_faces, length, scale, (tilt, axis))
            yield lead_faces, side_faces, length

    def gen_branch(self, side_faces, length, branch_depth=0):
        # Subdivide the face
        branch_face = choice(side_faces)
        res = bmesh.ops.subdivide_edges(self.bm,
            edges=[
                branch_face.edges[1],
                branch_face.edges[3]], cuts=1)

        # New face created from subdividing
        new_face = [v for v in res['geom'] if isinstance(v, BMFace) and len(v.verts) == 4][0]

        # Adjust the edge for the branch root face
        edge = [v for v in res['geom_inner'] if isinstance(v, BMEdge)][0]
        tangent = new_face.calc_tangent_edge_pair()
        target_root_length = length * self.params['BRANCH_ROOT_SIZE_PERCENT']
        shift = target_root_length - (0.5 * length) # after subdiving, the size is half the length of the trunk segment
        translate = tangent.normalized() * shift
        bmesh.ops.translate(self.bm, vec=translate, verts=edge.verts)

        # Branches generally get smaller the deeper the branchings occurs are
        branch_depth_multiplier = self.params['BRANCH_DEPTH_MULTIPLIER']**branch_depth

        # Branch start params
        length = uniform(*self.params['BRANCH_START_LENGTH_RANGE']) * branch_depth_multiplier
        tilt = uniform(*self.params['BRANCH_START_TILT_RANGE'])
        scale = uniform(*self.params['BRANCH_START_GIRTH_RANGE']) * branch_depth_multiplier

        # Figure out normal X-axis for rotation
        tangent = new_face.calc_tangent_edge_pair()
        axis = -new_face.normal.cross(tangent)

        lead_faces, side_faces = extrude_faces(self.bm, [new_face], length, scale, (tilt, axis))

        # Extend branch
        n_segments = randrange(*self.params['BRANCH_SEGMENTS_RANGE'])
        for _ in range(n_segments):
            length = uniform(*self.params['BRANCH_SEGMENT_LENGTH_RANGE']) * branch_depth_multiplier
            scale = uniform(*self.params['BRANCH_SEGMENT_GIRTH_RANGE'])
            tilt = uniform(*self.params['BRANCH_SEGMENT_TILT_RANGE'])
            axis = (random(), random(), random())
            lead_faces, side_faces = extrude_faces(self.bm, lead_faces,
                    length=length,
                    scale=scale, rotation=(tilt, axis))
            if branch_depth < self.params['MAX_BRANCH_DEPTH'] and random() < self.params['BRANCH_PROB']:
                self.gen_branch(side_faces, length, branch_depth=branch_depth+1)

        # Pinch the end of branches
        length = uniform(*self.params['BRANCH_SEGMENT_LENGTH_RANGE']) * branch_depth_multiplier
        lead_faces, _ = extrude_faces(self.bm, lead_faces, length=length)
        pinch_faces(self.bm, lead_faces)
        return face_idx


# ---


class QuickTreeOperator(bpy.types.Operator):
    """Export as GLTF to the corrresponding game models folder"""
    bl_idname = "mesh.quick_tree"
    bl_label = "Quick Tree"
    bl_description = "Generate a tree"
    bl_options = {'REGISTER', 'UNDO'}

    TRUNK_BASE_BOTTOM_SCALE_RANGE: FloatVectorProperty(
        name="Scaling of bottom of the trunk base",
        description="Scaling of bottom of the trunk base",
        min=0.1, max=5, size=2,
        default=[1, 1.5])

    TRUNK_BASE_TOP_SCALE_RANGE: FloatVectorProperty(
        name="Scaling of top of the trunk base",
        description="Scaling of top of the trunk base",
        min=0.1, max=5, size=2,
        default=[0.8, 1.1])

    TRUNK_TILT_RANGE: FloatVectorProperty(
        name="Trunk segment tilting angle",
        description="Trunk segment tilting angle",
        min=-3.2, max=3.2, size=2,
        default=[-0.1, 0.1])

    TRUNK_SEGMENTS_RANGE: IntVectorProperty(
        name="Number of trunk segments",
        description="Number of trunk segments",
        min=1, max=200, size=2,
        default=[10, 15])

    TRUNK_SEGMENT_LENGTH_RANGE: FloatVectorProperty(
        name="Trunk segment length",
        description="Trunk segment length",
        min=1, max=20, size=2,
        default=[1, 1.5])

    TRUNK_GIRTH_SCALE_RANGE: FloatVectorProperty(
        name="Scaling of trunk girth between segments",
        description="Scaling of trunk girth between segments",
        min=1, max=2, size=2,
        default=[0.8, 1])

    TRUNK_TIP_LENGTH: FloatProperty(
        name="Trunk tip length",
        description="Length of the trunk tip",
        min=1, max=10,
        default=1.25)

    BRANCH_PROB: FloatProperty(
        name="Branch probability",
        description="Probability of creating a branch from a trunk segment",
        min=0, max=1.,
        default=0.5)

    MIN_BRANCH_HEIGHT: IntProperty(
        name="Min branch height",
        description="Minimum number of trunk segments before branches can spawn",
        min=1, max=200,
        default=2)

    MAX_BRANCH_DEPTH: IntProperty(
        name="Max branch depth",
        description="Maximum number of times to branch recursively",
        min=1, max=10,
        default=1)

    BRANCH_DEPTH_MULTIPLIER: FloatProperty(
        name="Branch depth multiplier",
        description="Compounding multiplier across branch depths, such that deeper branches are shorter/thinner",
        min=0.1, max=1.,
        default=0.5)

    BRANCH_ROOT_SIZE_PERCENT: FloatProperty(
        name="Branch root size percent",
        description="Branch root size (face that starts a branch) as a percentage of its parent segment",
        min=0, max=1.,
        default=0.25)

    BRANCH_START_LENGTH_RANGE: FloatVectorProperty(
        name="Branch root length",
        description="Length of the branch root (the first segment of a branch)",
        min=0.1, max=20, size=2,
        default=[0.6, 0.8])

    BRANCH_START_TILT_RANGE: FloatVectorProperty(
        name="Branch root tilt",
        description="Tilt of the branch root (the first segment of a branch)",
        min=-3.2, max=3.2, size=2,
        default=[0.7, 0.8])

    BRANCH_START_GIRTH_RANGE: FloatVectorProperty(
        name="Branch root girth",
        description="How much to scale the end of the branch root (the first segment of a branch)",
        min=0.1, max=5, size=2,
        default=[0.5, 0.6])

    BRANCH_SEGMENTS_RANGE: IntVectorProperty(
        name="Number of branch segments",
        description="Number of branch segments",
        min=1, max=20, size=2,
        default=[3, 5])

    BRANCH_SEGMENT_LENGTH_RANGE: FloatVectorProperty(
        name="Branch segment length",
        description="Lengths of branch segments (not including the root)",
        min=0.1, max=5, size=2,
        default=[0.8, 1.2])

    BRANCH_SEGMENT_GIRTH_RANGE: FloatVectorProperty(
        name="Branch segment girth",
        description="How much to scale the end of branch segments (not including the root)",
        min=0.1, max=5, size=2,
        default=[0.9, 1])

    BRANCH_SEGMENT_TILT_RANGE: FloatVectorProperty(
        name="Branch segment tilt",
        description="How much to tilt the end of branch segments (not including the root)",
        min=-3.2, max=3.2, size=2,
        default=[-0.75, 0.75])

    props = [
        'TRUNK_BASE_BOTTOM_SCALE_RANGE',
        'TRUNK_BASE_TOP_SCALE_RANGE',
        'TRUNK_TILT_RANGE',
        'TRUNK_SEGMENTS_RANGE',
        'TRUNK_SEGMENT_LENGTH_RANGE',
        'TRUNK_GIRTH_SCALE_RANGE',
        'TRUNK_TIP_LENGTH',
        'BRANCH_PROB',
        'MIN_BRANCH_HEIGHT',
        'MAX_BRANCH_DEPTH',
        'BRANCH_DEPTH_MULTIPLIER',
        'BRANCH_ROOT_SIZE_PERCENT',
        'BRANCH_START_LENGTH_RANGE',
        'BRANCH_START_TILT_RANGE',
        'BRANCH_START_GIRTH_RANGE',
        'BRANCH_SEGMENTS_RANGE',
        'BRANCH_SEGMENT_LENGTH_RANGE',
        'BRANCH_SEGMENT_GIRTH_RANGE',
        'BRANCH_SEGMENT_TILT_RANGE'
    ]

    def draw(self, context):
        box = self.layout.box()
        for prop in self.props:
            if getattr(self, prop).__class__.__name__ == 'bpy_prop_array':
                row = box.row()
                row.prop(self, prop)
            else:
                box.prop(self, prop)

    def execute(self, context):
        # Create mesh and object
        mesh = bpy.data.meshes.new('Tree')
        obj = bpy.data.objects.new("Tree", mesh)

        # Add the object into the scene.
        bpy.context.collection.objects.link(obj)
        bpy.context.view_layer.objects.active = obj

        params = {p: getattr(self, p) for p in self.props}
        tree = Tree(params)
        bm = tree.grow()
        bm.to_mesh(mesh)
        bm.free()
        return {'FINISHED'}

def quick_tree_menu(self, context):
    layout = self.layout
    layout.separator()
    layout.operator(
        QuickTreeOperator.bl_idname,
        text="Quick Tree",
        icon="MESH_ICOSPHERE")

def register():
    bpy.utils.register_class(QuickTreeOperator)
    bpy.types.VIEW3D_MT_mesh_add.append(quick_tree_menu)

def unregister():
    bpy.utils.unregister_class(QuickTreeOperator)
    bpy.types.VIEW3D_MT_mesh_add.remove(quick_tree_menu)


if __name__ == "__main__":
    register()