Update inference.py

inference.py

@@ -3,6 +3,7 @@ import math
 import cv2
 import os
 import numpy as np
+from PIL import Image
 import rasterio
 import tensorflow as tf
 from keras.models import load_model
@@ -111,18 +112,19 @@ def save_results(prediction, outputPath, map_name, legend):
     Save the prediction results to a specified output path.
 
     Parameters:
-    - prediction: The prediction result.
+    - prediction: The prediction result (should be a 2D or 3D numpy array).
     - outputPath: The directory where the results should be saved.
     - map_name: The name of the map.
     - legend: The legend associated with the prediction.
     """
-    global h5_image
-    
     output_image_path = os.path.join(outputPath, f"{map_name}_{legend}.tif")
 
-    # Save the prediction to the output path
-    with open(output_image_path, 'w') as f:
-        f.write(str(prediction))
+    # Convert the prediction to an image
+    # Note: The prediction array may need to be scaled or converted before saving as an image
+    prediction_image = Image.fromarray(prediction.astype(np.uint8))
+
+    # Save the prediction as a tiff image
+    prediction_image.save(output_image_path, 'TIFF')
 
     ### Waiting for georeferencing data
     # This section will be used in future releases to save georeferenced images.