Update inference.py

inference.py

@@ -28,7 +28,7 @@ def save_plot_as_png(prediction_result, map_name, legend, outputPath):
     true_seg = h5_image.get_layer(map_name, legend)
     full_map = h5_image.get_map(map_name)
     legend_patch = h5_image.get_legend(map_name, legend)
-    legend_resized = tf.image.resize(legend_patch, (h5_image.patch_size, h5_image.patch_size))
+    legend_resized = tf.image.resize(legend_patch, (true_seg.shape[0], true_seg.shape[1]))
 
     output_image_path = os.path.join(outputPath, f"{map_name}_{legend}_visual.png")