Added optional data fields to JSON file. Fixed rotation bug. Added da…

…ta validation script

data_generation/blenderproc_data_gen/README.md

@@ -14,15 +14,20 @@ If you run into troubles, please consult the [project's own github page](https:/
 
 
 
-## Usage example:
+### Usage example:
 
-Run the blenderproc script 5 times, each time generating 1000 frames. Each frame will have five copies of the object and ten randomly chosen distractor objects:
+Run the blenderproc script in five parallel jobs, each generating 1000 frames. Each frame will have six instances of the object and ten randomly chosen distractor objects:
 ```
-./run_blenderproc_datagen.py --nb_runs 1 --nb_frames 10 --path_single_obj ../models/Ketchup/google_16k/textured.obj --nb_objects 5 --distractors_folder ~/data/google_scanned_models/ --nb_distractors 10 --backgrounds_folder ../dome_hdri_haven/
+./run_blenderproc_datagen.py --nb_runs 5 --nb_frames 1000 --path_single_obj ../models/Ketchup/google_16k/textured.obj --nb_objects 6 --distractors_folder ~/data/google_scanned_models/ --nb_distractors 10 --backgrounds_folder ../dome_hdri_haven/
 ```
 
-All parameters can be shown by running
+Parameters of the top-level script can be shown by running
 ```
 python ./run_blenderproc_datagen.py --help
 ```
-Note that, as a blenderproc script, `generate_training_data.py` cannot be invoked with Python. It must be run via the `blenderproc` launch script.
+
+Note that, as a blenderproc script, `generate_training_data.py` cannot be invoked with Python. It must be run via the `blenderproc` launch script.  To discover its command-line parameters, you must look
+at the source-code itself; `blenderproc run ./generate_training_data.py --help` will not report
+them properly.
+
+Blenderproc searches for python modules in a different order than when invoking Python by itself.  If you run into an issue where `generate_training_data.py` fails to import modules that you have installed, you may have to re-install them via blenderproc; e.g. `blenderproc pip install pyquaternion`.

data_generation/blenderproc_data_gen/generate_training_data.py

@@ -4,7 +4,6 @@
 from blenderproc.python.utility.Utility import Utility
 import bpy
 
-
 import argparse
 import cv2
 import glob
@@ -13,9 +12,11 @@
 import numpy as np
 import os
 from PIL import Image, ImageDraw
+from pyquaternion import Quaternion
 import random
 import sys
 
+
 def random_object_position(near=5.0, far=40.0):
     # Specialized function to randomly place the objects in a visible
     # location
@@ -50,15 +51,13 @@ def random_depth_in_frustrum(tw, th, bw, bh, depth):
 
 
 def point_in_frustrum(camera, near=10, far=20):
-     fov_w, fov_h = camera.get_fov()
+    fov_w, fov_h = camera.get_fov()
 
     tw = sin(fov_w)*near # top (nearest to camera) width of frustrum
     th = sin(fov_h)*near # top (nearest to camera) height of frustrum
     bw = sin(fov_w)*far  # bottom width
     bh = sin(fov_h)*far  # bottom height
 
-    print(tw, th, bw, bh)
-
     # calculate random inverse depth: 0 at the 'far' plane and 1 at the 'near' plane
     inv_depth = random_depth_in_frustrum(tw, th, bw, bh, far-near)
     depth = far-inv_depth
@@ -182,9 +181,9 @@ def get_cuboid_image_space(mesh, camera):
     # object aligned bounding box coordinates in world coordinates
     bbox = mesh.get_bound_box()
     '''
-    bbox is a list of the world-space coordinates of the corners of the
-    object's oriented bounding box
-    https://blender.stackexchange.com/questions/32283/what-are-all-values-in-bound-box
+    bbox is a list of the world-space coordinates of the corners of a
+    blender object's oriented bounding box
+     https://blender.stackexchange.com/questions/32283/what-are-all-values-in-bound-box
 
                    TOP
            3 +-----------------+ 7
@@ -201,7 +200,6 @@ def get_cuboid_image_space(mesh, camera):
         1 +-----------------+ 5
                 FRONT
 
-    Point '8' (the ninth entry) is the centroid
     '''
 
     centroid = np.array([0.,0.,0.])
@@ -217,7 +215,7 @@ def get_cuboid_image_space(mesh, camera):
 
     # However these points are in a different order than the original DOPE data format,
     # so we must reorder them
-    dope_order = [5, 1, 2, 6, 4, 0, 3, 7]
+    dope_order = [6, 2, 1, 5, 7, 3, 0, 4]
     cuboid = [None for ii in range(9)]
     for ii in range(8):
         cuboid[dope_order[ii]] = cv2.projectPoints(bbox[ii], rvec, tvec, K, np.array([]))[0][0][0]
@@ -281,7 +279,11 @@ def write_json(outf, args, camera, objects, objects_data, seg_map):
             'class': objects_data[ii]['class'],
             'name': objects_data[ii]['name'],
             'visibility': num_pixels,
-            'projected_cuboid': projected_keypoints
+            'projected_cuboid': projected_keypoints,
+            ## 'location' and 'quaternion_xyzw' are both optional data fields,
+            ## not used for training
+            'location': objects_data[ii]['location'],
+            'quaternion_xyzw': objects_data[ii]['quaternion_xyzw']
         })
 
     with open(outf, "w") as write_file:
