No public description

PiperOrigin-RevId: 755907301

Author: tensorflower-gardener

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

@@ -0,0 +1,105 @@
+# Copyright 2024 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.
+
+"""Prediction from the Triton server."""
+
+from typing import Any
+import cv2
+import numpy as np
+import tritonclient
+
+_API_URL = 'localhost:8000'
+_OUTPUT_KEYS = (
+    'detection_classes',
+    'detection_masks',
+    'detection_boxes',
+    'image_info',
+    'num_detections',
+    'detection_scores',
+)
+
+# Setting up the Triton client
+_TRITON_CLIENT = tritonclient.http.InferenceServerClient(
+    url=_API_URL, network_timeout=1200, connection_timeout=1200
+)
+
+# Outputs setup based on constants
+_OUTPUTS = [
+    tritonclient.http.InferRequestedOutput(key, binary_data=True)
+    for key in _OUTPUT_KEYS
+]
+
+
+def model_input(
+    path: str, height: int, width: int
+) -> tritonclient.http.InferInput:
+  """Prepares an image for input to a Triton model server.
+
+  It reads it from a path, resizes it, normalizes it, and converts it to the
+  format required by the server.
+
+  Args:
+      path: The file path to the image that needs to be processed.
+      height: The height of the image to be resized.
+      width: The width of the image to be resized.
+
+  Returns:
+      A Triton inference server input object containing the processed image.
+  """
+  original_image = cv2.imread(path)
+  image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
+  image_resized = cv2.resize(
+      image, (width, height), interpolation=cv2.INTER_AREA
+  )
+  expanded_image = np.expand_dims(image_resized, axis=0)
+  inputs = tritonclient.http.InferInput(
+      'inputs', expanded_image.shape, datatype='UINT8'
+  )
+  inputs.set_data_from_numpy(expanded_image, binary_data=True)
+  return inputs, image, image_resized
+
+
+def _query_model(
+    client: tritonclient.http.InferenceServerClient,
+    model_name: str,
+    inputs: tritonclient.http.InferInput,
+) -> tritonclient.http.InferResult:
+  """Sends an inference request to the Triton server.
+
+  Args:
+      client: The Triton server client.
+      model_name: Name of the model for which inference is requested.
+      inputs: The input data for inference.
+
+  Returns:
+      The result of the inference request.
+  """
+  return client.infer(model_name=model_name, inputs=[inputs], outputs=_OUTPUTS)
+
+
+def prediction(
+    model_name: str, inputs: tritonclient.http.InferInput
+) -> dict[str, Any]:
+  """Model name for prediction.
+
+  Args:
+      model_name: Model name in Triton Server.
+      inputs: The input data for inference.
+
+  Returns:
+      prediction output from the model.
+  """
+  result = _query_model(_TRITON_CLIENT, model_name, inputs)
+  result_dict = {key: result.as_numpy(key) for key in _OUTPUT_KEYS}
+  return result_dict

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

@@ -0,0 +1,70 @@
+# Copyright 2024 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
+from unittest import mock
+import numpy as np
+# Import the functions to be tested
+from official.projects.waste_identification_ml.Triton_TF_Cloud_Deployment.client import triton_server_prediction
+
+
+class TestTritonPrediction(unittest.TestCase):
+
+  @mock.patch("cv2.imread")
+  @mock.patch("cv2.cvtColor")
+  @mock.patch("cv2.resize")
+  def test_model_input(self, mock_resize, mock_convert_color, mock_imread):
+    """Test the model_input function."""
+
+    # Mocking image loading and processing
+    mock_imread.return_value = np.ones((500, 500, 3), dtype=np.uint8)
+    mock_convert_color.return_value = np.ones((500, 500, 3), dtype=np.uint8)
+    mock_resize.return_value = np.ones((224, 224, 3), dtype=np.uint8)
+
+    _, image, image_resized = triton_server_prediction.model_input(
+        "dummy_path.jpg", 224, 224
+    )
+
+    self.assertEqual(image.shape, (500, 500, 3))
+    self.assertEqual(image_resized.shape, (224, 224, 3))
+
+  @mock.patch("your_module._query_model")
+  def test_prediction(self, mock_query_model):
+    """Test the prediction function."""
+
+    mock_result = mock.MagicMock()
+    mock_result.as_numpy.side_effect = lambda key: np.array(
+        [1]
+    )
+    mock_query_model.return_value = mock_result
+
+    mock_inputs = mock.MagicMock()
+    model_name = "dummy_model"
+
+    result = triton_server_prediction.prediction(model_name, mock_inputs)
+
+    for key in (
+        "detection_classes",
+        "detection_masks",
+        "detection_boxes",
+        "image_info",
+        "num_detections",
+        "detection_scores",
+    ):
+      self.assertIn(key, result)
+      self.assertIsInstance(result[key], np.ndarray)
+
+
+if __name__ == "__main__":
+  unittest.main()