Blender Add-on: Export animation renders using bmild/NeRF's Blender data format

nerf_blender_data_exporter.py

bl_info = {
    "name": "NeRF Data Exporter",
    "description": "Outputs images and camera parameters in bmild/NeRF's Blender data format.",
    "author": "ecmjohnson",
    "version": (0, 3),
    "blender": (2, 80, 0),
    "support": "TESTING",
    "category": "Import-Export",
}

import bpy
from bpy.app.handlers import persistent
import json
import os
from os.path import join as pjoin
import numpy as np
import bmesh

### Constants

# Expected data split in NeRF code
train = 'train'
test = 'test'
val = 'val'
# Expected filenames in NeRF code
train_json = 'transforms_{}.json'.format(train)
test_json = 'transforms_{}.json'.format(test)
val_json = 'transforms_{}.json'.format(val)
# Filename for novel view camera
novel_json = 'novel_view.json'

### User preferences

class NerfDataExporterPreferences(bpy.types.AddonPreferences):
    bl_idname = __name__

    # Defaults to directory of Blender .blend file
    output_path: bpy.props.StringProperty(
        name = 'Data output folder',
        subtype = 'FILE_PATH'
    )

    # train -> test -> val
    train_end: bpy.props.IntProperty(
        name = 'Index of last training frame',
        default = 50,
        min = 1
    )
    # TODO some way to enforce test_end > train_end ?
    test_end: bpy.props.IntProperty(
        name = 'Index of last testing frame (must be greater than last training frame)',
        default = 75,
        min = 1
    )

    novel_cam_name: bpy.props.StringProperty(
        name = 'Novel view camera'
    )

    def draw(self, context):
        layout = self.layout
        layout.label(text = 'Preferences for NeRF data exporter')
        layout.prop(self, 'output_path')
        layout.prop(self, 'train_end')
        layout.prop(self, 'test_end')
        layout.prop(self, 'novel_cam_name')

### Helper functions

def get_prefs():
    return bpy.context.preferences.addons[__name__].preferences

def is_train(f):
    prefs = get_prefs()
    return f <= prefs.train_end

def is_test(f):
    prefs = get_prefs()
    return (f > prefs.train_end) and (f <= prefs.test_end)

def is_val(f):
    prefs = get_prefs()
    return f > prefs.test_end

def get_name(f):
    if is_train(f):
        return train
    elif is_test(f):
        return test
    elif is_val(f):
        return val
    else:
        print('Failure in get_name, f =', f)
        return 'fail'

def sanitize_path(p):
    return p.replace('\\', '/')

### Callback functions

root_folder = ''
f = 1
frames = []
out = {}
novel = []

# On start of rendering
@persistent
def render_init_fn(scene):
    # Clear previous state
    frames.clear()
    out.clear()
    novel.clear()
    # Setup output folders
    global root_folder
    prefs = get_prefs()
    if prefs.output_path != '':
        root_folder = prefs.output_path
    else:
        root_folder = os.path.dirname(bpy.data.filepath)
    os.makedirs(pjoin(root_folder, train), exist_ok=True)
    os.makedirs(pjoin(root_folder, test), exist_ok=True)
    os.makedirs(pjoin(root_folder, val), exist_ok=True)
    # Training images get written first
    bpy.context.scene.render.filepath = pjoin(root_folder, train, '')
    # Compute the assumed static camera information
    cam = scene.camera
    cam_angle_x = cam.data.angle_x # in radians
    out['camera_angle_x'] = cam_angle_x
    out['clip_start'] = cam.data.clip_start
    out['clip_end'] = cam.data.clip_end

# After a rendered frame is written
@persistent
def render_write_fn(scene):
    prefs = get_prefs()
    # Get the required file name and camera parameters
    global f
    f = scene.frame_current
    fout = {'id': f}
    imgfilepath = scene.render.frame_path(frame=f)
    imgfilename = bpy.path.basename(imgfilepath)
    # NeRF relative path doesn't include file extension!
    imgrelpath = pjoin('.', get_name(f), imgfilename.split('.')[0])
    # Change any Windows slashes to Unix slashes
    fout['file_path'] = sanitize_path(imgrelpath)
    mat = scene.camera.matrix_world
    # Blender's Matrix datatype is not directly JSON serializable
    fout['transform_matrix'] = [list(row) for row in mat]
    frames.append(fout)
    # train -> test -> val
    if f == prefs.train_end:
        out['frames'] = frames
        with open(pjoin(root_folder, train_json), 'w') as w:
            json.dump(out, w, indent=4)
        frames.clear()
        # Moving on to test images
        bpy.context.scene.render.filepath = pjoin(root_folder, test, '')
    elif f == prefs.test_end:
        out['frames'] = frames # overwrites last set of frames
        with open(pjoin(root_folder, test_json), 'w') as w:
            json.dump(out, w, indent=4)
        frames.clear()
        # Moving on to val images
        bpy.context.scene.render.filepath = pjoin(root_folder, val, '')
    # Add the novel view transform for this frame
    if prefs.novel_cam_name in scene.objects:
        novel_cam = scene.objects[prefs.novel_cam_name]
        mat = novel_cam.matrix_world
        nout = {
            'id': f,
            'transform_matrix': [list(row) for row in mat]
        }
        novel.append(nout)

# After all frames are rendered
@persistent
def render_complete_fn(scene):
    # Save out the camera parameters of the last set
    if is_train(f):
        final_json = pjoin(root_folder, train_json)
    elif is_test(f):
        final_json = pjoin(root_folder, test_json)
    elif is_val(f):
        final_json = pjoin(root_folder, val_json)
    else:
        print('Failure in render_complete_fn')
        final_json = pjoin(root_folder, 'failure.json')
    out['frames'] = frames # overwrites last set of frames
    with open(final_json, 'w') as w:
        json.dump(out, w, indent=4)
    # Save out novel view transforms
    if len(novel) > 0:
        nout = {'frames': novel}
        with open(pjoin(root_folder, novel_json), 'w') as w:
            json.dump(nout, w, indent=4)

### Register/unregister addon

def register():
    bpy.utils.register_class(NerfDataExporterPreferences)
    bpy.app.handlers.render_init.append(render_init_fn)
    bpy.app.handlers.render_write.append(render_write_fn)
    bpy.app.handlers.render_complete.append(render_complete_fn)

def unregister():
    bpy.utils.unregister_class(NerfDataExporterPreferences)
    bpy.app.handlers.render_init.remove(render_init_fn)
    bpy.app.handlers.render_write.remove(render_write_fn)
    bpy.app.handlers.render_complete.remove(render_complete_fn)