No public description

PiperOrigin-RevId: 764603158

Author: tensorflower-gardener

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

@@ -342,3 +342,116 @@ def resize_each_mask(
     )
     combined_masks.append(mask)
   return np.array(combined_masks)
+
+
+def extract_and_resize_objects(
+    results: Mapping[str, Any],
+    masks: str,
+    boxes: str,
+    image: np.ndarray,
+    resize_factor: float = 0.5,
+) -> Sequence[np.ndarray]:
+  """Extract and resize objects from the detection results.
+
+  Args:
+    results: The detection results from the model.
+    masks: The masks to extract objects from.
+    boxes: The bounding boxes of the objects.
+    image: The image to extract objects from.
+    resize_factor: The factor by which to resize the objects.
+
+  Returns:
+    A list of cropped objects.
+  """
+  cropped_objects = []
+
+  for i, mask in enumerate(results[masks]):
+    ymin, xmin, ymax, xmax = results[boxes][0][i]
+    mask = np.expand_dims(mask, axis=-1)
+
+    # Crop the object using the mask and bounding box
+    cropped_object = np.where(
+        mask[ymin:ymax, xmin:xmax], image[ymin:ymax, xmin:xmax], 0
+    )
+
+    # Calculate new dimensions
+    new_width = int(cropped_object.shape[1] * resize_factor)
+    new_height = int(cropped_object.shape[0] * resize_factor)
+    cropped_object = cv2.resize(
+        cropped_object, (new_width, new_height), interpolation=cv2.INTER_AREA
+    )
+    cropped_objects.append(cropped_object)
+
+  return cropped_objects
+
+
+def adjust_image_size(
+    height: int, width: int, min_size: int
+) -> tuple[int, int]:
+  """Adjust the image size to ensure both dimensions are at least of min_size.
+
+  Args:
+    height: The height of the image.
+    width: The width of the image.
+    min_size: Minimum size of the image dimension needed.
+
+  Returns:
+    The adjusted height and width of the image.
+  """
+  if height < min_size or width < min_size:
+    return height, width
+
+  # Calculate the scale factor to ensure both dimensions remain at least 1024
+  scale_factor = min(height / min_size, width / min_size)
+  return int(height / scale_factor), int(width / scale_factor)
+
+
+def filter_detections(
+    results: Mapping[str, np.ndarray],
+    valid_indices: Sequence[int] | Sequence[bool],
+) -> Mapping[str, np.ndarray]:
+  """Filter the detection results based on the valid indices.
+
+  Args:
+    results: The detection results from the model.
+    valid_indices: The indices of the valid detections.
+
+  Returns:
+    The filtered detection results.
+  """
+  if np.array(valid_indices).dtype == bool:
+    new_num_detections = int(np.sum(valid_indices))
+  else:
+    new_num_detections = len(valid_indices)
+
+  # Define the keys to filter
+  keys_to_filter = [
+      'detection_masks',
+      'detection_masks_resized',
+      'detection_masks_reframed',
+      'detection_classes',
+      'detection_boxes',
+      'normalized_boxes',
+      'detection_scores',
+  ]
+
+  filtered_output = {}
+
+  for key in keys_to_filter:
+    if key in results:
+      if key == 'detection_masks':
+        filtered_output[key] = results[key][:, valid_indices, :, :]
+      elif key in ['detection_masks_resized', 'detection_masks_reframed']:
+        filtered_output[key] = results[key][valid_indices, :, :]
+      elif key in ['detection_boxes', 'normalized_boxes']:
+        filtered_output[key] = results[key][:, valid_indices, :]
+      elif key in [
+          'detection_classes',
+          'detection_scores',
+          'detection_classes_names',
+      ]:
+        filtered_output[key] = results[key][:, valid_indices]
+  filtered_output['image_info'] = results['image_info']
+  filtered_output['num_detections'] = np.array([new_num_detections])
+
+  return filtered_output

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

@@ -113,6 +113,115 @@ def test_keeps_biggest_mask(self):
     # Expect only the largest mask (index 2) to remain
     self.assertEqual(result, [2])
 
