Merge pull request #1154 from DLR-RM/iss_bop_singleproc

fix(bop): Adds option back to use bop writer via single processing

blenderproc/python/writer/BopWriterUtility.py

@@ -34,7 +34,7 @@ def write_bop(output_dir: str, target_objects: Optional[List[MeshObject]] = None
  depth_scale: float = 1.0, jpg_quality: int = 95, save_world2cam: bool = True,
  ignore_dist_thres: float = 100., m2mm: Optional[bool] = None, annotation_unit: str = 'mm',
  frames_per_chunk: int = 1000, calc_mask_info_coco: bool = True, delta: float = 0.015,
- num_worker: Optional[int] = None):
+ num_worker: Optional[int] = 0):
  """Write the BOP data
 
  :param output_dir: Path to the output directory.
@@ -60,6 +60,7 @@ def write_bop(output_dir: str, target_objects: Optional[List[MeshObject]] = None
  :param calc_mask_info_coco: Whether to calculate gt masks, gt info and gt coco annotations.
  :param delta: Tolerance used for estimation of the visibility masks (in [m]).
  :param num_worker: The number of processes to use to calculate gt_masks and gt_info. If None is given, number of cores is used.
+ If 0 is given, no multiprocessing at all is used (default).
  """
 
  # Output paths.
@@ -169,7 +170,11 @@ def write_bop(output_dir: str, target_objects: Optional[List[MeshObject]] = None
  # Create pool and init each worker
  width = bpy.context.scene.render.resolution_x
  height = bpy.context.scene.render.resolution_y
- pool = Pool(num_worker, initializer=_BopWriterUtility._pyrender_init, initargs=[width, height, trimesh_objects])
+ if num_worker == 0:
+ pool = None
+ _BopWriterUtility._pyrender_init(width, height, trimesh_objects)
+ else:
+ pool = Pool(num_worker, initializer=_BopWriterUtility._pyrender_init, initargs=[width, height, trimesh_objects])
 
  _BopWriterUtility.calc_gt_masks(chunk_dirs=chunk_dirs, starting_frame_id=starting_frame_id,
  annotation_scale=annotation_scale, delta=delta, pool=pool)
@@ -180,9 +185,12 @@ def write_bop(output_dir: str, target_objects: Optional[List[MeshObject]] = None
  _BopWriterUtility.calc_gt_coco(chunk_dirs=chunk_dirs, dataset_objects=dataset_objects,
  starting_frame_id=starting_frame_id)
 
- pool.close()
- pool.join()
-
+ if pool is not None:
+ pool.close()
+ pool.join()
+ else:
+ # Make sure the renderer get destroyed
+ _BopWriterUtility._pyrender_cleanup()
 
 
 def bop_pose_to_pyrender_coordinate_system(cam_R_m2c: np.ndarray, cam_t_m2c: np.ndarray) -> np.ndarray:
@@ -547,6 +555,16 @@ def _pyrender_init(ren_width: int, ren_height: int, trimesh_objects: Dict[int, t
  metallicFactor=0.2, roughnessFactor=0.8, doubleSided=True)
  dataset_objects[key] = pyrender.Mesh.from_trimesh(mesh=trimesh_objects[key], material=material)
 
+ @staticmethod
+ def _pyrender_cleanup():
+ """ Cleans up global renderer 
+ 
+ This is only necessary when not using multiprocessing,.
+ """
+ global renderer, renderer_large, dataset_objects
+ del renderer
+ del renderer_large
+ del dataset_objects
 
  @staticmethod
  def _calc_gt_masks_iteration(annotation_scale: float, K: np.ndarray, delta: float, dist_im: np.ndarray, chunk_dir: str, im_id: int, gt_data: Tuple[int, Dict[str, int]]):
@@ -667,8 +685,8 @@ def calc_gt_masks(pool: Pool, chunk_dirs: List[str], starting_frame_id: int = 0,
  depth_im /= 1000. # to [m]
  dist_im = misc.depth_im_to_dist_im_fast(depth_im, K)
 
- pool.map(partial(_BopWriterUtility._calc_gt_masks_iteration, annotation_scale, K, delta, dist_im, chunk_dir, im_id), enumerate(scene_gt[im_id]))
-
+ map_fun = map if pool is None else pool.map
+  list(map_fun(partial(_BopWriterUtility._calc_gt_masks_iteration, annotation_scale, K, delta, dist_im, chunk_dir, im_id), enumerate(scene_gt[im_id])))
 
 
  @staticmethod
@@ -834,7 +852,8 @@ def calc_gt_info(pool, chunk_dirs: List[str], starting_frame_id: int = 0,
 
  K = np.array(scene_camera[im_id]['cam_K']).reshape(3, 3)
 
- scene_gt_info[im_id] = pool.map(partial(_BopWriterUtility._calc_gt_info_iteration, annotation_scale, ren_cy_offset, ren_cx_offset, im_height, im_width, K, delta, depth), scene_gt[im_id])
+ map_fun = map if pool is None else pool.map
+ scene_gt_info[im_id] = list(map_fun(partial(_BopWriterUtility._calc_gt_info_iteration, annotation_scale, ren_cy_offset, ren_cx_offset, im_height, im_width, K, delta, depth), scene_gt[im_id]))
 
 
  # Save the info for the current scene.