forked custom scene replication script and adjusted boploader for xyzibd

Author: maivnov

blenderproc/python/loader/BopLoader.py

@@ -122,32 +122,64 @@ def load_bop_scene(bop_dataset_path: str, scene_id: int, model_type: str = "", c
     """
 
     bop_path, bop_dataset_name = _BopLoader.setup_bop_toolkit(bop_dataset_path)
-
+    
     # This import is done inside to avoid having the requirement that BlenderProc depends on the bop_toolkit
     # pylint: disable=import-outside-toplevel
     from bop_toolkit_lib import dataset_params, inout
-
+    
     # pylint: enable=import-outside-toplevel
 
     if source_frame is None:
         source_frame = ["X", "-Y", "-Z"]
 
     model_p = dataset_params.get_model_params(bop_path, bop_dataset_name, model_type=model_type if model_type else None)
+    
+    # TODO: uncomment to enable debug output
+    """
+    print("------- Model Params -------")
+    for k,v in model_p.items():
+        print(f"[{k}]: {v}")
+    print("----------------------------")"""
+    
     try:
         split_p = dataset_params.get_split_params(bop_path, bop_dataset_name, split=split,
                                                   split_type=cam_type if cam_type else None)
+            
     except ValueError as e:
         raise RuntimeError(f"Wrong path or {split} split does not exist in {bop_dataset_path}.") from e
+    
+    # TODO: uncomment to enable debug output
+    """
+    print("------- Split Params -------")
+    for k,v in split_p.items():
+        if v is not None and "xyzibd" in v:
+            rel_v_path = v.split("xyzibd")[1]
+        else:
+            rel_v_path = v
+        print(f"[{k}]: {rel_v_path}")
+        
+        if rel_v_path is not None and "scene_gt.json" in rel_v_path:
+            print("found it")
+            break
+    print("----------------------------")"""
+    
+    # add actual paths by hand, since dataset_params.get_split_params provides no acceess to internal structure
+    if bop_dataset_name == "xyzibd":
+        split_p["scene_gt_tpath"] = os.path.join(bop_path, bop_dataset_name, "train", "{scene_id:06d}", "scene_gt.json")
+        split_p["scene_camera_tpath"] = os.path.join(bop_path, bop_dataset_name, "train", "{scene_id:06d}", "scene_camera.json")
+    
     sc_gt = inout.load_scene_gt(split_p['scene_gt_tpath'].format(**{'scene_id': scene_id}))
     sc_camera = inout.load_json(split_p['scene_camera_tpath'].format(**{'scene_id': scene_id}))
-
+    
     assert object_model_unit in ['m', 'dm', 'cm', 'mm'], (f"Invalid object model unit: `{object_model_unit}`. "
                                                           f"Supported are 'm', 'dm', 'cm', 'mm'")
     scale = {'m': 1., 'dm': 0.1, 'cm': 0.01, 'mm': 0.001}[object_model_unit]
     if mm2m is not None:
         warnings.warn("WARNING: `mm2m` is deprecated, please use `object_model_unit='mm'` instead!")
         scale = 0.001
-
+    
+    
+    camera_poses = []
     for i, (cam_id, insts) in enumerate(sc_gt.items()):
         cam_K, cam_H_m2c_ref = _BopLoader.get_ref_cam_extrinsics_intrinsics(sc_camera, cam_id, insts, scale)
 
@@ -162,8 +194,13 @@ def load_bop_scene(bop_dataset_path: str, scene_id: int, model_type: str = "", c
                 _BopLoader.set_object_pose(cur_objs[-1], inst, scale)
 
         cam_H_c2w = _BopLoader.compute_camera_to_world_trafo(cam_H_m2w_ref, cam_H_m2c_ref, source_frame)
+        camera_poses.append(cam_H_c2w)
         # set camera intrinsics
-        CameraUtility.set_intrinsics_from_K_matrix(cam_K, split_p['im_size'][0], split_p['im_size'][1])
+        if bop_dataset_name == "xyzibd":
+            # Quick and dirty fix for xyzibd
+            CameraUtility.set_intrinsics_from_K_matrix(cam_K, split_p['im_size']['xyz'][0], split_p['im_size']['xyz'][1])
+        else:
+            CameraUtility.set_intrinsics_from_K_matrix(cam_K, split_p['im_size'][0], split_p['im_size'][1])
 
         # set camera extrinsics as next frame
         frame_id = CameraUtility.add_camera_pose(cam_H_c2w)
@@ -176,10 +213,11 @@ def load_bop_scene(bop_dataset_path: str, scene_id: int, model_type: str = "", c
         # Copy object poses to key frame (to be sure)
         for cur_obj in cur_objs:
             _BopLoader.insert_key_frames(cur_obj, frame_id)
+        
+    return cur_objs, camera_poses
 
-    return cur_objs
-
-
+        
+        
 def load_bop_intrinsics(bop_dataset_path: str, split: str = "test", cam_type: str = "") -> Tuple[np.ndarray, int, int]:
     """
     Load and set the camera matrix and image resolution of a specified BOP dataset
@@ -217,6 +255,7 @@ def load_bop_intrinsics(bop_dataset_path: str, split: str = "test", cam_type: st
     return cam_p['K'], split_p['im_size'][0], split_p['im_size'][1]
 
 
+
 class _BopLoader:
     CACHED_OBJECTS = {}
     @staticmethod
@@ -256,10 +295,10 @@ def compute_camera_to_world_trafo(cam_H_m2w_ref: np.array, cam_H_m2c_ref: np.arr
         """
 
         cam_H_c2w = np.dot(cam_H_m2w_ref, np.linalg.inv(cam_H_m2c_ref))
-
-        print('-----------------------------')
+        # TODO: uncomment
+        """print('-----------------------------')
         print(f"Cam: {cam_H_c2w}")
-        print('-----------------------------')
+        print('-----------------------------')"""
 
         # transform from OpenCV to blender coords
         cam_H_c2w = change_source_coordinate_frame_of_transformation_matrix(cam_H_c2w, source_frame)
@@ -371,3 +410,4 @@ def load_mesh(obj_id: int, model_p: dict, bop_dataset_name: str, scale: float =
         cur_obj.set_cp("bop_dataset_name", bop_dataset_name)
 
         return cur_obj
+

examples/advanced/rerender_edges/main.py

@@ -0,0 +1,80 @@
+import blenderproc as bproc
+from pathlib import Path
+import argparse
+import os
+
+# This script loads a BOP dataset scene and rerenders it with freestyle edges
+# Then write the results back using the BOP folder structure
+
+def normalize_path(path_str: str) -> Path:
+    """
+    Normalize a filesystem path in an OS-agnostic way.
+    
+    - Expands user (~) and environment variables.
+    - Converts relative paths to absolute.
+    - Normalizes slashes and resolves ".." and "." segments.
+    """
+    if not path_str:
+        return None
+    return Path(path_str).expanduser().resolve()
+
+parser = argparse.ArgumentParser()
+parser.add_argument('bop_parent_path', nargs='?', help="Path to the bop datasets parent directory")
+parser.add_argument('bop_dataset_name', nargs='?', help="Main BOP dataset")
+parser.add_argument('output_dir', nargs='?', help="Path to where the final files will be saved ")
+args = parser.parse_args()
+
+# Normalize paths to ensure compatibility across different operating systems
+args.bop_parent_path = normalize_path(args.bop_parent_path)
+args.output_dir = normalize_path(args.output_dir)
+
+bproc.init()
+
+# load specified bop objects into the scene
+
+render_id = 0
+while True:
+    print(f"Rendering scene id {render_id}...")
+    try:
+        bop_objs, camera_poses = bproc.loader.load_bop_scene(
+            bop_dataset_path = os.path.join(args.bop_parent_path, args.bop_dataset_name),
+            object_model_unit = "mm",
+            scene_id = render_id,
+            split = 'train')
+
+        # set shading
+        print("----------Object Info----------")
+        for j, obj in enumerate(bop_objs):
+            obj.set_shading_mode('auto')
+            print(j, obj.get_name())
+        print("-------------------------------")
+                
+        # Set light source
+        light_point = bproc.types.Light()
+        light_point.set_energy(1000)
+        light_point.set_location([0, 0, -0.8])
+
+        # activate depth rendering
+        bproc.renderer.enable_depth_output(activate_antialiasing=False)
+        bproc.renderer.set_max_amount_of_samples(50)
+
+        bproc.renderer.set_render_devices(desired_gpu_device_type=["HIP", "CUDA"])
+        # render the cameras of the current scene
+        data = bproc.renderer.render(verbose=True)
+
+        # Render the edges
+        data["edges"] = bproc.renderer.render_edges(
+            target_objects=bop_objs, camera_poses=camera_poses
+        )
+
+        # Write data to bop format
+        bproc.writer.write_bop(os.path.join(args.output_dir, 'bop_data'),
+                            dataset=args.bop_dataset_name,
+                            depths = data["depth"],
+                            colors=data["colors"],
+                            edges=data["edges"],
+                            save_world2cam=False) # world coords are arbitrary in most real BOP datasets
+        render_id += 1
+    except Exception as e:
+        print(f"Finished rendering all scenes. Stopped at scene id {render_id}.")
+        break