Allow for bboxes in other formats.

PiperOrigin-RevId: 686530258

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/features/bounding_boxes.py

@@ -37,6 +37,17 @@
 class BBoxFeature(tensor_feature.Tensor):
   """`FeatureConnector` for a normalized bounding box.
 
+  By default, TFDS uses normalized YXYX bbox format. This can be changed by
+  passing the `bbox_format` argument, e.g.
+  ```
+    features=features.FeatureDict({
+        'bbox': tfds.features.BBox(bbox_format=bb_utils.BBoxFormat.XYWH),
+    })
+  ```
+  If you don't know the format of the bbox, you can use `bbox_format=None`. In
+  this case, the only check that is done is that 4 floats coordinates are
+  provided.
+
   Note: If you have multiple bounding boxes, you may want to wrap the feature
   inside a `tfds.features.Sequence`.
 
@@ -69,12 +80,17 @@ def __init__(
       self,
       *,
       doc: feature_lib.DocArg = None,
+      bbox_format: (
+          bb_utils.BBoxFormatType | None
+      ) = bb_utils.BBoxFormat.REL_YXYX,
   ):
-    super(BBoxFeature, self).__init__(shape=(4,), dtype=np.float32, doc=doc)
+    if isinstance(bbox_format, str):
+      bbox_format = bb_utils.BBoxFormat(bbox_format)
+    self.bbox_format = bbox_format
+    super().__init__(shape=(4,), dtype=np.float32, doc=doc)
 
   def encode_example(self, bbox: Union[bb_utils.BBox, np.ndarray]):
     """See base class for details."""
-
     if isinstance(bbox, np.ndarray):
       if bbox.shape != (4,):
         raise ValueError(
@@ -88,14 +104,22 @@ def encode_example(self, bbox: Union[bb_utils.BBox, np.ndarray]):
 
     # Validate the coordinates
     for coordinate in bbox:
-      if not isinstance(coordinate, (float, np.floating)):
-        raise ValueError(
-            'BBox coordinates should be float. Got {}.'.format(bbox)
-        )
-      if not 0.0 <= coordinate <= 1.0:
-        raise ValueError(
-            'BBox coordinates should be between 0 and 1. Got {}.'.format(bbox)
-        )
+      if (
+          self.bbox_format == bb_utils.BBoxFormat.REL_YXYX
+          or self.bbox_format == bb_utils.BBoxFormat.REL_XYXY
+      ):
+        if not isinstance(coordinate, (float, np.floating)):
+          raise ValueError(
+              'BBox coordinates should be float. Got {}.'.format(bbox)
+          )
+        if not 0.0 <= coordinate <= 1.0:
+          raise ValueError(
+              'BBox coordinates should be between 0 and 1. Got {}.'.format(bbox)
+          )
+    if (
+        self.bbox_format == bb_utils.BBoxFormat.YXYX
+        or self.bbox_format == bb_utils.BBoxFormat.REL_YXYX
+    ):
       if bbox.xmax < bbox.xmin or bbox.ymax < bbox.ymin:
         raise ValueError(
             'BBox coordinates should have min <= max. Got {}.'.format(bbox)
@@ -108,28 +132,51 @@ def encode_example(self, bbox: Union[bb_utils.BBox, np.ndarray]):
   def repr_html(self, ex: np.ndarray) -> str:
     """Returns the HTML str representation of an Image with BBoxes."""
     ex = np.expand_dims(ex, axis=0)  # Expand single bounding box to batch.
-    return _repr_html(ex)
+    return _repr_html(ex, bbox_format=self.bbox_format)
 
   def repr_html_batch(self, ex: np.ndarray) -> str:
     """Returns the HTML str representation of an Image with BBoxes (Sequence)."""
-    return _repr_html(ex)
+    return _repr_html(ex, bbox_format=self.bbox_format)
 
   @classmethod
   def from_json_content(
       cls, value: Union[Json, feature_pb2.BoundingBoxFeature]
   ) -> 'BBoxFeature':
-    del value  # Unused
-    return cls()
+    if isinstance(value, dict):
+      return cls(**value)
+    return cls(
+        bbox_format=bb_utils.BBoxFormat(value.bbox_format)
+        if value.bbox_format
+        else None
+    )
 
-  def to_json_content(self) -> feature_pb2.BoundingBoxFeature:  # pytype: disable=signature-mismatch  # overriding-return-type-checks
+  def to_json_content(
+      self,
+  ) -> (
+      feature_pb2.BoundingBoxFeature
+  ):  # pytype: disable=signature-mismatch  # overriding-return-type-checks
+    bbox_format = None
+    if self.bbox_format:
+      bbox_format = (
+          self.bbox_format
+          if isinstance(self.bbox_format, str)
+          else self.bbox_format.value
+      )
     return feature_pb2.BoundingBoxFeature(
         shape=feature_lib.to_shape_proto(self._shape),
         dtype=feature_lib.dtype_to_str(self._dtype),
+        bbox_format=bbox_format,
     )
 
 
-def _repr_html(ex: np.ndarray) -> str:
+def _repr_html(
+    ex: np.ndarray, bbox_format: bb_utils.BBoxFormatType | None
+) -> str:
   """Returns the HTML str representation of an Image with BBoxes."""
+  # If the bbox format is not normalized, we don't draw the bbox on a blank
+  # image but we return a string representation of the bbox instead.
+  if bbox_format != bb_utils.BBoxFormat.REL_YXYX:
+    return repr(ex)
   img = _build_thumbnail_with_bbox(ex)
   img_str = utils.get_base64(lambda buff: img.save(buff, format='PNG'))
   return f'<img src="data:image/png;base64,{img_str}" alt="Img" />'

tensorflow_datasets/core/features/bounding_boxes_test.py

@@ -51,6 +51,24 @@ def test_feature(self):
         ],
     )
 
+  def test_unspecified_bbox_format_feature(self):
+    self.assertFeature(
+        feature=features.BBoxFeature(bbox_format=None),
+        shape=(4,),
+        dtype=np.float32,
+        tests=[
+            testing.FeatureExpectationItem(
+                # 1D numpy array, float32 unnormalized bbox
+                value=np.array(
+                    [200.46, 199.84, 77.71, 70.88], dtype=np.float32
+                ),
+                expected=np.array(
+                    [200.46, 199.84, 77.71, 70.88], dtype=np.float32
+                ),
+            ),
+        ],
+    )
+
 
 if __name__ == '__main__':
   testing.test_main()

tensorflow_datasets/core/features/bounding_boxes_utils.py

@@ -46,6 +46,8 @@ class BBoxFormat(enum.Enum):
   XYXY = 'XYXY'
   YXYX = 'YXYX'
   XYWH = 'XYWH'
+  REL_XYXY = 'REL_XYXY'
+  REL_YXYX = 'REL_YXYX'
 
 
 BBoxFormatType = Union[BBoxFormat, str]
@@ -77,7 +79,6 @@ def convert_coordinates_to_bbox(
   """
   if len(coordinates) != 4:
     raise ValueError(f'Expected 4 coordinates, got {coordinates}.')
