updated args passing

yolov5.py

@@ -76,14 +76,13 @@ def xai_yolov5(image):
     # Grad-CAM visualization
     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)
+    rgb_img_float, batch_explanations, result = dff_nmf(image, target_lyr = -5, n_components = 8)
     im = visualize_batch_explanations(rgb_img_float, batch_explanations)
     
     # Combine results
-    print('shapes........', image.shape, detections_img.shape, renormalized_cam_image.shape, im.shape)
     final_image = np.hstack((image, detections_img, renormalized_cam_image, im))
     caption = "Results using YOLOv5"
-    return Image.fromarray(final_image), caption
+    return Image.fromarray(final_image), caption, result
 
 
 
@@ -194,8 +193,8 @@ def dff_nmf(image, target_lyr, n_components):
     yaml_data = requests.get(yolov5_categories_url).text
     labels = yaml.safe_load(yaml_data)['names']  # Parse the YAML file to get class names
     num_classes = model.model.model.model[-1].nc 
-
-    for indx in range( explanations[0].shape[0]):
+    results = []
+    for indx in range(explanations[0].shape[0]):
         upsampled_input =  explanations[0][indx]
         upsampled_input = torch.tensor(upsampled_input)
         device = next(model.parameters()).device
@@ -234,7 +233,9 @@ def dff_nmf(image, target_lyr, n_components):
         plt.show()
         plt.savefig("test_" + str(indx) + ".png" )
         plt.clf()
-        return rgb_img_float, batch_explanations
+        results.append(Image.open(f"test_{indx}.png"))
+
+    return rgb_img_float, batch_explanations, results
 
 
 def visualize_batch_explanations(rgb_img_float, batch_explanations, image_weight=0.7):