Merge pull request #27 from shenyunhang/lvis_challenge_2021

Reduce memory usage and time consumption with large cpu cores.

lvis/boundary_utils.py

@@ -8,6 +8,46 @@
 logger = logging.getLogger(__name__)
 
 
+def ann_to_rle(ann, imgs):
+    """Convert annotation which can be polygons, uncompressed RLE to RLE.
+    Args:
+        ann (dict) : annotation object
+        imgs (dict) : image dicts
+
+    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
+        imgs (dict) : image dicts
+
+    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, max_cpu_num=80):
+    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/eval.py

@@ -7,23 +7,25 @@
 
 from lvis.lvis import LVIS
 from lvis.results import LVISResults
+from lvis.boundary_utils import ann_to_rle
 
 import pycocotools.mask as mask_utils
 
 
 class LVISEval:
-    def __init__(self, lvis_gt, lvis_dt, iou_type="segm", mode="default", dilation_ratio=0.02):
+    def __init__(self, lvis_gt, lvis_dt, iou_type="segm", mode="default", dilation_ratio=0.02, max_cpu_num=80):
         """Constructor for LVISEval.
         Args:
             lvis_gt (LVIS class instance, or str containing path of annotation file)
             lvis_dt (LVISResult class instance, or str containing path of result file,
             or list of dict)
             iou_type (str): segm, bbox, or boundary evaluation. Ignored if `mode` is set to
                 'challenge2021'.
-            dilation_ratio (float): ratio to calculate dilation = dilation_ratio * image_diagonal
             mode (str): Either 'default' or 'challenge2021'. Specifying 'challenge2021'
                 uses iou_type=boundary and limits detections to 10,000 per class
                 (instead of 300 per image).
+            dilation_ratio (float): ratio to calculate dilation = dilation_ratio * image_diagonal
+            max_cpu_num (int): max number of cpu cores to compute mask boundary before evaluation
         """
         self.logger = logging.getLogger(__name__)
 
@@ -73,12 +75,14 @@ def __init__(self, lvis_gt, lvis_dt, iou_type="segm", mode="default", dilation_r
             if not self.lvis_gt.precompute_boundary:
                 self.lvis_gt.precompute_boundary = self.use_boundary_iou
                 self.lvis_gt.dilation_ratio = dilation_ratio
+                self.lvis_gt.max_cpu_num = max_cpu_num
                 self.lvis_gt._create_index()
             else:
                 assert self.lvis_gt.dilation_ratio == dilation_ratio, "Dilation ratio not consistent"
             if not self.lvis_dt.precompute_boundary:
                 self.lvis_dt.precompute_boundary = self.use_boundary_iou
                 self.lvis_dt.dilation_ratio = dilation_ratio
+                self.lvis_dt.max_cpu_num = max_cpu_num
                 self.lvis_dt._create_index()
             else:
                 assert self.lvis_gt.dilation_ratio == dilation_ratio, "Dilation ratio not consistent"
@@ -100,7 +104,7 @@ def __init__(self, lvis_gt, lvis_dt, iou_type="segm", mode="default", dilation_r
 
     def _to_mask(self, anns, lvis):
         for ann in anns:
-            rle = lvis.ann_to_rle(ann)
+            rle = ann_to_rle(ann, lvis.imgs)
             ann["segmentation"] = rle
 
     def _prepare(self):

lvis/lvis.py

@@ -19,12 +19,13 @@
 
 
 class LVIS:
-    def __init__(self, annotation_path, precompute_boundary=False, dilation_ratio=0.02):
+    def __init__(self, annotation_path, precompute_boundary=False, dilation_ratio=0.02, max_cpu_num=80):
         """Class for reading and visualizing annotations.
         Args:
             annotation_path (str): location of annotation file
             precompute_boundary (bool): whether to precompute mask boundary before evaluation
             dilation_ratio (float): ratio to calculate dilation = dilation_ratio * image_diagonal
+            max_cpu_num (int): max number of cpu cores to compute mask boundary before evaluation
         """
         self.logger = logging.getLogger(__name__)
         self.logger.info("Loading annotations.")
@@ -33,6 +34,7 @@ def __init__(self, annotation_path, precompute_boundary=False, dilation_ratio=0.
 
         self.precompute_boundary = precompute_boundary
         self.dilation_ratio = dilation_ratio
+        self.max_cpu_num = max_cpu_num
 
         assert (
             type(self.dataset) == dict
@@ -61,8 +63,9 @@ 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,
-                                                                                       dilation_ratio=self.dilation_ratio)
+                                                                                       self.imgs,
+                                                                                       dilation_ratio=self.dilation_ratio,
+                                                                                       max_cpu_num=self.max_cpu_num)
 
             self.logger.info('`boundary` added! (t={:0.2f}s)'.format(time.time()- tic))
 
@@ -185,39 +188,3 @@ def download(self, save_dir, img_ids=None):
             file_name = os.path.join(save_dir, img["coco_url"].split("/")[-1])
             if not os.path.exists(file_name):
                 urlretrieve(img["coco_url"], file_name)
-
-    def ann_to_rle(self, ann):
-        """Convert annotation which can be polygons, uncompressed RLE to RLE.
-        Args:
-            ann (dict) : annotation object
-
-        Returns:
-            ann (rle)
-        """
-        img_data = self.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(self, ann):
-        """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 = self.ann_to_rle(ann)
-        return mask_utils.decode(rle)

lvis/results.py

@@ -15,6 +15,7 @@ def __init__(
         max_dets_per_im=300,
         precompute_boundary=False,
         dilation_ratio=0.02,
+        max_cpu_num=80,
     ):
         """Constructor for LVIS results.
         Args:
@@ -31,6 +32,7 @@ def __init__(
                 (i.e., -1).
             precompute_boundary (bool): whether to precompute mask boundary before evaluation
             dilation_ratio (float): ratio to calculate dilation = dilation_ratio * image_diagonal
+            max_cpu_num (int): max number of cpu cores to compute mask boundary before evaluation
         """
         if isinstance(lvis_gt, LVIS):
             self.dataset = deepcopy(lvis_gt.dataset)
@@ -42,6 +44,7 @@ def __init__(
 
         self.precompute_boundary = precompute_boundary
         self.dilation_ratio = dilation_ratio
+        self.max_cpu_num = max_cpu_num
 
         self.logger = logging.getLogger(__name__)
         self.logger.info("Loading and preparing results.")