Upload model.py

model.py

@@ -65,7 +65,7 @@ class TritonPythonModel:
                     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
+                        # This assumes the data might be raw numpy bytes as fallback
                         data_array = np.frombuffer(jp2_bytes, dtype=np.float32)
                         
                         # Try to guess square dimensions
@@ -163,11 +163,39 @@ class TritonPythonModel:
                     responses.append(response)
                     continue
 
-                red_bytes = jp2_bytes_list[0]
-                green_bytes = jp2_bytes_list[1]
-                nir_bytes = jp2_bytes_list[2]
+                # The input might be hex strings, decode them to bytes first
+                red_hex = jp2_bytes_list[0]
+                green_hex = jp2_bytes_list[1] 
+                nir_hex = jp2_bytes_list[2]
 
-                print(f"Processing JP2 data - sizes: Red={len(red_bytes)}, Green={len(green_bytes)}, NIR={len(nir_bytes)}")
+                # Convert hex strings to bytes
+                try:
+                    if isinstance(red_hex, str):
+                        red_bytes = bytes.fromhex(red_hex)
+                        green_bytes = bytes.fromhex(green_hex)
+                        nir_bytes = bytes.fromhex(nir_hex)
+                        print(f"Decoded hex strings to bytes")
+                    elif isinstance(red_hex, (bytes, np.bytes_)):
+                        # Already bytes, use directly
+                        red_bytes = bytes(red_hex)
+                        green_bytes = bytes(green_hex) 
+                        nir_bytes = bytes(nir_hex)
+                        print(f"Input already in bytes format")
+                    else:
+                        # Might be numpy string object
+                        red_bytes = bytes.fromhex(str(red_hex))
+                        green_bytes = bytes.fromhex(str(green_hex))
+                        nir_bytes = bytes.fromhex(str(nir_hex))
+                        print(f"Converted numpy strings to bytes")
+                except Exception as e:
+                    error_msg = f"Failed to decode input data: {str(e)}"
+                    print(f"Decode error: {error_msg}")
+                    error = pb_utils.TritonError(error_msg)
+                    response = pb_utils.InferenceResponse(output_tensors=[], error=error)
+                    responses.append(response)
+                    continue
+
+                print(f"Processing JP2 data - decoded 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)
@@ -192,6 +220,14 @@ class TritonPythonModel:
                     responses.append(response)
                     continue
 
+                # Check for valid dimensions
+                if red_data.shape[0] == 0 or red_data.shape[1] == 0:
+                    error_msg = f"Invalid band dimensions: {red_data.shape}"
+                    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}")