Update: class lables

yolov5.py

@@ -7,10 +7,13 @@ 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, classes = [], [], [], []
+    boxes, colors, names = [], [], []
     for i in range(len(detections["xmin"])):
         confidence = detections["confidence"][i]
         if confidence < 0.2:
@@ -21,8 +24,7 @@ def parse_detections(results):
         boxes.append((xmin, ymin, xmax, ymax))
         colors.append(COLORS[category])
         names.append(name)
-        classes.append(category)  # Add class ID
-    return boxes, colors, names, classes
+    return boxes, colors, names
 
 
 def draw_detections(boxes, colors, names, classes, img):
@@ -42,12 +44,17 @@ 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
 
 
@@ -61,8 +68,8 @@ def xai_yolov5(image):
 
     # Run YOLO detection
     results = model([image])
-    boxes, colors, names, classes = parse_detections(results)  # Extract classes
-    detections_img = draw_detections(boxes, colors, names, classes, image.copy())
+    boxes, colors, names = parse_detections(results)
+    detections_img = draw_detections(boxes, colors, names, image.copy())
 
     # Prepare input tensor for Grad-CAM
     img_float = np.float32(image) / 255
@@ -73,19 +80,6 @@ def xai_yolov5(image):
     cam_image, renormalized_cam_image = generate_cam_image(model, target_layers, tensor, image, boxes)
 
     # Combine results
-    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()
+    final_image = np.hstack((image, cam_image, renormalized_cam_image))
+    caption = "Results using YOLOv5"
+    return Image.fromarray(final_image), caption

yolov8.py

@@ -10,9 +10,8 @@ from ultralytics import YOLO
 
 
 COLORS = np.random.uniform(0, 255, size=(80, 3))
-
 def parse_detections(detections, model):
-    boxes, colors, names, classes = [], [], [], []
+    boxes, colors, names = [], [], []
     for detection in detections.boxes:
         xmin, ymin, xmax, ymax = map(int, detection.xyxy[0].tolist())
         confidence = detection.conf.item()
@@ -23,8 +22,7 @@ def parse_detections(detections, model):
         boxes.append((xmin, ymin, xmax, ymax))
         colors.append(COLORS[class_id])
         names.append(name)
-        classes.append(class_id)  # Add class ID
-    return boxes, colors, names, classes
+    return boxes, colors, names
 
 def draw_detections(boxes, colors, names, classes, img):
     for box, color, name, cls in zip(boxes, colors, names, classes):
@@ -53,33 +51,17 @@ 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')  
+    model = YOLO('yolov8s.pt')  # Ensure the model weights are available
     model.eval()
     results = model(image)
     detections = results[0]
-    boxes, colors, names, classes = parse_detections(detections, model)
-    detections_img = draw_detections(boxes, colors, names, classes, image.copy())
+    boxes, colors, names = parse_detections(detections, model)
+    detections_img = draw_detections(boxes, colors, names, 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)
-    
-    # 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()
+    final_image = np.hstack((image, cam_image, renormalized_cam_image))
+    caption = "Results using YOLOv8"
+    return Image.fromarray(final_image), caption