Upload 3 files

config.pbtxt

@@ -1,25 +1,18 @@
 backend: "python"
-max_batch_size: 0  # Keep batching disabled as per original config
+max_batch_size: 0
 
 input [
-  {
-    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
-  }
+{
+  name: "input_jp2_bytes"
+  data_type: TYPE_STRING
+  dims: [ 3 ]
+}
 ]
 
 output [
-  {
-    name: "output_mask"
-    data_type: TYPE_UINT8
-    dims: [-1, -1]     # Variable height, width
-  }
-]
-
-# Optional: Specify instance_group if running on GPU
-# instance_group [
-#   {
-#     kind: KIND_GPU
-#   }
-# ]
+{
+  name: "output_mask"
+  data_type: TYPE_UINT8
+  dims: [-1, -1]
+}
+]

model.py

@@ -4,98 +4,223 @@ from omnicloudmask import predict_from_array
 import rasterio
 from rasterio.io import MemoryFile
 from rasterio.enums import Resampling
+import tempfile
+import os
+from io import BytesIO
 
 class TritonPythonModel:
     def initialize(self, args):
         """
         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')
+        print('Initialized Cloud Detection model with JP2 input and robust GDAL handling')
 
-    def execute(self, requests):
+    def safe_read_jp2_bytes(self, jp2_bytes):
         """
