blender_utils.py

# Adapted from https://github.com/5agado/data-science-learning/blob/master/ds_utils/blender_utils.py,
# which is generously provided by 5agado under the Apache license.

import bpy
import bpy_extras


import numpy as np
from math import sin, cos

assert bpy.app.version[1] >= 80

def delete_all_materials():
    for mat in bpy.data.materials:
        bpy.data.materials.remove(mat)

def delete_all(obj_type: str):
    """Delete all objects of the given type from the current scene"""

    for obj in bpy.data.objects:
        obj.hide_select = False
        obj.select_set(obj.type == obj_type)

    bpy.ops.object.delete(use_global=True)


def create_line(p0: tuple, p1: tuple, name: str):
    mesh = bpy.data.meshes.new(name)

    mesh.from_pydata([p0, p1], [(0, 1)], [])
    mesh.update()

    return add_mesh(mesh, name)


def add_mesh(mesh, obj_name: str, collection="Collection"):
    obj = bpy.data.objects.new(obj_name, mesh)
    bpy.data.collections[collection].objects.link(obj)
    return obj


def add_text(text: str, location: tuple = (0, 0, 0)):
    bpy.ops.object.text_add(location=location)
    text_obj = bpy.context.view_layer.objects.active
    text_obj.data.body = text
    return text_obj


def render(filepath: str, animation=False, file_format=None):
    bpy.context.scene.render.filepath = filepath
    if animation:
        bpy.context.scene.render.image_settings.file_format = file_format if file_format is not None else 'FFMPEG'
        if bpy.context.scene.render.image_settings.file_format == 'FFMPEG':
            bpy.context.scene.render.ffmpeg.format = 'MPEG4'
        bpy.ops.render.render(animation=True)
    else:
        bpy.context.scene.render.image_settings.file_format = file_format if file_format is not None else 'PNG'
        bpy.ops.render.render(write_still=True)


def blender_run(nb_frames, frames_spacing=1, frame_update_print_mod=1):
    print("----------------------------")

    bpy.context.scene.frame_start = 0
    bpy.context.scene.frame_end = nb_frames * frames_spacing

    delete_all()
    for frame in range(0, nb_frames):
        if frame % frame_update_print_mod == 0:
            print("Frame {}".format(frame))


def create_grease_pencil(name: str):
    if name not in bpy.context.scene.objects:
        bpy.ops.object.gpencil_add(location=(0, 0, 0), type='EMPTY')
        new_gp = bpy.context.scene.objects[-1]
        new_gp.name = name

        return new_gp
    else:
        Exception("Grease Pencil named {} already exists.".format(name))


def get_grease_pencil(name: str):
    if name in bpy.context.scene.objects:
        return bpy.context.scene.objects[name]


def create_grease_pencil_layer(gp: bpy.types.GreasePencil, name: str):
    return gp.data.layers.new(name, set_active=True)


def get_grease_pencil_layer(gp: bpy.types.GreasePencil, name: str):
    if gp.data.layers and name in gp.data.layers:
        return gp.data.layers[name]


def init_grease_pencil(name: str, layer_name: str):
    gp = create_grease_pencil(name)
    layer = create_grease_pencil_layer(gp, layer_name)
    return gp, layer


def create_frame(gp_layer, frame_num):
    return gp_layer.frames.new(frame_num)


def get_frame(gp_layer, frame_num):
    for f in gp_layer.frames:
        if f.frame_num == frame_num:
            return f


def create_stroke(gp_frame, material_index, line_width=10):
    # Init new stroke
    gp_stroke = gp_frame.strokes.new()
    gp_stroke.display_mode = '3DSPACE'  # allows for editing
    gp_stroke.material_index = material_index
    gp_stroke.line_width = line_width

    return gp_stroke


def draw_line_segments(gp_stroke, points, cyclic=False, pressures=[], strengths=[]):
    # Define stroke geometry
    gp_stroke.draw_cyclic = cyclic  # closes the stroke

    gp_stroke.points.add(count=len(points))
    for i, p in enumerate(points):
        gp_stroke.points[i].co = p

    for p in pressures:
        i, pr = p
        gp_stroke.points[i].pressure = pr

    for s in strengths:
        i, st = s
        gp_stroke.points[i].strength = st

    return gp_stroke


def x_rotor(theta):
    return np.array([[1, 0, 0],
                     [0, cos(theta), -sin(theta)],
                     [0, sin(theta), cos(theta)]])


def y_rotor(theta):
    return np.array([[cos(theta), 0, sin(theta)],
                     [0, 1, 0],
                     [-sin(theta), 0, cos(theta)]])


def z_rotor(theta):
    return np.array([[cos(theta), -sin(theta, 0)],
                     [sin(theta), cos(theta), 0],
                     [0, 0, 1]])


def rotate_stroke(stroke, rotor):
    for i, p in enumerate(stroke.points):
        p.co = rotor @ np.array(p.co).reshape(3, 1)


def translate_stroke(stroke, vector):
    for i, p in enumerate(stroke.points):
        p.co = np.array(p.co) + vector

def create_stroke_material(
# Create material for grease pencil
if "Bright Material" in bpy.data.materials.keys():
    gp_mat = bpy.data.materials["Bright Material"]
else:
    gp_mat = bpy.data.materials.new("Bright Material")

if not gp_mat.is_grease_pencil:
    bpy.data.materials.create_gpencil_data(gp_mat)
    gp_mat.grease_pencil.color = (1, 0, 0.818649, 1)

delete_all("MESH")
delete_all("GPENCIL")

gp, gpl = init_grease_pencil("test", "0")
gp_frame = gpl.frames.new(0)
gp_stroke = create_stroke(gp_frame, 2, 10)

draw_line_segments(gp_stroke, [[0, 0, 0], [1, 1, 1]])

print(0)