No public description

PiperOrigin-RevId: 794717467

Author: tensorflower-gardener

official/projects/waste_identification_ml/llm_applications/milk_pouch_detection/models_utils.py

@@ -0,0 +1,184 @@
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utility functions for milk pouch detection."""
+
+from collections.abc import Mapping
+import dataclasses
+from typing import Any
+
+import numpy as np
+import torch
+import torchvision
+
+
+@dataclasses.dataclass(frozen=True)
+class _BoundingBox:
+  """A class representing a bounding box."""
+  x1: float
+  y1: float
+  x2: float
+  y2: float
+
+
+def _box_area(box: _BoundingBox) -> float:
+  """Calculates the area of a bounding box.
+
+  Args:
+    box: A _BoundingBox object.
+
+  Returns:
+    The area of the bounding box.
+  """
+  return max(0, box.x2 - box.x1) * max(0, box.y2 - box.y1)
+
+
+def _calculate_iou(
+    box1: _BoundingBox,
+    box2: _BoundingBox
+) -> float:
+  """Calculates the Intersection over Union (IoU) of two bounding boxes.
+
+  Args:
+    box1: The first bounding box in (x1, y1, x2, y2) format.
+    box2: The second bounding box in (x1, y1, x2, y2) format.
+
+  Returns:
+    The IoU score, a float between 0.0 and 1.0.
+  """
+  # Determine the coordinates of the intersection rectangle
+  x1 = max(box1.x1, box2.x1)
+  y1 = max(box1.y1, box2.y1)
+  x2 = min(box1.x2, box2.x2)
+  y2 = min(box1.y2, box2.y2)
+
+  # Calculate the area of intersection
+  inter_area = max(0, x2 - x1) * max(0, y2 - y1)
+
+  # Calculate the area of both bounding boxes
+  box1_area = _box_area(box1)
+  box2_area = _box_area(box2)
+
+  # Calculate the area of the union
+  union_area = box1_area + box2_area - inter_area
+
+  # Compute the IoU score
+  return inter_area / union_area if union_area != 0 else 0.0
+
+
+def _is_contained(
+    inner_box: _BoundingBox,
+    outer_box: _BoundingBox,
+    margin: int = 5,
+) -> bool:
+  """Checks if one bounding box is contained within another, with a margin.
+
+  Args:
+    inner_box: The bounding box that is potentially inside.
+    outer_box: The bounding box that is potentially surrounding.
+    margin: An optional pixel margin to allow for slight inaccuracies.
+
+  Returns:
+    True if the inner box is contained within the outer box, False
+    otherwise.
+  """
+  return (
+      inner_box.x1 >= outer_box.x1 - margin
+      and inner_box.y1 >= outer_box.y1 - margin
+      and inner_box.x2 <= outer_box.x2 + margin
+      and inner_box.y2 <= outer_box.y2 + margin
+  )
+
+
+def filter_boxes_keep_smaller(
+    data: Mapping[str, list[Any]],
+    iou_threshold: float = 0.8,
+    area_threshold: int | None = None,
+    min_area: int = 1000,
+    margin: int = 5,
+) -> dict[str, list[Any]]:
+  """Filters overlapping bounding boxes, preferentially keeping smaller ones.
+
+  This function sorts boxes by area and iterates through them, discarding any
+  box that has a high IoU with an already-kept box or is contained within one.
+  This is useful for eliminating duplicate or redundant detections.
+
+  Args:
+      data: A dictionary containing 'boxes' and 'masks' lists.
+      iou_threshold: The IoU value above which a box is considered an overlap.
+      area_threshold: An optional maximum area to consider for a box.
+      min_area: The minimum area required for a box to be kept.
+      margin: The pixel margin used for the containment check.
+
+  Returns:
+      A dictionary with the filtered 'boxes' and their corresponding 'masks'.
+  """
+  # Check if the input data is valid
+  bounding_boxes = [_BoundingBox(*b) for b in data['boxes']]
+
+  areas = ([_box_area(b) for b in bounding_boxes])
+
+  # Sort boxes from smallest to largest area
+  sorted_indices = np.argsort(areas)
+  sorted_bounding_boxes = [bounding_boxes[i] for i in sorted_indices]
+
+  masks = np.array(data['masks'])
+  sorted_masks = masks[sorted_indices]
+
+  kept_boxes = []
+  kept_masks = []
+  kept_bounding_boxes_for_check = []
+
+  for i, box in enumerate(sorted_bounding_boxes):
+    current_area = _box_area(box)
+    if (
+        area_threshold is not None and current_area > area_threshold
+    ) or current_area < min_area:
+      continue
+
+    keep = True
+    for kept_box in kept_bounding_boxes_for_check:
+      if _calculate_iou(box, kept_box) > iou_threshold or _is_contained(
+          kept_box, box, margin
+      ):
+        keep = False
+        break
+
+    if keep:
+      kept_boxes.append([box.x1, box.y1, box.x2, box.y2])
+      kept_masks.append(sorted_masks[i])
+      kept_bounding_boxes_for_check.append(box)
+
+  return {'boxes': kept_boxes, 'masks': kept_masks}
+
+
+def convert_boxes_cxcywh_to_xyxy(
+    boxes: torch.Tensor, image_shape: tuple[int, int, int]
+) -> np.ndarray:
+  """Converts bounding boxes from center-based to corner-based format.
+
+  Args:
+    boxes: A tensor of bounding boxes in (cx, cy, w, h) format.
+    image_shape: A tuple representing the image dimensions (h, w, c).
+
+  Returns:
+    A NumPy array of bounding boxes in (x1, y1, x2, y2) format.
+  """
+  h, w, _ = image_shape
+  scale_factors = torch.tensor([w, h, w, h], device=boxes.device)
+  scaled_boxes = boxes * scale_factors
+  xyxy_boxes = torchvision.ops.box_convert(
+      boxes=scaled_boxes, in_fmt='cxcywh', out_fmt='xyxy'
+  )
+  return xyxy_boxes.cpu().numpy().astype(int)

official/projects/waste_identification_ml/llm_applications/milk_pouch_detection/models_utils_test.py

@@ -0,0 +1,97 @@
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import numpy as np
+import torch
+from official.projects.waste_identification_ml.llm_applications.milk_pouch_detection import models_utils
+
+
+class UtilsTest(unittest.TestCase):
+  """Tests for the utility functions."""
+
+  def test_filter_boxes_keep_smaller_on_given_data(self):
+    boxes = [
+        [402.24, 0.54, 1343.04, 350.46],
+        [402.24, 0.54, 955.20, 333.18],
+        [930.24, 0.54, 1343.04, 351.54],
+        [611.52, 0.54, 955.20, 334.26],
+        [402.24, 0.54, 751.68, 305.10],
+        [749.76, 592.38, 1055.04, 875.34],
+        [941.76, 1012.50, 1039.68, 1078.38],
+    ]
+    masks = [f"mask_{i}" for i in range(len(boxes))]
+    data = {"boxes": boxes, "masks": masks}
+
+    expected_boxes = [
+        [941.76, 1012.50, 1039.68, 1078.38],
+        [749.76, 592.38, 1055.04, 875.34],
+        [402.24, 0.54, 751.68, 305.10],
+        [611.52, 0.54, 955.20, 334.26],
+        [930.24, 0.54, 1343.04, 351.54],
+    ]
+
+    result = models_utils.filter_boxes_keep_smaller(data)
+    actual_boxes = result["boxes"]
+
+    # Sort both lists for comparison (optional depending on importance of order)
+    actual_sorted = sorted(actual_boxes)
+    expected_sorted = sorted(expected_boxes)
+
+    self.assertEqual(len(actual_sorted), len(expected_sorted))
+    for box1, box2 in zip(actual_sorted, expected_sorted):
+      np.testing.assert_almost_equal(box1, box2, decimal=2)
+
+  def test_convert_boxes_cxcywh_to_xyxy_withsinglebox_returnscorrectcoordinates(
+      self,
+  ):
+    """Tests that a single box is converted correctly."""
+    boxes = torch.tensor([[0.5, 0.5, 0.2, 0.4]])  # cx, cy, w, h
+    image_shape = (100, 200, 3)  # h, w, c
+    expected_boxes = np.array([[80, 30, 120, 70]])  # x1, y1, x2, y2
+
+    converted_boxes = models_utils.convert_boxes_cxcywh_to_xyxy(
+        boxes, image_shape
+    )
+    np.testing.assert_array_equal(converted_boxes, expected_boxes)
+
+  def test_convert_boxes_cxcywh_to_xyxy_withmultipleboxes_returnscorrectcoordinates(
+      self,
+  ):
+    """Tests that multiple boxes are converted correctly."""
+    boxes = torch.tensor([
+        [0.5, 0.5, 0.2, 0.4],
+        [0.25, 0.25, 0.1, 0.1],
+    ])
+    image_shape = (100, 200, 3)
+    expected_boxes = np.array([
+        [80, 30, 120, 70],
+        [40, 20, 60, 30],
+    ])
+    converted_boxes = models_utils.convert_boxes_cxcywh_to_xyxy(
+        boxes, image_shape
+    )
+    np.testing.assert_array_equal(converted_boxes, expected_boxes)
+
+    boxes = torch.empty((0, 4))
+    image_shape = (100, 200, 3)
+    expected_boxes = np.empty((0, 4), dtype=int)
+    converted_boxes = models_utils.convert_boxes_cxcywh_to_xyxy(
+        boxes, image_shape
+    )
+    np.testing.assert_array_equal(converted_boxes, expected_boxes)
+
+
+if __name__ == "__main__":
+  unittest.main()