Update yolov5.py

yolov5.py

@@ -235,7 +235,7 @@ def dff_nmf(image, target_lyr, n_components):
         image_array = np.array(fig.canvas.renderer.buffer_rgba())  # Convert to numpy array
         
         temp = image_array.reshape((rgb_img_float.shape[0],rgb_img_float.shape[1]) )
-        visualization = show_factorization_on_image(rgb_img_float,np.stack([temp] * 3, axis=-1) , image_weight=0.3)
+        visualization = show_factorization_on_image(rgb_img_float, image_array.resize((rgb_img_float.shape)) , image_weight=0.3)
         results.append(visualization)
         plt.clf()  
         #return image_array