Add utility to map COCO IDs to class names

Similar to the ImageNet utility.

Author: mattdangerw

keras_hub/api/utils/__init__.py

@@ -4,6 +4,10 @@
 since your modifications would be overwritten.
 """
 
+from keras_hub.src.utils.coco.coco_utils import coco_id_to_name
+from keras_hub.src.utils.coco.coco_utils import coco_name_to_id
 from keras_hub.src.utils.imagenet.imagenet_utils import (
     decode_imagenet_predictions,
 )
+from keras_hub.src.utils.imagenet.imagenet_utils import imagenet_id_to_name
+from keras_hub.src.utils.imagenet.imagenet_utils import imagenet_name_to_id

keras_hub/src/utils/coco/__init__.py

@@ -0,0 +1,137 @@
+from keras_hub.src.api_export import keras_hub_export
+
+
+@keras_hub_export("keras_hub.utils.coco_id_to_name")
+def coco_id_to_name(id):
+    """Convert a single COCO class name to a class ID.
+
+    Args:
+        id: An integer class id from 0 to 91.
+
+    Returns:
+        The human readable image class name, e.g. "bicycle".
+
+    Example:
+    ```python
+    >>> keras_hub.utils.coco_id_to_name(2)
+    "bicycle"
+    ```
+    """
+    return COCO_NAMES[id]
+
+
+@keras_hub_export("keras_hub.utils.coco_name_to_id")
+def coco_name_to_id(name):
+    """Convert a single COCO class name to a class ID.
+
+    Args:
+        name: A human readable image class name, e.g. "bicycle".
+
+    Returns:
+        The integer class id from 0 to 999.
+
+    Example:
+    ```python
+    >>> keras_hub.utils.coco_name_to_id("bicycle")
+    2
+    ```
+    """
+    return COCO_IDS[name]
+
+
+COCO_NAMES = {
+    0: "unlabeled",
+    1: "person",
+    2: "bicycle",
+    3: "car",
+    4: "motorcycle",
+    5: "airplane",
+    6: "bus",
+    7: "train",
+    8: "truck",
+    9: "boat",
+    10: "traffic_light",
+    11: "fire_hydrant",
+    12: "street_sign",
+    13: "stop_sign",
+    14: "parking_meter",
+    15: "bench",
+    16: "bird",
+    17: "cat",
+    18: "dog",
+    19: "horse",
+    20: "sheep",
+    21: "cow",
+    22: "elephant",
+    23: "bear",
+    24: "zebra",
+    25: "giraffe",
+    26: "hat",
+    27: "backpack",
+    28: "umbrella",
+    29: "shoe",
+    30: "eye_glasses",
+    31: "handbag",
+    32: "tie",
+    33: "suitcase",
+    34: "frisbee",
+    35: "skis",
+    36: "snowboard",
+    37: "sports_ball",
+    38: "kite",
+    39: "baseball_bat",
+    40: "baseball_glove",
+    41: "skateboard",
+    42: "surfboard",
+    43: "tennis_racket",
+    44: "bottle",
+    45: "plate",
+    46: "wine_glass",
+    47: "cup",
+    48: "fork",
+    49: "knife",
+    50: "spoon",
+    51: "bowl",
+    52: "banana",
+    53: "apple",
+    54: "sandwich",
+    55: "orange",
+    56: "broccoli",
+    57: "carrot",
+    58: "hot_dog",
+    59: "pizza",
+    60: "donut",
+    61: "cake",
+    62: "chair",
+    63: "couch",
+    64: "potted_plant",
+    65: "bed",
+    66: "mirror",
+    67: "dining_table",
+    68: "window",
+    69: "desk",
+    70: "toilet",
+    71: "door",
+    72: "tv",
+    73: "laptop",
+    74: "mouse",
+    75: "remote",
+    76: "keyboard",
+    77: "cell_phone",
+    78: "microwave",
+    79: "oven",
+    80: "toaster",
+    81: "sink",
+    82: "refrigerator",
+    83: "blender",
+    84: "book",
+    85: "clock",
+    86: "vase",
+    87: "scissors",
+    88: "teddy_bear",
+    89: "hair_drier",
+    90: "toothbrush",
+    91: "hair_brush",
+}
+
+COCO_IDS = {v: k for k, v in COCO_NAMES.items()}

keras_hub/src/utils/coco/coco_utils.py

@@ -0,0 +1,17 @@
+from keras_hub.src.tests.test_case import TestCase
+from keras_hub.src.utils.coco.coco_utils import coco_id_to_name
+from keras_hub.src.utils.coco.coco_utils import coco_name_to_id
+
+
+class CocoUtilsTest(TestCase):
+    def test_coco_id_to_name(self):
+        self.assertEqual(coco_id_to_name(0), "unlabeled")
+        self.assertEqual(coco_id_to_name(24), "zebra")
+        with self.assertRaises(KeyError):
+            coco_id_to_name(2001)
+
+    def test_coco_name_to_id(self):
+        self.assertEqual(coco_name_to_id("unlabeled"), 0)
+        self.assertEqual(coco_name_to_id("zebra"), 24)
+        with self.assertRaises(KeyError):
+            coco_name_to_id("whirligig")

keras_hub/src/utils/coco/coco_utils_test.py

@@ -3,6 +3,44 @@
 from keras_hub.src.api_export import keras_hub_export
 
 
+@keras_hub_export("keras_hub.utils.imagenet_id_to_name")
+def imagenet_id_to_name(id):
+    """Convert a single ImageNet class ID to a class name.
+
+    Args:
+        id: An integer class id from 0 to 999.
+
+    Returns:
+        The human readable image class name, e.g. "goldfish".
+
+    Example:
+    ```python
+    >>> keras_hub.utils.imagenet_id_to_name(1)
+    "goldfish"
+    ```
+    """
+    return IMAGENET_NAMES[id][1]
+
+
+@keras_hub_export("keras_hub.utils.imagenet_name_to_id")
+def imagenet_name_to_id(name):
+    """Convert a single ImageNet class name to a class ID.
+
+    Args:
+        name: A human readable image class name, e.g. "goldfish".
+
+    Returns:
+        The integer class id from 0 to 999.
+
+    Example:
+    ```python
+    >>> keras_hub.utils.imagenet_name_to_id("goldfish")
+    1
+    ```
+    """
+    return IMAGENET_IDS[name]
+
+
 @keras_hub_export("keras_hub.utils.decode_imagenet_predictions")
 def decode_imagenet_predictions(preds, top=5, include_synset_ids=False):
     """Decodes the predictions for an ImageNet-1k prediction.
