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)