@@ -470,13 +472,21 @@ def main(args):
         #                    150+random.random()*100])
 
         # Place object(s)
-        for oo in objects:
+        for idx, oo in enumerate(objects):
             # Set a random pose
             xform = np.eye(4)
             xform[0:3,3] = random_object_position(near=20, far=100)
             xform[0:3,0:3] = random_rotation_matrix()
             oo.set_local2world_mat(xform)
 
+            # 'location' and 'quaternion_xyzw' describe the position and orientation of the
+            # object in the camera coordinate system
+            xform_in_cam = np.linalg.inv(bp.camera.get_camera_pose()) @ xform
+            objects_data[idx]['location'] = xform_in_cam[0:3,3].tolist()
+            tmp_wxyz = Quaternion(matrix=xform_in_cam[0:3,0:3]).elements  # [scalar, x, y, z]
+            q_xyzw = [tmp_wxyz[1], tmp_wxyz[2], tmp_wxyz[3], tmp_wxyz[0]] # [x, y, z, scalar]
+            objects_data[idx]['quaternion_xyzw'] = q_xyzw
+
             # Scale 3D model to cm
             oo.set_scale([args.scale, args.scale, args.scale])
 
@@ -529,7 +539,7 @@ def main(args):
                 im = background
 
         if args.debug:
-            im = draw_cuboid_markers(objects, bp.camera, background)
+            im = draw_cuboid_markers(objects, bp.camera, im)
 
         filename = os.path.join(out_directory, str(frame).zfill(6) + ".png")
         im.save(filename)

data_generation/validate_data.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+
+import json
+import numpy as np
+import os
+from PIL import Image, ImageDraw
+from pyquaternion import Quaternion
+import sys
+
+
+def main(json_files):
+    for json_fn in json_files:
+        # Find corresponding PNG
+        base, _ = os.path.splitext(json_fn)
+        img_fn = base+'.png'
+        if not os.path.isfile(img_fn):
+            print(f"Could not locate '{img_fn}'. Skipping..")
+            continue
+
+        # Load JSON data
+        with open(json_fn, 'r') as F:
+            data_json = json.load(F)
+        up = np.array(data_json['camera_data']['camera_look_at']['up'])
+        at = np.array(data_json['camera_data']['camera_look_at']['at'])
+        eye = np.array(data_json['camera_data']['camera_look_at']['eye'])
+
+        cam_matrix = np.eye(4)
+        cam_matrix[0:3,0] = up
+        cam_matrix[0:3,1] = np.cross(up, -at)
+        cam_matrix[0:3,2] = -at
+        cam_matrix[0:3,3] = -eye
+
+        img = Image.open(img_fn)
+
+        objects = data_json['objects']
+        # draw projected cuboid dots
+        for oo in objects:
+            draw = ImageDraw.Draw(img)
+            pts = oo['projected_cuboid']
+            for pt in pts:
+                draw.ellipse((pt[0]-2, pt[1]-2, pt[0]+2, pt[1]+2), fill = 'cyan',
+                             outline ='cyan')
+
+            line_order = [[0, 1], [1, 2], [3, 2], [3, 0], # front
+                          [4, 5], [6, 5], [6, 7], [4, 7], # back
+                          [0, 4], [7, 3], [5, 1], [2, 6], # sides
+                          [0, 5], [1,4]]                  # 'x' on top
+            for ll in line_order:
+                draw.line([(pts[ll[0]][0],pts[ll[0]][1]), (pts[ll[1]][0],pts[ll[1]][1])],
+                           fill='cyan', width=1)
+
+        img.save(base+'-output.png')
+
+
+def usage_msg(script_name):
+    print(f"Usage: {script_name} _JSON FILES_")
+    print("  The basename of the JSON files in _FILES_ will be used to find its")
+    print("  corresponding image file; i.e. if `00001.json` is provided, the code")
+    print("  will look for an image named `00001.png`")
+
+
+if __name__ == "__main__":
+    # Print out usage information if there are no arguments
+    if len(sys.argv) < 2:
+        usage_msg(sys.argv[0])
+        exit(0)
+
+    # ..or if the first argument is a request for help
+    s = sys.argv[1].lstrip('-')
+    if s == "h" or s == "help":
+        usage_msg(sys.argv[0])
+        exit(0)
+
+    main(sys.argv[1:])
+
+