Repository: fabiofumarola/polimorfo

Issue:
add method to remove categories

Discussion:

Patch:
diff --git a/polimorfo/datasets/coco.py b/polimorfo/datasets/coco.py
index 5dc46b1..b491fe0 100644
--- a/polimorfo/datasets/coco.py
+++ b/polimorfo/datasets/coco.py
@@ -200,6 +200,24 @@ class CocoDataset():
 
         self.reindex()
 
+    def remove_categories(self, idxs: List[int]) -> None:
+        """Remove the categories with the relative annotations
+
+        Args:
+            idxs (List[int]): [description]
+        """
+        for cat_idx in idxs:
+            if cat_idx not in self.cats:
+                continue
+
+            for idx, ann_meta in tqdm(self.anns.items(), 'process annotations'):
+                if ann_meta['category_id'] == cat_idx:
+                    del self.anns[idx]
+
+            del self.cats[cat_idx]
+
+        self.reindex()
+
     def remove_images_without_annotations(self):
         idx_images_with_annotations = {
             ann['image_id'] for ann in self.anns.values()


Test patch:
diff --git a/tests/datasets/test_coco.py b/tests/datasets/test_coco.py
index bc3c309..7b42d33 100644
--- a/tests/datasets/test_coco.py
+++ b/tests/datasets/test_coco.py
@@ -55,4 +55,14 @@ def test_create_dataset():
     img_id = ds.add_image((BASE_PATH / 'test_nodamage.jpg').as_posix())
     ds.add_annotation(img_id, cat_id, [1, 2, 3, 4, 5], 10, [0, 0, 256, 256], 0)
 
-    print(ds.dumps())
+    assert len(ds) == 1
+    assert len(ds.anns) == 1
+
+
+def test_remove_categories():
+    ds = CocoDataset()
+    cat_id = ds.add_category('dog', 'animal')
+    assert len(ds.cats) == 1
+    ds.remove_categories([cat_id])
+
+    assert len(ds.cats) == 0