avoid code repetition and avoid hard coding the max_cpu_num limit

lvis/boundary_utils.py

@@ -7,13 +7,12 @@
 
 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
+        imgs (dict) : image dicts
 
     Returns:
         ann (rle)
@@ -40,6 +39,7 @@ def ann_to_mask(ann, imgs):
     to binary mask.
     Args:
         ann (dict) : annotation object
+        imgs (dict) : image dicts
 
     Returns:
         binary mask (numpy 2D array)
@@ -86,8 +86,8 @@ def augment_annotations_with_boundary_single_core(proc_id, annotations, imgs, di
     return new_annotations
 
 
-def augment_annotations_with_boundary_multi_core(annotations, imgs, dilation_ratio=0.02):
-    cpu_num = min(multiprocessing.cpu_count(), MAX_CPU_NUM)
+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)

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
@@ -62,7 +64,8 @@ def _create_index(self):
             tic = time.time()
             self.dataset["annotations"] = augment_annotations_with_boundary_multi_core(self.dataset["annotations"],
                                                                                        self.imgs,
-                                                                                       dilation_ratio=self.dilation_ratio)
+                                                                                       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.")