Update yolov5.py

yolov5.py

@@ -213,7 +213,7 @@ def dff_nmf(image, target_lyr, n_components):
         boxes = boxes[confidence_mask]  # Filter boxes by confidence mask
         fig, ax = plt.subplots(1, figsize=(8, 8))
         ax.axis("off")
-        ax.imshow(torch.tensor(batch_explanations[0][indx]).cpu().numpy(), cmap="Spectral")  # Display image
+        ax.imshow(torch.tensor(batch_explanations[0][indx]).cpu().numpy(), cmap="bwr")  # Display image
         top_score_idx = scores.argmax(dim=0)  # Get the index of the max score
         top_score = scores[top_score_idx].item()
         top_class_id = class_ids[top_score_idx].item()
@@ -242,7 +242,7 @@ def dff_nmf(image, target_lyr, n_components):
         
         #temp = image_array.reshape((rgb_img_float.shape[0],rgb_img_float.shape[1]) )
         #visualization = show_factorization_on_image(rgb_img_float, image_array.resize((rgb_img_float.shape)) , image_weight=0.3)
-        visualization = show_factorization_on_image(rgb_img_float, np.transpose(rgb_channels, (2, 0, 1)), image_weight=0.7)
+        visualization = show_factorization_on_image(rgb_img_float, np.transpose(rgb_channels, (2, 0, 1)), image_weight=0.3)
         results.append(visualization)
         plt.clf()  
         #return image_array