-
   coordinates = coordinates.astype(np.float64)
 
   try:
@@ -88,15 +89,24 @@ def convert_coordinates_to_bbox(
         f'Unsupported bbox format: {format}. Currently supported bounding box'
         f' formats are: {[format.value for format in BBoxFormat]}'
     ) from e
-
-  if input_format == BBoxFormat.YXYX:
+  if (
+      input_format == BBoxFormat.REL_XYXY or input_format == BBoxFormat.REL_YXYX
+  ) and normalize:
+    raise ValueError(
+        'If the input format is normalized, then normalize should be False.'
+    )
+  if normalize and img_shape is None:
+    raise ValueError(
+        'If normalize is True, img_shape must be provided, but got None.'
+    )
+  if input_format == BBoxFormat.YXYX or input_format == BBoxFormat.REL_YXYX:
     bbox = BBox(
         ymin=coordinates[0],
         xmin=coordinates[1],
         ymax=coordinates[2],
         xmax=coordinates[3],
     )
-  elif input_format == BBoxFormat.XYXY:
+  elif input_format == BBoxFormat.XYXY or input_format == BBoxFormat.REL_XYXY:
     bbox = BBox(
         ymin=coordinates[1],
         xmin=coordinates[0],

tensorflow_datasets/core/features/bounding_boxes_utils_test.py

@@ -67,6 +67,19 @@ def test_convert_to_bbox(
     [
         ([1, 2, 3], None, 'Expected 4'),
         (TEST_INPUT_LIST, 'NonExistentConverter', 'Unsupported bbox format'),
+        (
+            TEST_INPUT_LIST,
+            bb_utils.BBoxFormat.XYXY,
+            'If normalize is True, img_shape must be provided, but got None.',
+        ),
+        (
+            TEST_INPUT_LIST,
+            bb_utils.BBoxFormat.REL_XYXY,
+            (
+                'If the input format is normalized, then normalize should be'
+                ' False.'
+            ),
+        ),
     ],
 )
 def test_convert_coordinates_to_bbox_valueerror(

tensorflow_datasets/core/proto/feature.proto

@@ -105,10 +105,12 @@ message AudioFeature {
 
 // A bounding box around an object in an image. Typically, bounding boxes are
 // tensors of type `tf.float32` and shape `[4,]` and contain the normalized
-// coordinates of the bounding box `[ymin, xmin, ymax, xmax]`.
+// coordinates of the bounding box `[ymin, xmin, ymax, xmax]`. Different formats
+// are supported through the `bbox_format` field.
 message BoundingBoxFeature {
   Shape shape = 1;
   string dtype = 2;
+  string bbox_format = 3;
 }
 
 message TextFeature {}

tensorflow_datasets/core/proto/feature_generated_pb2.py

@@ -72,8 +72,9 @@
     b' \x01(\t\x12\x13\n\x0bsample_rate\x18\x04'
     b' \x01(\x03\x12\x10\n\x08\x65ncoding\x18\x05'
     b' \x01(\t\x12\x13\n\x0blazy_decode\x18\x06'
-    b' \x01(\x08"N\n\x12\x42oundingBoxFeature\x12)\n\x05shape\x18\x01'
+    b' \x01(\x08"c\n\x12\x42oundingBoxFeature\x12)\n\x05shape\x18\x01'
     b' \x01(\x0b\x32\x1a.tensorflow_datasets.Shape\x12\r\n\x05\x64type\x18\x02'
+    b' \x01(\t\x12\x13\n\x0b\x62\x62ox_format\x18\x03'
     b' \x01(\t"\r\n\x0bTextFeature"O\n\x12TranslationFeature\x12\x11\n\tlanguages\x18\x01'
     b' \x03(\t\x12&\n\x1evariable_languages_per_example\x18\x02'
     b' \x01(\x08"I\n\x08Sequence\x12-\n\x07\x66\x65\x61ture\x18\x01'
@@ -110,11 +111,11 @@
   _AUDIOFEATURE._serialized_start = 1500
   _AUDIOFEATURE._serialized_end = 1653
   _BOUNDINGBOXFEATURE._serialized_start = 1655
-  _BOUNDINGBOXFEATURE._serialized_end = 1733
-  _TEXTFEATURE._serialized_start = 1735
-  _TEXTFEATURE._serialized_end = 1748
-  _TRANSLATIONFEATURE._serialized_start = 1750
-  _TRANSLATIONFEATURE._serialized_end = 1829
-  _SEQUENCE._serialized_start = 1831
-  _SEQUENCE._serialized_end = 1904
+  _BOUNDINGBOXFEATURE._serialized_end = 1754
+  _TEXTFEATURE._serialized_start = 1756
+  _TEXTFEATURE._serialized_end = 1769
+  _TRANSLATIONFEATURE._serialized_start = 1771
+  _TRANSLATIONFEATURE._serialized_end = 1850
+  _SEQUENCE._serialized_start = 1852
+  _SEQUENCE._serialized_end = 1925
 # @@protoc_insertion_point(module_scope)