No public description

PiperOrigin-RevId: 764863901

Author: tensorflower-gardener

official/projects/waste_identification_ml/Triton_TF_Cloud_Deployment/client/utils.py

@@ -16,6 +16,7 @@
 
 from collections.abc import Mapping, Sequence
 import csv
+import dataclasses
 import logging
 import os
 from typing import Any, TypedDict
@@ -31,6 +32,20 @@ class ItemDict(TypedDict):
   supercategory: str
 
 
+@dataclasses.dataclass
+class BoundingBox:
+  y1: int | float
+  x1: int | float
+  y2: int | float
+  x2: int | float
+
+
+@dataclasses.dataclass
+class ImageSize:
+  height: int
+  width: int
+
+
 def _reframe_image_corners_relative_to_boxes(boxes: tf.Tensor) -> tf.Tensor:
   """Reframe the image corners ([0, 0, 1, 1]) to be relative to boxes.
 
@@ -455,3 +470,27 @@ def filter_detections(
   filtered_output['num_detections'] = np.array([new_num_detections])
 
   return filtered_output
+
+
+def resize_bbox(
+    bbox: BoundingBox, old_size: ImageSize, new_size: ImageSize
+) -> tuple[int, int, int, int]:
+  """Resize bounding box coordinates based on new image size.
+
+  Args:
+      bbox: BoundingBox with original coordinates.
+      old_size: Original image size.
+      new_size: New image size.
+
+  Returns:
+      Rescaled bounding box coordinates.
+  """
+  scale_x = new_size.width / old_size.width
+  scale_y = new_size.height / old_size.height
+
+  new_y1 = int(bbox.y1 * scale_y)
+  new_x1 = int(bbox.x1 * scale_x)
+  new_y2 = int(bbox.y2 * scale_y)
+  new_x2 = int(bbox.x2 * scale_x)
+
+  return new_y1, new_x1, new_y2, new_x2

official/projects/waste_identification_ml/Triton_TF_Cloud_Deployment/client/utils_test.py

@@ -17,6 +17,8 @@
 import unittest
 import numpy as np
 from official.projects.waste_identification_ml.Triton_TF_Cloud_Deployment.client import utils
+from official.projects.waste_identification_ml.Triton_TF_Cloud_Deployment.client.utils import BoundingBox
+from official.projects.waste_identification_ml.Triton_TF_Cloud_Deployment.client.utils import ImageSize
 
 
 class TestLoadLabels(unittest.TestCase):
@@ -222,6 +224,26 @@ def test_extract_and_resize_single_object(self):
     # The output pixels in mask area should be non-zero
     self.assertTrue(np.any(obj > 0))
 
+  def test_resize_bbox_scaling(self):
+    bbox = BoundingBox(y1=50, x1=100, y2=150, x2=200)
+    old_size = ImageSize(height=200, width=400)
+    new_size = ImageSize(height=400, width=800)
+
+    expected = (100, 200, 300, 400)
+    result = utils.resize_bbox(bbox, old_size, new_size)
+
+    self.assertEqual(result, expected)
+
+  def test_resize_bbox_no_scaling(self):
+    bbox = BoundingBox(y1=10, x1=20, y2=30, x2=40)
+    old_size = ImageSize(height=100, width=100)
+    new_size = ImageSize(height=100, width=100)
+
+    expected = (10, 20, 30, 40)
+    result = utils.resize_bbox(bbox, old_size, new_size)
+
+    self.assertEqual(result, expected)
+
 
 if __name__ == '__main__':
   unittest.main()