Explicit method for adding boundaries

boundary_iou/coco_instance_api/coco.py

@@ -322,13 +322,13 @@ def showAnns(self, anns, draw_bbox=False):
             for ann in anns:
                 print(ann['caption'])
 
-    def loadRes(self, resFile):
+    def loadRes(self, resFile, get_boundary=False):
         """
         Load result file and return a result api object.
         :param   resFile (str)     : file name of result file
         :return: res (obj)         : result api object
         """
-        res = COCO(get_boundary=self.get_boundary, dilation_ratio=self.dilation_ratio)
+        res = COCO(get_boundary=get_boundary, dilation_ratio=self.dilation_ratio)
         res.dataset['images'] = [img for img in self.dataset['images']]
 
         print('Loading and preparing results...')

boundary_iou/coco_instance_api/cocoeval.py

@@ -203,7 +203,7 @@ def computeIoU(self, imgId, catId):
             raise Exception('unknown iouType for iou computation')
 
         # compute iou between each dt and gt region
-        iscrowd = [int(o['iscrowd']) for o in gt]
+        iscrowd = [int(o.get('iscrowd', 0)) for o in gt]
         ious = maskUtils.iou(d,g,iscrowd)
         return ious
 
@@ -231,7 +231,7 @@ def computeBoundaryIoU(self, imgId, catId):
         d_b = [d['boundary'] for d in dt]
 
         # compute iou between each dt and gt region
-        iscrowd = [int(o['iscrowd']) for o in gt]
+        iscrowd = [int(o.get('iscrowd', 0)) for o in gt]
         mask_ious = maskUtils.iou(d_m,g_m,iscrowd)
         boundary_ious = maskUtils.iou(d_b,g_b,iscrowd)
         # combine mask and boundary iou
@@ -314,7 +314,7 @@ def evaluateImg(self, imgId, catId, aRng, maxDet):
         gt = [gt[i] for i in gtind]
         dtind = np.argsort([-d['score'] for d in dt], kind='mergesort')
         dt = [dt[i] for i in dtind[0:maxDet]]
-        iscrowd = [int(o['iscrowd']) for o in gt]
+        iscrowd = [int(o.get('iscrowd', 0)) for o in gt]
         # load computed ious
         ious = self.ious[imgId, catId][:, gtind] if len(self.ious[imgId, catId]) > 0 else self.ious[imgId, catId]
 
@@ -431,8 +431,8 @@ def accumulate(self, p = None):
                     tps = np.logical_and(               dtm,  np.logical_not(dtIg) )
                     fps = np.logical_and(np.logical_not(dtm), np.logical_not(dtIg) )
 
-                    tp_sum = np.cumsum(tps, axis=1).astype(dtype=np.float)
-                    fp_sum = np.cumsum(fps, axis=1).astype(dtype=np.float)
+                    tp_sum = np.cumsum(tps, axis=1).astype(dtype=float)
+                    fp_sum = np.cumsum(fps, axis=1).astype(dtype=float)
                     for t, (tp, fp) in enumerate(zip(tp_sum, fp_sum)):
                         tp = np.array(tp)
                         fp = np.array(fp)

boundary_iou/utils/boundary_utils.py

@@ -1,13 +1,15 @@
 from collections import defaultdict
 import multiprocessing
 import math
+import time
+from pathlib import Path
+import json
 
 import cv2
 import numpy as np
 
 import pycocotools.mask as mask_utils
 
-
 # General util function to get the boundary of a binary mask.
 def mask_to_boundary(mask, dilation_ratio=0.02):
     """
@@ -42,7 +44,7 @@ def augment_annotations_with_boundary_single_core(proc_id, annotations, ann_to_m
         ann['boundary'] = mask_utils.encode(
             np.array(boundary[:, :, None], order="F", dtype="uint8"))[0]
         new_annotations.append(ann)
-    
+
     return new_annotations
 
 
@@ -57,12 +59,69 @@ def augment_annotations_with_boundary_multi_core(annotations, ann_to_mask, dilat
         p = workers.apply_async(augment_annotations_with_boundary_single_core,
                                 (proc_id, annotation_set, ann_to_mask, dilation_ratio))
         processes.append(p)
-    
+
     new_annotations = []
     for p in processes:
         new_annotations.extend(p.get())
-    
+
     workers.close()
     workers.join()
-
-    return new_annotations
+
+    return new_annotations
+
+
+def add_boundary_multi_core(coco, cpu_num=16, dilation_ratio=0.02):
+    if coco.get_boundary == True:
+        print('Found existing boundaries, skipping...')
+        return
+
+    print('Adding `boundary` to annotation.')
+    tic = time.time()
+    cpu_num = min(cpu_num, multiprocessing.cpu_count())
+
+    annotations = coco.dataset["annotations"]
+    annotations_split = np.array_split(annotations, cpu_num)
+    print("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, coco.annToMask, dilation_ratio))
+        processes.append(p)
+
+    new_annotations = []
+    for p in processes:
+        new_annotations.extend(p.get())
+
+    workers.close()
+    workers.join()
+
+    coco.dataset["annotations"] = new_annotations
+    coco.createIndex()
+    coco.get_boundary = True
+    print('`boundary` added! (t={:0.2f}s)'.format(time.time()- tic))
+
+
+def coco_add_boundaries_and_save_as_annotation_file(annotation_file: str, stem_addon='_b'):
+    # NOTE: we only need coco at the moment, feel free to generalize this function
+    from ..coco_instance_api.coco import COCO
+
+    annotation_file = Path(annotation_file).expanduser().resolve()
+    coco = COCO(annotation_file)
+    add_boundary_multi_core(coco)
+
+    with open(annotation_file, 'r') as f:
+        coco_json = json.load(f)
+
+    # boundary counts are byte strings, we need to decode them so they are json serializable
+    anns_b = list(coco.anns.values())
+    for a in anns_b:
+        a['boundary']['counts'] = a['boundary']['counts'].decode("utf-8")
+    coco_json['annotations'] = anns_b
+
+    new_annotation_file = annotation_file.parent / (annotation_file.stem + stem_addon + annotation_file.suffix)
+    with open(new_annotation_file, 'w', encoding='utf-8') as f:
+        json.dump(coco_json, f)
+
+    print(f"\nAdded boundaries to ({annotation_file.parent}) '{annotation_file.name}' and saved as: '{new_annotation_file.name}'")