Update: class lables

yolov5.py

@@ -45,11 +45,7 @@ def generate_cam_image(model, target_layers, tensor, rgb_img, boxes):
     cam = EigenCAM(model, target_layers)
     grayscale_cam = cam(tensor)[0, :, :]
     img_float = np.float32(rgb_img) / 255
-
-    # Generate Grad-CAM
     cam_image = show_cam_on_image(img_float, grayscale_cam, use_rgb=True)
-
-    # Renormalize Grad-CAM inside bounding boxes
     renormalized_cam = np.zeros(grayscale_cam.shape, dtype=np.float32)
     for x1, y1, x2, y2 in boxes:
         renormalized_cam[y1:y2, x1:x2] = scale_cam_image(grayscale_cam[y1:y2, x1:x2].copy())
@@ -60,7 +56,6 @@ def generate_cam_image(model, target_layers, tensor, rgb_img, boxes):
 
 
 def xai_yolov5(image):
-    # Load YOLOv5 model
     model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
     model.eval()
     model.cpu()
@@ -81,6 +76,9 @@ def xai_yolov5(image):
     cam_image, renormalized_cam_image = generate_cam_image(model, target_layers, tensor, image, boxes)
 
     # Combine results
-    final_image = np.hstack((image, cam_image, renormalized_cam_image))
+    final_image = np.hstack((image, detections_img, renormalized_cam_image))
     caption = "Results using YOLOv5"
-    return Image.fromarray(final_image), caption
+    return Image.fromarray(final_image), caption
+
+
+