diff --git a/blenderproc/python/writer/BopWriterUtility.py b/blenderproc/python/writer/BopWriterUtility.py index ab16a1ac0..3a748ed52 100644 --- a/blenderproc/python/writer/BopWriterUtility.py +++ b/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.