update rendering script

utils/blender_cam.py

@@ -1,5 +1,7 @@
+import os
 import bpy
 import json
+import argparse
 import numpy as np
 
 def add_lighting():
@@ -8,14 +10,23 @@ def add_lighting():
     bpy.ops.object.light_add(type="POINT", location=(0.80797, -7.77557, 4.78247))
     bpy.ops.object.light_add(type="POINT", location=(-4.96121, 1.9155, 9.01307))
 
+def remove_cube():
+    bpy.ops.object.select_all(action='DESELECT')
+    bpy.data.objects['Cube'].select_set(True)
+    bpy.ops.object.delete()
+
 
 def camera_setting(exp_dir, n_frames=100, stage='train', track_to='start'):
     '''Set the render setting for the camera and the scene'''
     bpy.context.scene.render.engine = 'BLENDER_WORKBENCH'
 
+    remove_cube()
+
     add_lighting()
 
     cam_obj = bpy.data.objects['Camera']
+    constraint = cam_obj.constraints.new('TRACK_TO')
+    constraint.target = bpy.data.objects[track_to]
     cam_obj.constraints["Track To"].target = bpy.data.objects[track_to]
     bpy.context.scene.frame_end = n_frames - 1
     bpy.context.scene.frame_start = 0
@@ -123,26 +134,45 @@ def get_camera_dict():
     return cam_dict
 
 
-def camera_export(stage='train'):
-    file_path = f'{ROOT}/camera_{stage}.json'
+def camera_export(stage='train', dst_dir=''):
+    file_path = f'{dst_dir}/camera_{stage}.json'
     with open(file_path, 'w') as fh:
         cam = get_camera_dict()
         json.dump(cam, fh)
         fh.close()
 
 if __name__ == "__main__":
-    # render configs
-    ROOT = '/path/to/save/rendered/images/and/camera/parameters/'
-    n_frames = 100  # number of frames to render
-    stage = 'train' # ['train', 'val', 'test']
-    track_to = 'start'  # the object name that camera should track to
+    '''
+    Script to render images and export camera parameters using Blender (tested on version 4.1.0)
+    
+    Structure of the output folder:
+    <dst_dir>
+        ├── <stage>
+        │   ├── 0000.png
+        │   ├── 0001.png
+        │   ├── ...
+        │── camera_<stage>.json
+    '''
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--src_dir', type=str, required=True, help='folder containing .obj files')
+    parser.add_argument('--dst_dir', type=str, required=True, help='folder to save rendered images and camera parameters')
+    parser.add_argument('--n_frames', type=int, default=100, help='number of frames to render')
+    parser.add_argument('--stage', type=str, default='train', choices=['train', 'val', 'test'], help='render images for which set')
+    args = parser.parse_args()
+
+    assert os.path.exists(args.src_dir), f'{args.src_dir} does not exist'
+    os.makedirs(args.dst_dir, exist_ok=True)
 
     # STEP 0: Load .obj files into the blender scene
+    bpy.ops.wm.obj_import(filepath=args.src_dir)
 
     # STEP 1: Config the camera setting
-    camera_setting(exp_dir=ROOT, n_frames=n_frames, stage=stage, track_to=track_to)
+    fname = os.path.basename(args.src_dir)
+    track_to = fname[:-4]  # the object name that camera should track to
+    camera_setting(exp_dir=args.dst_dir, n_frames=args.n_frames, stage=args.stage, track_to=track_to)
+
     # STEP 2: Export the camera parameters
-    camera_export(stage)
+    camera_export(args.stage, args.dst_dir)
 
     # STEP 3: Render Images
     bpy.ops.render.render(animation=True)