add category based annotation area filtering (#86)

* add intervals_per_category argument to get_area_filtered_coco function * add docs for intervals_per_category argument
obss · May 8, 2021 · d0f8e4b · d0f8e4b
1 parent 02c50e9
commit d0f8e4b
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/docs/COCO.md b/docs/COCO.md
@@ -136,6 +136,12 @@ coco = Coco.from_coco_dict_or_path("coco.json")
 area_filtered_coco = coco.get_area_filtered_coco(min=50)
 # filter out images that contain annotations with smaller area than 50 and larger area than 10000
 area_filtered_coco = coco.get_area_filtered_coco(min=50, max=10000)
+# filter out images with seperate area intervals per category
+intervals_per_category = {
+  "human": {"min": 20, "max": 10000},
+  "vehicle": {"min": 50, "max": 15000},
+}
+area_filtered_coco = coco.get_area_filtered_coco(intervals_per_category=intervals_per_category)
 
 # export filtered COCO dataset
 save_json(area_filtered_coco.json, "area_filtered_coco.json")

diff --git a/sahi/utils/coco.py b/sahi/utils/coco.py
@@ -1236,7 +1236,7 @@ def get_subsampled_coco(self, subsample_ratio=2):
 
         return subsampled_coco
 
-    def get_area_filtered_coco(self, min=0, max=float('inf')):
+    def get_area_filtered_coco(self, min=0, max=float('inf'), intervals_per_category=None):
         """
         Filters annotation areas with given min and max values and returns remaining
         images as sahi.utils.coco.Coco object.
@@ -1246,6 +1246,11 @@ def get_area_filtered_coco(self, min=0, max=float('inf')):
                 minimum allowed area
             max: int
                 maximum allowed area
+            intervals_per_category: dict of dicts
+                {
+                    "human": {"min": 20, "max": 10000},
+                    "vehicle": {"min": 50, "max": 15000},
+                }
         Returns:
             area_filtered_coco: sahi.utils.coco.Coco
         """
@@ -1259,6 +1264,13 @@ def get_area_filtered_coco(self, min=0, max=float('inf')):
         for image in self.images:
             is_valid_image = True
             for annotation in image.annotations:
+                if intervals_per_category is not None:
+                    category_based_min = intervals_per_category[annotation.category_name]["min"]
+                    category_based_max = intervals_per_category[annotation.category_name]["max"]
+                    if (annotation.area < category_based_min
+                        or
+                        annotation.area > category_based_max):
+                        is_valid_image = False
                 if annotation.area < min or annotation.area > max:
                     is_valid_image = False
             if is_valid_image: