No public description

PiperOrigin-RevId: 862264383

Author: tensorflower-gardener

official/projects/waste_identification_ml/Deploy/detr_cloud_deployment/client/requirements.sh

@@ -32,6 +32,8 @@ sudo apt-get install -y python3-venv python3-pip
 python3.10 -m venv myenv
 source myenv/bin/activate
 
+echo "Activated python environment, installing dependencies."
+
 pip install --no-cache-dir natsort absl-py opencv-python pandas pandas-gbq \
   google-cloud-bigquery google-auth trackpy google-cloud-storage \
   scikit-image scikit-learn webcolors==1.13 ffmpeg-python tritonclient[all] \

official/projects/waste_identification_ml/Deploy/detr_cloud_deployment/client/triton_server_inference.py

@@ -81,10 +81,18 @@ def _resize_mask_batch(
   ) -> np.ndarray:
     """Resizes a batch of masks to the target dimensions."""
     target_w, target_h = target_dims
-    return cv2.resize(
-        masks, (target_w, target_h), interpolation=cv2.INTER_NEAREST
+    masks_transposed = np.transpose(masks, (1, 2, 0))
+
+    resized_batch = cv2.resize(
+        masks_transposed, (target_w, target_h), interpolation=cv2.INTER_NEAREST
     )
 
+    # If N=1, cv2.resize might drop the last dim, so we ensure 3D
+    if resized_batch.ndim == 2:
+      return resized_batch[np.newaxis, ...]
+
+    return np.transpose(resized_batch, (2, 0, 1))
+
   def _scale_bbox_and_masks(
       self, results: Dict[str, Any], target_dims: Tuple[int, int]
   ) -> Dict[str, Any]:

official/projects/waste_identification_ml/Deploy/detr_cloud_deployment/client/triton_server_inference_test.py

@@ -97,14 +97,12 @@ def test_box_cxcywh_to_xyxyn(self):
     expected = np.array([[0.4, 0.4, 0.6, 0.6], [0.0, 0.0, 1.0, 1.0]])
     np.testing.assert_allclose(result, expected)
 
-  @mock.patch(f"{MODULE_PATH}.Image.fromarray")
-  def test_scale_bbox_and_masks(self, mock_fromarray):
+  @mock.patch(f"{MODULE_PATH}.cv2.resize")
+  def test_scale_bbox_and_masks(self, mock_cv2_resize):
     # Arrange
-    mock_img_instance = mock.Mock()
-    mock_fromarray.return_value = mock_img_instance
-    mock_img_instance.resize.return_value = np.ones(
+    mock_cv2_resize.return_value = np.ones(
         (20, 10)
-    )  # w=10, h=20 -> PIL resize returns (width, height) numpy array
+    )  # h=20, w=10. cv2 resize returns h, w
 
     results = {
         "xyxy": np.array([[0.1, 0.1, 0.5, 0.5]]),
@@ -118,8 +116,14 @@ def test_scale_bbox_and_masks(self, mock_fromarray):
     # Assert
     np.testing.assert_allclose(scaled_results["xyxy"], [[1.0, 2.0, 5.0, 10.0]])
     self.assertEqual(scaled_results["masks"].shape, (1, 20, 10))
-    mock_fromarray.assert_called_once_with(results["masks"][0])
-    mock_img_instance.resize.assert_called_once_with((10, 20))
+    self.assertEqual(scaled_results["masks"].dtype, bool)
+    mock_cv2_resize.assert_called_once()
+    args, kwargs = mock_cv2_resize.call_args
+    np.testing.assert_array_equal(args[0], np.ones((5, 5, 1)))
+    self.assertEqual(args[1], (10, 20))
+    self.assertEqual(
+        kwargs["interpolation"], triton_server_inference.cv2.INTER_NEAREST
+    )
 
   @mock.patch(f"{MODULE_PATH}.cv2.imread")
   @mock.patch(f"{MODULE_PATH}.cv2.cvtColor")

official/projects/waste_identification_ml/Deploy/detr_cloud_deployment/server/triton_inference_server.sh

@@ -33,6 +33,8 @@ command -v screen >/dev/null 2>&1 || { \
   sudo apt update && sudo apt install -y screen; \
 }
 
+echo "Starting Triton server in a screen session."
+
 # Start Triton server
 screen -dmS server bash -c '
 sudo docker run --gpus all --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 \