-        Process inference requests.
+        Safely read JP2 bytes with multiple fallback methods
         """
-        responses = []
-        # Every request must contain three JP2 byte strings (Red, Green, NIR).
-        for request in requests:
-            # 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:
+            # Method 1: Try direct MemoryFile approach (works if GDAL drivers are properly configured)
+            with MemoryFile(jp2_bytes) as memfile:
+                with memfile.open() as src:
+                    data = src.read(1).astype(np.float32)
+                    height, width = src.height, src.width
+                    profile = src.profile
+                    return data, height, width, profile
+                    
+        except Exception as e1:
+            print(f"Method 1 (MemoryFile) failed: {e1}")
             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)
+                # Method 2: Write to temporary file and read from disk
+                with tempfile.NamedTemporaryFile(delete=False, suffix='.jp2') as tmp_file:
+                    tmp_file.write(jp2_bytes)
+                    tmp_file.flush()
+                    
+                    with rasterio.open(tmp_file.name) as src:
+                        data = src.read(1).astype(np.float32)
+                        height, width = src.height, src.width
+                        profile = src.profile
+                    
+                    # Clean up temporary file
+                    os.unlink(tmp_file.name)
+                    return data, height, width, profile
+                    
+            except Exception as e2:
+                print(f"Method 2 (temporary file) failed: {e2}")
+                try:
+                    # Method 3: Try with different suffix and basic profile
+                    with tempfile.NamedTemporaryFile(delete=False, suffix='.tiff') as tmp_file:
+                        tmp_file.write(jp2_bytes)
+                        tmp_file.flush()
+                        
+                        with rasterio.open(tmp_file.name) as src:
+                            data = src.read(1).astype(np.float32)
+                            height, width = src.height, src.width
+                            profile = {'driver': 'GTiff', 'height': height, 'width': width, 'count': 1, 'dtype': 'float32'}
+                        
+                        os.unlink(tmp_file.name)
+                        return data, height, width, profile
+                        
+                except Exception as e3:
+                    print(f"Method 3 (tiff fallback) failed: {e3}")
+                    # Method 4: Final fallback - try to interpret as raw numpy array
+                    try:
+                        # This assumes the client is sending raw numpy bytes as fallback
+                        data_array = np.frombuffer(jp2_bytes, dtype=np.float32)
+                        
+                        # Try to guess square dimensions
+                        side_length = int(np.sqrt(len(data_array)))
+                        if side_length * side_length == len(data_array):
+                            data = data_array.reshape(side_length, side_length)
+                            height, width = side_length, side_length
+                            profile = {'driver': 'GTiff', 'height': height, 'width': width, 'count': 1, 'dtype': 'float32'}
+                            return data, height, width, profile
                         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)
+                            # Try common satellite image dimensions
+                            common_dims = [(10980, 10980), (5490, 5490), (1024, 1024), (512, 512)]
+                            for h, w in common_dims:
+                                if h * w == len(data_array):
+                                    data = data_array.reshape(h, w)
+                                    height, width = h, w
+                                    profile = {'driver': 'GTiff', 'height': height, 'width': width, 'count': 1, 'dtype': 'float32'}
+                                    return data, height, width, profile
+                            
+                            raise ValueError(f"Cannot interpret data array of length {len(data_array)} as image")
+                            
+                    except Exception as e4:
+                        raise Exception(f"All fallback methods failed: MemoryFile({e1}), TempFile({e2}), TiffFallback({e3}), RawBytes({e4})")
+
+    def safe_resample_data(self, data, current_height, current_width, target_height, target_width, profile):
+        """
+        Safely resample data to target dimensions with fallback methods
+        """
+        if current_height == target_height and current_width == target_width:
+            return data
+            
+        try:
+            # Method 1: Use rasterio resampling
+            temp_profile = profile.copy()
+            temp_profile.update({
+                'height': current_height,
+                'width': current_width,
+                'count': 1,
+                'dtype': 'float32'
+            })
+            
+            with MemoryFile() as memfile:
+                with memfile.open(**temp_profile) as temp_dataset:
+                    temp_dataset.write(data, 1)
+                    
+                    resampled = temp_dataset.read(
+                        out_shape=(1, target_height, target_width),
+                        resampling=Resampling.bilinear
+                    )[0].astype(np.float32)
+                    
+                    return resampled
+                    
+        except Exception as e1:
+            print(f"Rasterio resampling failed: {e1}")
+            try:
+                # Method 2: Use scipy if available
+                from scipy import ndimage
+                zoom_factors = (target_height / current_height, target_width / current_width)
+                resampled = ndimage.zoom(data, zoom_factors, order=1)
+                return resampled.astype(np.float32)
+                
+            except ImportError:
+                print("Scipy not available for resampling")
+                # Method 3: Simple nearest-neighbor resampling
+                h_indices = np.round(np.linspace(0, current_height - 1, target_height)).astype(int)
+                w_indices = np.round(np.linspace(0, current_width - 1, target_width)).astype(int)
+                
+                resampled = data[np.ix_(h_indices, w_indices)]
+                return resampled.astype(np.float32)
+                
+            except Exception as e2:
+                print(f"Scipy resampling failed: {e2}")
+                # Method 3: Simple nearest-neighbor resampling
+                h_indices = np.round(np.linspace(0, current_height - 1, target_height)).astype(int)
+                w_indices = np.round(np.linspace(0, current_width - 1, target_width)).astype(int)
+                
+                resampled = data[np.ix_(h_indices, w_indices)]
+                return resampled.astype(np.float32)
 
-                # Perform inference using the original function
-                pred_mask = predict_from_array(input_array)
+    def execute(self, requests):
+        """
+        Process inference requests with robust error handling.
+        """
+        responses = []
 
-                # Create output tensor
-                output_tensor = pb_utils.Tensor(
-                    "output_mask",
-                    pred_mask.astype(np.uint8)
-                )
-                response = pb_utils.InferenceResponse([output_tensor])
+        for request in requests:
+            try:
+                input_tensor = pb_utils.get_input_tensor_by_name(request, "input_jp2_bytes")
+                jp2_bytes_list = input_tensor.as_numpy()
+
+                if len(jp2_bytes_list) != 3:
+                    error_msg = f"Expected 3 JP2 byte strings, received {len(jp2_bytes_list)}"
+                    error = pb_utils.TritonError(error_msg)
+                    response = pb_utils.InferenceResponse(output_tensors=[], error=error)
+                    responses.append(response)
+                    continue
+
+                red_bytes = jp2_bytes_list[0]
+                green_bytes = jp2_bytes_list[1]
+                nir_bytes = jp2_bytes_list[2]
+
+                print(f"Processing JP2 data - sizes: Red={len(red_bytes)}, Green={len(green_bytes)}, NIR={len(nir_bytes)}")
+
+                # Read red band data (use as reference for dimensions)
+                red_data, target_height, target_width, red_profile = self.safe_read_jp2_bytes(red_bytes)
+                print(f"Red band: {red_data.shape}, target dimensions: {target_height}x{target_width}")
+
+                # Read and resample green band
+                green_data, green_height, green_width, green_profile = self.safe_read_jp2_bytes(green_bytes)
+                green_data = self.safe_resample_data(green_data, green_height, green_width, target_height, target_width, green_profile)
+                print(f"Green band after resampling: {green_data.shape}")
+
+                # Read and resample NIR band
+                nir_data, nir_height, nir_width, nir_profile = self.safe_read_jp2_bytes(nir_bytes)
+                nir_data = self.safe_resample_data(nir_data, nir_height, nir_width, target_height, target_width, nir_profile)
+                print(f"NIR band after resampling: {nir_data.shape}")
+
+                # Verify all bands have the same shape
+                if not (red_data.shape == green_data.shape == nir_data.shape):
+                    shapes = [red_data.shape, green_data.shape, nir_data.shape]
+                    error_msg = f"Band shape mismatch after resampling: {shapes}"
+                    error = pb_utils.TritonError(error_msg)
+                    response = pb_utils.InferenceResponse(output_tensors=[], error=error)
+                    responses.append(response)
+                    continue
+
+                # Stack bands in CHW format for prediction (channels, height, width)
+                prediction_array = np.stack([red_data, green_data, nir_data], axis=0)
+                print(f"Final prediction array shape: {prediction_array.shape}")
+
+                # Run cloud detection prediction
+                cloud_mask = predict_from_array(prediction_array)
+                print(f"Cloud mask shape: {cloud_mask.shape}")
+
+                # Flatten the mask for output
+                if cloud_mask.ndim > 1:
+                    cloud_mask = cloud_mask.flatten()
+
+                # Create output tensor (config expects TYPE_UINT8)
+                output_tensor = pb_utils.Tensor("output_mask", cloud_mask.astype(np.uint8))
+                response = pb_utils.InferenceResponse(output_tensors=[output_tensor])
+                responses.append(response)
 
             except Exception as e:
-                # Handle errors during processing (e.g., invalid JP2 data)
-                error = pb_utils.TritonError(f"Error processing JP2 data: {str(e)}")
+                # Enhanced error reporting
+                error_msg = f"Error processing JP2 data: {str(e)}"
+                print(f"Model execution error: {error_msg}")
+                error = pb_utils.TritonError(error_msg)
                 response = pb_utils.InferenceResponse(output_tensors=[], error=error)
+                responses.append(response)
 
-            responses.append(response)
-
-        # Return a list of responses
         return responses
 
     def finalize(self):
         """
-        Called when the model is unloaded. Perform any necessary cleanup.
+        Clean up when the model is unloaded.
         """
-        print('Finalizing Cloud Detection model')
-
+        print('Cloud Detection model finalized')

requirements.txt

@@ -5,3 +5,5 @@ timm>=0.9
 tqdm>=4.0
 gdown>=5.1.0
 torch>=2.2
+scipy>=1.9.0
+numpy>=1.21.0