Reduce memory usage and time consumption with large cpu cores.

lvis/boundary_utils.py

@@ -7,6 +7,46 @@
 
 logger = logging.getLogger(__name__)
 
+MAX_CPU_NUM = 80
+
+
+def ann_to_rle(ann, imgs):
+    """Convert annotation which can be polygons, uncompressed RLE to RLE.
+    Args:
+        ann (dict) : annotation object
+
+    Returns:
+        ann (rle)
+    """
+    img_data = imgs[ann["image_id"]]
+    h, w = img_data["height"], img_data["width"]
+    segm = ann["segmentation"]
+    if isinstance(segm, list):
+        # polygon -- a single object might consist of multiple parts
+        # we merge all parts into one mask rle code
+        rles = mask_utils.frPyObjects(segm, h, w)
+        rle = mask_utils.merge(rles)
+    elif isinstance(segm["counts"], list):
+        # uncompressed RLE
+        rle = mask_utils.frPyObjects(segm, h, w)
+    else:
+        # rle
+        rle = ann["segmentation"]
+    return rle
+
+
+def ann_to_mask(ann, imgs):
+    """Convert annotation which can be polygons, uncompressed RLE, or RLE
+    to binary mask.
+    Args:
+        ann (dict) : annotation object
+
+    Returns:
+        binary mask (numpy 2D array)
+    """
+    rle = ann_to_rle(ann, imgs)
+    return mask_utils.decode(rle)
+
 
 # General util function to get the boundary of a binary mask.
 def mask_to_boundary(mask, dilation_ratio=0.02):
@@ -31,11 +71,11 @@ def mask_to_boundary(mask, dilation_ratio=0.02):
 
 
 # COCO/LVIS related util functions, to get the boundary for every annotations.
-def augment_annotations_with_boundary_single_core(proc_id, annotations, ann_to_mask, dilation_ratio=0.02):
+def augment_annotations_with_boundary_single_core(proc_id, annotations, imgs, dilation_ratio=0.02):
     new_annotations = []
 
     for ann in annotations:
-        mask = ann_to_mask(ann)
+        mask = ann_to_mask(ann, imgs)
         # Find mask boundary.
         boundary = mask_to_boundary(mask, dilation_ratio)
         # Add boundary to annotation in RLE format.
@@ -46,16 +86,16 @@ def augment_annotations_with_boundary_single_core(proc_id, annotations, ann_to_m
     return new_annotations
 
 
-def augment_annotations_with_boundary_multi_core(annotations, ann_to_mask, dilation_ratio=0.02):
-    cpu_num = multiprocessing.cpu_count()
+def augment_annotations_with_boundary_multi_core(annotations, imgs, dilation_ratio=0.02):
+    cpu_num = min(multiprocessing.cpu_count(), MAX_CPU_NUM)
     annotations_split = np.array_split(annotations, cpu_num)
     logger.info("Number of cores: {}, annotations per core: {}".format(cpu_num, len(annotations_split[0])))
     workers = multiprocessing.Pool(processes=cpu_num)
     processes = []
 
     for proc_id, annotation_set in enumerate(annotations_split):
         p = workers.apply_async(augment_annotations_with_boundary_single_core,
-                                (proc_id, annotation_set, ann_to_mask, dilation_ratio))
+                                (proc_id, annotation_set, imgs, dilation_ratio))
         processes.append(p)
 
     new_annotations = []

lvis/lvis.py

@@ -61,7 +61,7 @@ def _create_index(self):
             self.logger.info('Adding `boundary` to annotation.')
             tic = time.time()
             self.dataset["annotations"] = augment_annotations_with_boundary_multi_core(self.dataset["annotations"],
-                                                                                       self.ann_to_mask,
+                                                                                       self.imgs,
                                                                                        dilation_ratio=self.dilation_ratio)
 
             self.logger.info('`boundary` added! (t={:0.2f}s)'.format(time.time()- tic))