Support custom classes (#2443)

* Add regression test for custom classes * remove assert
WonderNoot · Apr 14, 2020 · 4ece0d8 · 4ece0d8
1 parent c64a80d
commit 4ece0d8
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/mmdet/datasets/custom.py b/mmdet/datasets/custom.py
@@ -184,7 +184,6 @@ def get_classes(cls, classes=None):
             raise ValueError('Unsupported type {} of classes.'.format(
                 type(classes)))
 
-        assert set(class_names).issubset(set(cls.CLASSES))
         return class_names
 
     def get_subset_by_classes(self):

diff --git a/tests/test_dataset.py b/tests/test_dataset.py
@@ -40,6 +40,18 @@ def test_custom_classes_override_default(dataset):
     assert custom_dataset.CLASSES == ['bus', 'car']
     assert custom_dataset.custom_classes
 
+    # Test overriding not a subset
+    custom_dataset = dataset_class(
+        ann_file=MagicMock(),
+        pipeline=[],
+        classes=['foo'],
+        test_mode=True,
+        img_prefix='VOC2007' if dataset == 'VOCDataset' else '')
+
+    assert custom_dataset.CLASSES != original_classes
+    assert custom_dataset.CLASSES == ['foo']
+    assert custom_dataset.custom_classes
+
     # Test default behavior
     custom_dataset = dataset_class(
         ann_file=MagicMock(),