@@ -1052,3 +1090,5 @@ def decode_imagenet_predictions(preds, top=5, include_synset_ids=False):
     998: ("n13133613", "ear"),
     999: ("n15075141", "toilet_tissue"),
 }
+
+IMAGENET_IDS = {v[1]: k for k, v in IMAGENET_NAMES.items()}

keras_hub/src/utils/imagenet/imagenet_utils.py

@@ -4,9 +4,23 @@
 from keras_hub.src.utils.imagenet.imagenet_utils import (
     decode_imagenet_predictions,
 )
+from keras_hub.src.utils.imagenet.imagenet_utils import imagenet_id_to_name
+from keras_hub.src.utils.imagenet.imagenet_utils import imagenet_name_to_id
 
 
 class ImageNetUtilsTest(TestCase):
+    def test_imagenet_id_to_name(self):
+        self.assertEqual(imagenet_id_to_name(0), "tench")
+        self.assertEqual(imagenet_id_to_name(21), "kite")
+        with self.assertRaises(KeyError):
+            imagenet_id_to_name(2001)
+
+    def test_imagenet_name_to_id(self):
+        self.assertEqual(imagenet_name_to_id("tench"), 0)
+        self.assertEqual(imagenet_name_to_id("kite"), 21)
+        with self.assertRaises(KeyError):
+            imagenet_name_to_id(2001)
+
     def test_decode_imagenet_predictions(self):
         preds = np.array(
             [