Chongyangbai/multilabel2kvp

README.md

@@ -96,10 +96,13 @@ print(target)
 You can convert an existing IC/OD VisionDataset to the generalized KVP format using the following adapter:
 
 ```{python}
-# For IC dataset
-from vision_datasets.image_classification import ClassificationAsKeyValuePairDataset
-sample_ic_dataset = VisionDataset(dataset_info, dataset_manifest)
-kvp_dataset = ClassificationAsKeyValuePairDataset(sample_ic_dataset)
+# For MultiClass and MultiLabel IC dataset
+from vision_datasets.image_classification import MultiClassAsKeyValuePairDataset, MultiLabelAsKeyValuePairDataset
+sample_multiclass_ic_dataset = VisionDataset(dataset_info, dataset_manifest)
+kvp_dataset = MultiClassAsKeyValuePairDataset(sample_multiclass_ic_dataset)
+sample_multilabel_ic_dataset = VisionDataset(dataset_info, dataset_manifest)
+kvp_dataset = MultiLabelAsKeyValuePairDataset(sample_multilabel_ic_dataset)
+
 
 # For OD dataset
 from vision_datasets.image_object_detection import DetectionAsKeyValuePairDataset

setup.py

@@ -1,7 +1,7 @@
 import setuptools
 from os import path
 
-VERSION = '1.0.15'
+VERSION = '1.0.16'
 
 # Get the long description from the README file
 here = path.abspath(path.dirname(__file__))

tests/test_fixtures.py

@@ -58,7 +58,7 @@ def create_an_od_dataset(n_images=2, n_categories=4, coordinates='relative'):
         return dataset, tempdir
 
 
-class MultilcassClassificationTestFixtures:
+class MulticlassClassificationTestFixtures:
     DATASET_INFO_DICT = {
         "name": "dummy",
         "version": 1,
@@ -75,14 +75,14 @@ class MultilcassClassificationTestFixtures:
 
     @staticmethod
     def create_an_ic_dataset(n_images=2, n_categories=3):
-        dataset_dict = copy.deepcopy(MultilcassClassificationTestFixtures.DATASET_INFO_DICT)
+        dataset_dict = copy.deepcopy(MulticlassClassificationTestFixtures.DATASET_INFO_DICT)
         tempdir = tempfile.TemporaryDirectory()
         dataset_dict['root_folder'] = tempdir.name
         for i in range(n_images):
             Image.new('RGB', (100, 100)).save(pathlib.Path(tempdir.name) / f'{i + 1}.jpg')
 
         dataset_info = DatasetInfo(dataset_dict)
-        dataset_manifest = MultilcassClassificationTestFixtures.create_an_ic_manifest(tempdir.name, n_images, n_categories)
+        dataset_manifest = MulticlassClassificationTestFixtures.create_an_ic_manifest(tempdir.name, n_images, n_categories)
         dataset = VisionDataset(dataset_info, dataset_manifest)
         return dataset, tempdir
 
@@ -98,3 +98,44 @@ def create_an_ic_manifest(root_dir='', n_images=2, n_categories=3):
         coco_path = pathlib.Path(root_dir) / 'coco.json'
         coco_path.write_text(json.dumps(coco_dict))
         return CocoManifestAdaptorFactory.create(DatasetTypes.IMAGE_CLASSIFICATION_MULTICLASS).create_dataset_manifest(coco_path.name, root_dir)
+
+
+class MultilabelClassificationTestFixtures:
+    DATASET_INFO_DICT = {
+        "name": "dummy",
+        "version": 1,
+        "type": "image_classification_multilabel",
+        "root_folder": "dummy",
+        "format": "coco",
+        "test": {
+            "index_path": "train.json",
+            "files_for_local_usage": [
+                "train.zip"
+            ]
+        },
+    }
+
+    @staticmethod
+    def create_an_ic_dataset(n_images=2, n_categories=3):
+        dataset_dict = copy.deepcopy(MultilabelClassificationTestFixtures.DATASET_INFO_DICT)
+        tempdir = tempfile.TemporaryDirectory()
+        dataset_dict['root_folder'] = tempdir.name
+        for i in range(n_images):
+            Image.new('RGB', (100, 100)).save(pathlib.Path(tempdir.name) / f'{i + 1}.jpg')
+
+        dataset_info = DatasetInfo(dataset_dict)
+        dataset_manifest = MultilabelClassificationTestFixtures.create_an_ic_manifest(tempdir.name, n_images, n_categories)
+        dataset = VisionDataset(dataset_info, dataset_manifest)
+        return dataset, tempdir
+
+    @staticmethod
+    def create_an_ic_manifest(root_dir='', n_images=2, n_categories=3):
+        images = [{'id': i + 1, 'file_name': f'{i + 1}.jpg', 'width': 100, 'height': 100} for i in range(n_images)]
+
+        categories = [{'id': i + 1, 'name': f'{i + 1}-class', } for i in range(n_categories)]
+        annotations = [{'id': i + 1, 'image_id': i + 1, 'category_id': i + 1} for i in range(n_images)]
+        annotations.extend([{'id': n_images + i + 1, 'image_id': i + 1, 'category_id': n_images - i} for i in range(n_images)])
+        coco_dict = {'images': images, 'categories': categories, 'annotations': annotations}
+        coco_path = pathlib.Path(root_dir) / 'coco.json'
+        coco_path.write_text(json.dumps(coco_dict))
+        return CocoManifestAdaptorFactory.create(DatasetTypes.IMAGE_CLASSIFICATION_MULTILABEL).create_dataset_manifest(coco_path.name, root_dir)

tests/test_ic_od_to_kvp_wrapper/test_classification_as_kvp.py

@@ -1,24 +1,22 @@
 import unittest
 
-from tests.test_fixtures import MultilcassClassificationTestFixtures
+from tests.test_fixtures import MulticlassClassificationTestFixtures, MultilabelClassificationTestFixtures
 from vision_datasets.common import DatasetTypes
-from vision_datasets.image_classification import ClassificationAsKeyValuePairDataset
+from vision_datasets.image_classification import MultiClassAsKeyValuePairDataset, MultiLabelAsKeyValuePairDataset
 from vision_datasets.key_value_pair.manifest import KeyValuePairLabelManifest
 
 
 class TestClassificationAsKeyValuePairDataset(unittest.TestCase):
     def test_multiclass_classification(self):
-        sample_classification_dataset, _ = MultilcassClassificationTestFixtures.create_an_ic_dataset()
-        kvp_dataset = ClassificationAsKeyValuePairDataset(sample_classification_dataset)
+        sample_classification_dataset, _ = MulticlassClassificationTestFixtures.create_an_ic_dataset()
+        kvp_dataset = MultiClassAsKeyValuePairDataset(sample_classification_dataset)
 
-        self.assertIsInstance(kvp_dataset, ClassificationAsKeyValuePairDataset)
+        self.assertIsInstance(kvp_dataset, MultiClassAsKeyValuePairDataset)
         self.assertEqual(kvp_dataset.dataset_info.type, DatasetTypes.KEY_VALUE_PAIR)
         self.assertIn("name", kvp_dataset.dataset_info.schema)
         self.assertIn("description", kvp_dataset.dataset_info.schema)
         self.assertIn("fieldSchema", kvp_dataset.dataset_info.schema)
 
-        print(kvp_dataset.dataset_info.schema["fieldSchema"])
-
         self.assertEqual(kvp_dataset.dataset_info.schema["fieldSchema"],
                          {"className": {
                              "type": "string",
@@ -37,6 +35,38 @@ def test_multiclass_classification(self):
                          {"fields": {"className": {"value": "1-class"}}}
                          )
 
+    def test_multilabel_classification(self):
+        sample_classification_dataset, _ = MultilabelClassificationTestFixtures.create_an_ic_dataset(n_images=2, n_categories=2)
+        kvp_dataset = MultiLabelAsKeyValuePairDataset(sample_classification_dataset)
+
+        self.assertIsInstance(kvp_dataset, MultiLabelAsKeyValuePairDataset)
+        self.assertEqual(kvp_dataset.dataset_info.type, DatasetTypes.KEY_VALUE_PAIR)
+        self.assertIn("name", kvp_dataset.dataset_info.schema)
+        self.assertIn("description", kvp_dataset.dataset_info.schema)
+        self.assertIn("fieldSchema", kvp_dataset.dataset_info.schema)
+
+        self.assertEqual(kvp_dataset.dataset_info.schema["fieldSchema"],
+                         {'classNames': {
+                             'type': 'array',
+                             'items': {
+                                 'type': 'string',
+                                 'description': 'Class names that the image belongs to.',
+                                 'classes': {
+                                     '1-class': {'description': 'A single class name. Only output 1-class as the class name if present.'},
+                                     '2-class': {'description': 'A single class name. Only output 2-class as the class name if present.'}
+                                  }
+                                 }
+                             }
+                          }
+                         )
+
+        _, target, _ = kvp_dataset[0]
+        self.assertIsInstance(target, KeyValuePairLabelManifest)
+        self.assertEqual(target.label_data,
+                         {'fields': {
+                             'classNames': {'value': ['1-class', '2-class']}}
+                          })
+
 
 if __name__ == '__main__':
     unittest.main()

vision_datasets/image_classification/__init__.py

@@ -1,9 +1,9 @@
 from .coco_manifest_adaptor import MultiClassClassificationCocoManifestAdaptor, MultiLabelClassificationCocoManifestAdaptor
 from .operations import ImageClassificationCocoDictGenerator
 from .manifest import ImageClassificationLabelManifest
-from .classification_as_kvp_dataset import ClassificationAsKeyValuePairDataset
+from .classification_as_kvp_dataset import MultiClassAsKeyValuePairDataset, MultiLabelAsKeyValuePairDataset
 
 __all__ = ['MultiClassClassificationCocoManifestAdaptor', 'MultiLabelClassificationCocoManifestAdaptor',
            'ImageClassificationCocoDictGenerator',
            'ImageClassificationLabelManifest',
-           'ClassificationAsKeyValuePairDataset']
+           'MultiClassAsKeyValuePairDataset', 'MultiLabelAsKeyValuePairDataset']

vision_datasets/image_classification/classification_as_kvp_dataset.py

@@ -11,21 +11,7 @@
 logger = logging.getLogger(__name__)
 
 
-CLASS_NAME_KEY = "className"
-BASE_CLASSIFICATION_SCHEMA = {
-    "name": "Multiclass image classification",
-    "description": "Classify images into one of the provided classes.",
-    "fieldSchema": {
-            f"{CLASS_NAME_KEY}": {
-                "type": "string",
-                "description": "Class name that the image belongs to.",
-                "classes": {}
-            }
-    }
-}
-
-
-class ClassificationAsKeyValuePairDataset(VisionDataset):
+class ClassificationAsKeyValuePairDatasetBase(VisionDataset):
     """Dataset class that access Classification datset as KeyValuePair dataset."""
 
     def __init__(self, classification_dataset: VisionDataset):
@@ -35,8 +21,7 @@ def __init__(self, classification_dataset: VisionDataset):
             classification_dataset (VisionDataset): The classification dataset to convert to key-value pair dataset.
         """
 
-        if classification_dataset is None or classification_dataset.dataset_info.type not in {DatasetTypes.IMAGE_CLASSIFICATION_MULTICLASS}:
-            # TODO: Add support for multilabel classification
+        if classification_dataset is None or classification_dataset.dataset_info.type not in {DatasetTypes.IMAGE_CLASSIFICATION_MULTICLASS, DatasetTypes.IMAGE_CLASSIFICATION_MULTILABEL}:
             raise ValueError
 
         # Generate schema and update dataset info
@@ -55,10 +40,10 @@ def __init__(self, classification_dataset: VisionDataset):
         # Construct KeyValuePairDatasetManifest
         annotations = []
         for id, img in enumerate(classification_dataset.dataset_manifest.images, 1):
-            label_id = img.labels[0].label_data
-            label_name = self.class_id_to_names[label_id]
+            label_ids = [l.label_data for l in img.labels]
+            label_names = [self.class_id_to_names[id] for id in label_ids]
 
-            kvp_label_data = self.construct_kvp_label_data(label_name)
+            kvp_label_data = self.construct_kvp_label_data(label_names)
             img_ids = [self.img_id_to_pos[img.id]]  # 0-based index
             kvp_annotation = KeyValuePairLabelManifest(id, img_ids, label_data=kvp_label_data)
 
@@ -70,20 +55,79 @@ def __init__(self, classification_dataset: VisionDataset):
         super().__init__(dataset_info, dataset_manifest, dataset_resources=classification_dataset.dataset_resources)
 
     def construct_schema(self, class_names: typing.List[str]) -> typing.Dict[str, typing.Any]:
-        schema: typing.Dict[str, typing.Any] = BASE_CLASSIFICATION_SCHEMA  # initialize with base schema
-        schema["fieldSchema"][f"{CLASS_NAME_KEY}"]["classes"] = {c: {"description": f"A single class name. Only output {c} as the class name if present."} for c in class_names}
+        raise NotImplementedError
+
+    def construct_kvp_label_data(self, label_names: typing.List[str]) -> typing.Dict[str, typing.Union[typing.Dict[str, typing.Dict[str, str]], None]]:
+        raise NotImplementedError
+
+
+class MultiClassAsKeyValuePairDataset(ClassificationAsKeyValuePairDatasetBase):
+    CLASS_NAME_KEY = "className"
+    SCHEMA_BASE = {
+        "name": "Multiclass image classification",
+        "description": "Classify images into one of the provided classes.",
+        "fieldSchema": {
+            f"{CLASS_NAME_KEY}": {
+                "type": "string",
+                "description": "Class name that the image belongs to.",
+                "classes": {}
+            }
+        }
+    }
+
+    def construct_schema(self, class_names: typing.List[str]) -> typing.Dict[str, typing.Any]:
+        schema: typing.Dict[str, typing.Any] = MultiClassAsKeyValuePairDataset.SCHEMA_BASE  # initialize with base schema
+        schema["fieldSchema"][f"{MultiClassAsKeyValuePairDataset.CLASS_NAME_KEY}"]["classes"] = {c: {"description": f"A single class name. Only output {c} as the class name if present."} for c in class_names}
         return schema
 
-    def construct_kvp_label_data(self, label_name: str) -> typing.Dict[str, typing.Union[typing.Dict[str, typing.Dict[str, str]], None]]:
+    def construct_kvp_label_data(self, label_names: typing.List[str]) -> typing.Dict[str, typing.Union[typing.Dict[str, typing.Dict[str, str]], None]]:
         """
-        Convert the classification dataset label_name to the desired format for KVP annnotation as defined by the BASE_CLASSIFICATION_SCHEMA.
+        Convert the classification dataset label_name to the desired format for KVP annnotation as defined by the SCHEMA_BASE.
         E.g. {"fields": {"className": {"value": <label_name>}}}
 
         """
         return {
             f"{KeyValuePairLabelManifest.LABEL_KEY}": {
-                f"{CLASS_NAME_KEY}": {
-                    f"{KeyValuePairLabelManifest.LABEL_VALUE_KEY}": label_name
+                f"{MultiClassAsKeyValuePairDataset.CLASS_NAME_KEY}": {
+                    f"{KeyValuePairLabelManifest.LABEL_VALUE_KEY}": label_names[0]
                 }
             }
         }
+
+
+
+class MultiLabelAsKeyValuePairDataset(ClassificationAsKeyValuePairDatasetBase):
+    CLASS_NAME_KEY = "classNames"
+    SCHEMA_BASE = {
+        "name": "Multilabel image classification",
+        "description": "Classify images into one or more of the provided classes.",
+        "fieldSchema": {
+            f"{CLASS_NAME_KEY}": {
+                "type": "array",
+                "items": {
+                    "type": "string",
+                    "description": "Class names that the image belongs to.",
+                    "classes": {}
+                }
+            }
+        }
+    }
+
+    def construct_schema(self, class_names: typing.List[str]) -> typing.Dict[str, typing.Any]:
+        schema: typing.Dict[str, typing.Any] = MultiLabelAsKeyValuePairDataset.SCHEMA_BASE  # initialize with base schema
+        schema["fieldSchema"][f"{MultiLabelAsKeyValuePairDataset.CLASS_NAME_KEY}"]['items']["classes"] = {c: {"description": f"A single class name. Only output {c} as the class name if present."} for c in class_names}
+        return schema
+
+    def construct_kvp_label_data(self, label_names: typing.List[str]) -> typing.Dict[str, typing.Union[typing.Dict[str, typing.Dict[str, str]], None]]:
+        """
+        Convert the classification dataset label_name to the desired format for KVP annnotation as defined by the SCHEMA_BASE.
+        E.g. {"fields": {"className": {"value": <label_name>}}}
+
+        """
+        return {
+            f"{KeyValuePairLabelManifest.LABEL_KEY}": {
+                f"{MultiLabelAsKeyValuePairDataset.CLASS_NAME_KEY}": {
+                    f"{KeyValuePairLabelManifest.LABEL_VALUE_KEY}": label_names
+                }
+            }
+        }