Update inference.py

inference.py

@@ -212,6 +212,8 @@ def save_results(prediction, map_name, legend, outputPath):
 
     prediction_image = (prediction*255).astype(np.uint8)
 
+    prediction_image = np.expand_dims(prediction_image, axis=0)
+
     rasterio.open(output_image_path, 'w', driver='GTiff', compress='lzw',
                 height=prediction_image.shape[0], width=prediction_image.shape[1], count=1, dtype=prediction_image.dtype,
                 crs=h5_image.get_crs(map_name, legend), transform=h5_image.get_transform(map_name, legend)).write(prediction_image)