Update: class lables

yolov5.py

@@ -7,13 +7,10 @@ from pytorch_grad_cam import EigenCAM
 from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
 import gradio as gr
 
-# Global Color Palette
 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 +21,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)  # Add class ID
+    return boxes, colors, names, classes
 
 
 def draw_detections(boxes, colors, names, classes, img):
@@ -44,17 +42,12 @@ 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())
     renormalized_cam = scale_cam_image(renormalized_cam)
     renormalized_cam_image = show_cam_on_image(img_float, renormalized_cam, use_rgb=True)
-
     return cam_image, renormalized_cam_image
 
 
@@ -68,8 +61,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)  # Extract classes
+    detections_img = draw_detections(boxes, colors, names, classes, image.copy())
 
     # Prepare input tensor for Grad-CAM
     img_float = np.float32(image) / 255
@@ -80,6 +73,19 @@ 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))
-    caption = "Results using YOLOv5"
-    return Image.fromarray(final_image), caption
+    final_image = np.hstack((
+        cv2.cvtColor(image, cv2.COLOR_RGB2BGR),
+        detections_img,
+        cv2.cvtColor(cam_image, cv2.COLOR_RGB2BGR)
+    ))
+    caption = "Results using YOLOv5 with Class Labels and Bounding Boxes"
+    return Image.fromarray(final_image), caption
+
+
+# Example Gradio interface
+gr.Interface(
+    fn=xai_yolov5,
+    inputs=gr.Image(type="numpy"),
+    outputs=[gr.Image(type="pil"), gr.Text()],
+    title="YOLOv5 with XAI and Class Labels"
+).launch()

yolov8.py

@@ -10,8 +10,9 @@ 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 +23,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)  # Add 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):
@@ -51,17 +53,33 @@ def generate_cam_image(model, target_layers, tensor, rgb_img, boxes):
     return cam_image, renormalized_cam_image
 
 def xai_yolov8s(image):
-    model = YOLO('yolov8s.pt')  # Ensure the model weights are available
+    model = YOLO('yolov8s.pt')  
     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)
     target_layers = [model.model.model[-2]]  # Adjust to YOLOv8 architecture
     cam_image, renormalized_cam_image = generate_cam_image(model.model, target_layers, tensor, image, boxes)
-    final_image = np.hstack((image, cam_image, renormalized_cam_image))
-    caption = "Results using YOLOv8"
-    return Image.fromarray(final_image), caption
+    
+    # Combine the original, detection, and CAM images for visualization
+    combined_image = np.hstack((
+        cv2.cvtColor(image, cv2.COLOR_RGB2BGR),
+        detections_img,
+        cv2.cvtColor(cam_image, cv2.COLOR_RGB2BGR)
+    ))
+    
+    caption = "Results using YOLOv8 with Class Labels and Bounding Boxes"
+    return Image.fromarray(combined_image), caption
+
+
+# Example Gradio interface
+gr.Interface(
+    fn=xai_yolov8s,
+    inputs=gr.Image(type="numpy"),
+    outputs=[gr.Image(type="pil"), gr.Text()],
+    title="YOLOv8 with XAI and Class Labels"
+).launch()