jp2 as input type

config.pbtxt

@@ -1,16 +1,25 @@
 backend: "python"
-max_batch_size: 0  # Disable batching if not supported
+max_batch_size: 0  # Keep batching disabled as per original config
+
 input [
   {
-    name: "input_array"
-    data_type: TYPE_FP32
-    dims: [3, -1, -1]  # Channels, height, width
+    name: "input_jp2_bytes"  # New input name for JP2 bytes
+    data_type: TYPE_STRING  # Use TYPE_STRING for bytes
+    dims: [ 3 ]             # Expecting 3 elements: Red, Green, NIR bytes
   }
 ]
+
 output [
   {
     name: "output_mask"
-    data_type: TYPE_UINT8  # Adjust based on actual output type
-    dims: [-1, -1]     # Height, width
+    data_type: TYPE_UINT8
+    dims: [-1, -1]     # Variable height, width
   }
-]
+]
+
+# Optional: Specify instance_group if running on GPU
+# instance_group [
+#   {
+#     kind: KIND_GPU
+#   }
+# ]

model.py

@@ -1,28 +1,101 @@
 import numpy as np
 import triton_python_backend_utils as pb_utils
 from omnicloudmask import predict_from_array
+import rasterio
+from rasterio.io import MemoryFile
+from rasterio.enums import Resampling
 
 class TritonPythonModel:
     def initialize(self, args):
-        pass
+        """
+        Initialize the model. This function is called once when the model is loaded.
+        """
+        # You can load models or initialize resources here if needed.
+        # Ensure rasterio is installed in the Python backend environment.
+        print('Initialized Cloud Detection model with JP2 input')
 
     def execute(self, requests):
+        """
+        Process inference requests.
+        """
         responses = []
+        # Every request must contain three JP2 byte strings (Red, Green, NIR).
         for request in requests:
-            # Get input tensor
-            input_tensor = pb_utils.get_input_tensor_by_name(request, "input_array")
-            input_array = input_tensor.as_numpy()
-
-            # Perform inference
-            pred_mask = predict_from_array(input_array)
-
-            # Create output tensor
-            output_tensor = pb_utils.Tensor(
-                "output_mask", 
-                pred_mask.astype(np.uint8)
-            )
-            responses.append(pb_utils.InferenceResponse([output_tensor]))
+            # Get the input tensor containing the byte arrays
+            input_tensor = pb_utils.get_input_tensor_by_name(request, "input_jp2_bytes")
+            # as_numpy() for TYPE_STRING gives an ndarray of Python bytes objects
+            jp2_bytes_list = input_tensor.as_numpy()
+
+            if len(jp2_bytes_list) != 3:
+                # Send an error response if the input shape is incorrect
+                error = pb_utils.TritonError(f"Expected 3 JP2 byte strings, received {len(jp2_bytes_list)}")
+                response = pb_utils.InferenceResponse(output_tensors=[], error=error)
+                responses.append(response)
+                continue # Skip to the next request
+
+            # Assume order: Red, Green, NIR based on client logic
+            red_bytes = jp2_bytes_list[0]
+            green_bytes = jp2_bytes_list[1]
+            nir_bytes = jp2_bytes_list[2]
+
+            try:
+                # Process JP2 bytes using rasterio in memory
+                with MemoryFile(red_bytes) as memfile_red:
+                    with memfile_red.open() as src_red:
+                        red_data = src_red.read(1).astype(np.float32)
+                        target_height = src_red.height
+                        target_width = src_red.width
+
+                with MemoryFile(green_bytes) as memfile_green:
+                    with memfile_green.open() as src_green:
+                        # Ensure green band matches red band dimensions (should if B03)
+                        if src_green.height != target_height or src_green.width != target_width:
+                             # Optional: Resample green if necessary, though B03 usually matches B04
+                             green_data = src_green.read(
+                                 1,
+                                 out_shape=(1, target_height, target_width),
+                                 resampling=Resampling.bilinear
+                             ).astype(np.float32)
+                        else:
+                             green_data = src_green.read(1).astype(np.float32)
+
+
+                with MemoryFile(nir_bytes) as memfile_nir:
+                    with memfile_nir.open() as src_nir:
+                        # Resample NIR (B8A) to match Red/Green (B04/B03) resolution
+                        nir_data = src_nir.read(
+                            1, # Read the first band
+                            out_shape=(1, target_height, target_width),
+                            resampling=Resampling.bilinear
+                        ).astype(np.float32)
+
+                # Stack bands in CHW format (Red, Green, NIR) for the model
+                # Match the channel order expected by predict_from_array
+                input_array = np.stack([red_data, green_data, nir_data], axis=0)
+
+                # Perform inference using the original function
+                pred_mask = predict_from_array(input_array)
+
+                # Create output tensor
+                output_tensor = pb_utils.Tensor(
+                    "output_mask",
+                    pred_mask.astype(np.uint8)
+                )
+                response = pb_utils.InferenceResponse([output_tensor])
+
+            except Exception as e:
+                # Handle errors during processing (e.g., invalid JP2 data)
+                error = pb_utils.TritonError(f"Error processing JP2 data: {str(e)}")
+                response = pb_utils.InferenceResponse(output_tensors=[], error=error)
+
+            responses.append(response)
+
+        # Return a list of responses
         return responses
 
     def finalize(self):
-        pass
+        """
+        Called when the model is unloaded. Perform any necessary cleanup.
+        """
+        print('Finalizing Cloud Detection model')
+