debug: hstack dimensions inconsistency

yolov5.py

@@ -77,8 +77,7 @@ def xai_yolov5(image):
     cam_image, renormalized_cam_image = generate_cam_image(model, target_layers, tensor, image, boxes)
 
     rgb_img_float, batch_explanations = dff_nmf(image, target_lyr = -5, n_components = 8)
-    visualize_batch_explanations(rgb_img_float, batch_explanations)
-    im = cv2.imread("test_w.png")
+    im = visualize_batch_explanations(rgb_img_float, batch_explanations)
     
     # Combine results
     print('shapes........', image.shape, detections_img.shape, renormalized_cam_image.shape, im.shape)
@@ -248,3 +247,5 @@ def visualize_batch_explanations(rgb_img_float, batch_explanations, image_weight
         plt.axis('off')  # Hide axes
         plt.show()  # Show the plot
     plt.savefig("test_w.png")
+    print('viz', visualization.shape)
+    return visualization