+  def test_filter_with_boolean_indices(self):
+    results = {
+        'detection_masks': np.random.rand(1, 3, 5, 5),
+        'detection_masks_resized': np.random.rand(3, 5, 5),
+        'detection_boxes': np.random.rand(1, 3, 4),
+        'detection_classes': np.array([[1, 2, 3]]),
+        'detection_scores': np.array([[0.9, 0.8, 0.3]]),
+        'image_info': np.array([[640, 480]]),
+    }
+
+    valid_indices = [True, False, True]
+
+    output = utils.filter_detections(results, valid_indices)
+
+    self.assertEqual(output['detection_masks'].shape[1], 2)
+    self.assertEqual(output['detection_masks_resized'].shape[0], 2)
+    self.assertEqual(output['detection_boxes'].shape[1], 2)
+    self.assertEqual(output['detection_classes'].shape[1], 2)
+    self.assertEqual(output['detection_scores'].shape[1], 2)
+    self.assertTrue(np.array_equal(output['image_info'], results['image_info']))
+    self.assertEqual(output['num_detections'][0], 2)
+
+  def test_filter_with_integer_indices(self):
+    results = {
+        'detection_masks': np.random.rand(1, 4, 5, 5),
+        'detection_masks_resized': np.random.rand(4, 5, 5),
+        'detection_boxes': np.random.rand(1, 4, 4),
+        'detection_classes': np.array([[1, 2, 3, 4]]),
+        'detection_scores': np.array([[0.9, 0.8, 0.3, 0.6]]),
+        'image_info': np.array([[640, 480]]),
+    }
+
+    valid_indices = [0, 2]  # Keep detections at index 0 and 2
+
+    output = utils.filter_detections(results, valid_indices)
+
+    self.assertEqual(output['detection_masks'].shape[1], 2)
+    self.assertEqual(output['detection_masks_resized'].shape[0], 2)
+    self.assertEqual(output['detection_boxes'].shape[1], 2)
+    self.assertEqual(output['detection_classes'].shape[1], 2)
+    self.assertEqual(output['detection_scores'].shape[1], 2)
+    self.assertEqual(output['num_detections'][0], 2)
+
+  def test_both_dimensions_below_min_size(self):
+    height, width, min_size = 800, 900, 1024
+
+    result = utils.adjust_image_size(height, width, min_size)
+
+    self.assertEqual(result, (800, 900))  # No scaling should happen
+
+  def test_height_below_min_size(self):
+    height, width, min_size = 900, 1200, 1024
+
+    result = utils.adjust_image_size(height, width, min_size)
+
+    self.assertEqual(result, (900, 1200))  # No scaling
+
+  def test_width_below_min_size(self):
+    height, width, min_size = 1300, 800, 1024
+
+    result = utils.adjust_image_size(height, width, min_size)
+
+    self.assertEqual(result, (1300, 800))  # No scaling
+
+  def test_both_dimensions_above_min_size(self):
+    height, width, min_size = 2048, 1536, 1024
+    expected_scale = min(height / min_size, width / min_size)
+    expected_height = int(height / expected_scale)
+    expected_width = int(width / expected_scale)
+
+    result = utils.adjust_image_size(height, width, min_size)
+
+    self.assertEqual(result, (expected_height, expected_width))
+
+  def test_exact_min_size(self):
+    height, width, min_size = 1024, 1024, 1024
+
+    result = utils.adjust_image_size(height, width, min_size)
+
+    self.assertEqual(result, (1024, 1024))  # Already meets the requirement
+
+  def test_extract_and_resize_single_object(self):
+    image = np.ones((10, 10, 3), dtype=np.uint8) * 255  # white image
+
+    # Define a simple binary mask (1 in a 4x4 box)
+    mask = np.zeros((10, 10), dtype=np.uint8)
+    mask[2:6, 3:7] = 1
+
+    # Box coordinates match the mask
+    boxes = np.array([[[2, 3, 6, 7]]], dtype=np.int32)  # shape (1, 1, 4)
+
+    results = {'masks': [mask], 'boxes': boxes}
+
+    cropped_objects = utils.extract_and_resize_objects(
+        results, 'masks', 'boxes', image, resize_factor=0.5
+    )
+
+    self.assertEqual(len(cropped_objects), 1)
+    obj = cropped_objects[0]
+
+    # Original crop size is (4, 4), so resized should be (2, 2)
+    self.assertEqual(obj.shape[:2], (2, 2))
+
+    # Should still be 3 channels
+    self.assertEqual(obj.shape[2], 3)
+
+    # The output pixels in mask area should be non-zero
+    self.assertTrue(np.any(obj > 0))
+
 
 if __name__ == '__main__':
   unittest.main()