Update: class lables

yolov5.py

@@ -13,7 +13,7 @@ COLORS = np.random.uniform(0, 255, size=(80, 3))
 
 def parse_detections(results):
     detections = results.pandas().xyxy[0].to_dict()
-    boxes, colors, names = [], [], []
+    boxes, colors, names, classes = [], [], [], []
     for i in range(len(detections["xmin"])):
         confidence = detections["confidence"][i]
         if confidence < 0.2:
@@ -24,7 +24,8 @@ def parse_detections(results):
         boxes.append((xmin, ymin, xmax, ymax))
         colors.append(COLORS[category])
         names.append(name)
-    return boxes, colors, names
+        classes.append(category) 
+    return boxes, colors, names, classes
 
 
 def draw_detections(boxes, colors, names, classes, img):
@@ -68,8 +69,8 @@ def xai_yolov5(image):
 
     # Run YOLO detection
     results = model([image])
-    boxes, colors, names = parse_detections(results)
-    detections_img = draw_detections(boxes, colors, names, image.copy())
+    boxes, colors, names, classes = parse_detections(results) 
+    detections_img = draw_detections(boxes, colors, names,classes,  image.copy())
 
     # Prepare input tensor for Grad-CAM
     img_float = np.float32(image) / 255

yolov8.py

@@ -11,7 +11,7 @@ from ultralytics import YOLO
 
 COLORS = np.random.uniform(0, 255, size=(80, 3))
 def parse_detections(detections, model):
-    boxes, colors, names = [], [], []
+    boxes, colors, names, classes = [], [], [], []
     for detection in detections.boxes:
         xmin, ymin, xmax, ymax = map(int, detection.xyxy[0].tolist())
         confidence = detection.conf.item()
@@ -22,7 +22,8 @@ def parse_detections(detections, model):
         boxes.append((xmin, ymin, xmax, ymax))
         colors.append(COLORS[class_id])
         names.append(name)
-    return boxes, colors, names
+        classes.append(class_id) 
+    return boxes, colors, names, classes
 
 def draw_detections(boxes, colors, names, classes, img):
     for box, color, name, cls in zip(boxes, colors, names, classes):
@@ -55,8 +56,8 @@ def xai_yolov8s(image):
     model.eval()
     results = model(image)
     detections = results[0]
-    boxes, colors, names = parse_detections(detections, model)
-    detections_img = draw_detections(boxes, colors, names, image.copy())
+    boxes, colors, names, classes = parse_detections(detections, model)
+    detections_img = draw_detections(boxes, colors, names, classes, image.copy())
     img_float = np.float32(image) / 255
     transform = transforms.ToTensor()
     tensor = transform(img_float).unsqueeze(0)