Update inference.py

inference.py

@@ -145,7 +145,7 @@ def main(args):
 
     # Load the HDF5 file using the H5Image class
     print("Loading the HDF5 file.")
-    h5_image = H5Image(args.mapPath, mode='r')
+    h5_image = H5Image(args.mapPath, mode='r', patch_border=0)
 
     # Get map details
     print("Getting map details.")
@@ -207,10 +207,23 @@ def main(args):
                 prediction = perform_inference(legend_patch, map_patch, model)
 
                 
-                # Place the prediction in the corresponding position in the full_prediction array
+                # Calculate starting indices for rows and columns
                 x_start = row * h5_image.patch_size
                 y_start = col * h5_image.patch_size
-                full_prediction[x_start:x_start+h5_image.patch_size, y_start:y_start+h5_image.patch_size] = prediction
+
+                # Calculate ending indices for rows and columns
+                x_end = x_start + h5_image.patch_size
+                y_end = y_start + h5_image.patch_size
+
+                # Adjust the ending indices if they go beyond the image size
+                x_end = min(x_end, map_width)
+                y_end = min(y_end, map_height)
+
+                # Adjust the shape of the prediction if necessary
+                prediction_shape_adjusted = prediction[:x_end-x_start, :y_end-y_start]
+
+                # Assign the prediction to the correct location in the full_prediction array
+                full_prediction[x_start:x_end, y_start:y_end] = prediction_shape_adjusted
         
         # Mask out the map background pixels from the prediction
         print("Applying mask to the full prediction.")