improve .stats (#85)

* add num_negative_images to coco stats * add min_area_per_category and max_area_per_category to coco stats * fix min_annotation_area_per_category * fix typo * update stat docs
obss · May 8, 2021 · 02c50e9 · 02c50e9
1 parent f30efb3
commit 02c50e9
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 19 deletions.
diff --git a/docs/COCO.md b/docs/COCO.md
@@ -215,23 +215,20 @@ coco = Coco.from_coco_dict_or_path("coco.json")
 # get dataset stats
 coco.stats
 {
-    'avg_annotation_area': 2448.405738278109,
-    'avg_num_annotations_in_image': 53.037243084530985,
-    'max_annotation_area': 328640,
-    'max_num_annotations_in_image': 902,
-    'min_annotation_area': 3,
-    'min_num_annotations_in_image': 1,
-    'num_annotations': 343204,
-    'num_annotations_per_category': {
-        'human': 106396,
-        'vehicle': 236808
-    },
-    'num_categories': 2,
-    'num_images': 6471,
-    'num_images_per_category': {
-        'human': 5684,
-        'vehicle': 6323
-    }
+  'num_images': 6471,
+  'num_annotations': 343204,
+  'num_categories': 2,
+  'num_negative_images': 0,
+  'num_images_per_category': {'human': 5684, 'vehicle': 6323},
+  'num_annotations_per_category': {'human': 106396, 'vehicle': 236808},
+  'min_num_annotations_in_image': 1,
+  'max_num_annotations_in_image': 902,
+  'avg_num_annotations_in_image': 53.037243084530985,
+  'min_annotation_area': 3,
+  'max_annotation_area': 328640,
+  'avg_annotation_area': 2448.405738278109,
+  'min_annotation_area_per_category': {'human': 3, 'vehicle': 3},
+  'max_annotation_area_per_category': {'human': 72670, 'vehicle': 328640},
 }
 
 ```
diff --git a/sahi/utils/coco.py b/sahi/utils/coco.py
@@ -1016,12 +1016,17 @@ def calculate_stats(self):
         """
         Iterates over all annotations and calculates total number of
         """
+        # init all stats
         num_annotations = 0
         num_images = len(self.images)
+        num_negative_images = 0
         num_categories = len(self.json_categories)
         category_name_to_zero = {category["name"]:0 for category in self.json_categories}
+        category_name_to_inf = {category["name"]:float('inf') for category in self.json_categories}
         num_images_per_category = copy.deepcopy(category_name_to_zero)
         num_annotations_per_category = copy.deepcopy(category_name_to_zero)
+        min_annotation_area_per_category = copy.deepcopy(category_name_to_inf)
+        max_annotation_area_per_category = copy.deepcopy(category_name_to_zero)
         min_num_annotations_in_image = float('inf')
         max_num_annotations_in_image = 0
         total_annotation_area = 0
@@ -1039,6 +1044,13 @@ def calculate_stats(self):
                     max_annotation_area = annotation_area
                 if annotation_area<min_annotation_area:
                     min_annotation_area = annotation_area
+                if annotation_area>max_annotation_area_per_category[annotation.category_name]:
+                    max_annotation_area_per_category[annotation.category_name] = annotation_area
+                if annotation_area<min_annotation_area_per_category[annotation.category_name]:
+                    min_annotation_area_per_category[annotation.category_name] = annotation_area
+            # update num_negative_images
+            if len(image.annotations) == 0:
+                num_negative_images += 1
             # update num_annotations
             num_annotations += len(image.annotations)
             # update num_images_per_category
@@ -1051,21 +1063,24 @@ def calculate_stats(self):
                 max_num_annotations_in_image = num_annotations_in_image
             if num_annotations_in_image<min_num_annotations_in_image:
                 min_num_annotations_in_image = num_annotations_in_image
-        avg_num_annotations_in_image = num_annotations/num_images
+        avg_num_annotations_in_image = num_annotations/(num_images-num_negative_images)
         avg_annotation_area = total_annotation_area/num_annotations
 
         self._stats = {
             "num_images": num_images,
             "num_annotations": num_annotations,
             "num_categories": num_categories,
+            "num_negative_images": num_negative_images,
             "num_images_per_category": num_images_per_category,
             "num_annotations_per_category": num_annotations_per_category,
             "min_num_annotations_in_image": min_num_annotations_in_image,
             "max_num_annotations_in_image": max_num_annotations_in_image,
             "avg_num_annotations_in_image": avg_num_annotations_in_image,
             "min_annotation_area": min_annotation_area,
             "max_annotation_area": max_annotation_area,
-            "avg_annotation_area": avg_annotation_area
+            "avg_annotation_area": avg_annotation_area,
+            "min_annotation_area_per_category": min_annotation_area_per_category,
+            "max_annotation_area_per_category": max_annotation_area_per_category,
         }
 
     def split_coco_as_train_val(