Debug: yolov8 target lyr

app.py

@@ -26,7 +26,7 @@ interface = gr.Interface(
     inputs=[
         gr.Image(type="pil", label="Upload an Image"),
         gr.CheckboxGroup(
-            choices=["yolov5", "yolov8", "yolov10"],
+            choices=["yolov5", "yolov8n", "yolov10"],
             value=["yolov5"],  # Set default selection to YOLOv5
             label="Select Model(s)",
         )

yolov8.py

@@ -1,29 +1,32 @@
 from ultralytics import YOLO
+import torch
 import cv2
 import numpy as np
 from PIL import Image
-import torch
-from torchcam.methods import GradCAM
-from torchcam.utils import overlay_mask
+import torchvision.transforms as transforms
+from pytorch_grad_cam import EigenCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
+import gradio as gr
 
-# Set random colors for detection bounding boxes
+# Global Color Palette
 COLORS = np.random.uniform(0, 255, size=(80, 3))
 
-def parse_detections_yolov8(results):
+def parse_detections(results):
+    detections = results.pandas().xyxy[0].to_dict()
     boxes, colors, names = [], [], []
-    detections = results.boxes
-    for box in detections:
-        confidence = box.conf[0].item()
-        if confidence < 0.2:  # Filter out low-confidence detections
+    for i in range(len(detections["xmin"])):
+        confidence = detections["confidence"][i]
+        if confidence < 0.2:
             continue
-        xmin, ymin, xmax, ymax = map(int, box.xyxy[0].tolist())
-        category = int(box.cls[0].item())
-        name = results.names[category]
+        xmin, ymin = int(detections["xmin"][i]), int(detections["ymin"][i])
+        xmax, ymax = int(detections["xmax"][i]), int(detections["ymax"][i])
+        name, category = detections["name"][i], int(detections["class"][i])
         boxes.append((xmin, ymin, xmax, ymax))
         colors.append(COLORS[category])
         names.append(name)
     return boxes, colors, names
 
+
 def draw_detections(boxes, colors, names, img):
     for box, color, name in zip(boxes, colors, names):
         xmin, ymin, xmax, ymax = box
@@ -33,35 +36,47 @@ def draw_detections(boxes, colors, names, img):
                     lineType=cv2.LINE_AA)
     return img
 
-def xai_yolov8(image):
-    # Load YOLOv8 model
-    model = YOLO('yolov8n.pt')  # Load YOLOv8 nano model
-    model.to('cpu')
+
+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
+
+
+def xai_yolov8n(image):
+    # Load YOLOv8n model
+    model = YOLO('yolov8n.pt')  # Load YOLOv8n pre-trained weights
     model.eval()
+    model.cpu()
+
+    target_layers = [model.model.model[-2]]  # Grad-CAM target layer
 
     # Run YOLO detection
-    results = model(image)
-    boxes, colors, names = parse_detections_yolov8(results[0])
+    results = model([image])
+    boxes, colors, names = parse_detections(results)
     detections_img = draw_detections(boxes, colors, names, image.copy())
 
-    # Convert image to PyTorch tensor for Grad-CAM
-    image_tensor = torch.tensor(np.array(image)).permute(2, 0, 1).unsqueeze(0).float() / 255.0
-    image_tensor = image_tensor.to('cpu')
-    model(image_tensor)  # Forward pass
-    print(model.model)  # Output model layers to find the target layer
-    target_layer = model.model.model[-1]  # Target the last detection layer
-    
-    # Initialize Grad-CAM with the target layer
-    grad_cam = GradCAM(model.model, target_layer=target_layer)
-
-    # Perform Grad-CAM
-    cam_map = grad_cam(image_tensor)
-
-    # Overlay Grad-CAM mask onto original image
-    cam_image = overlay_mask(image, cam_map.squeeze(0).cpu().numpy(), alpha=0.5)
-
-    # Combine original image and Grad-CAM image
-    final_image = np.hstack((np.array(image), cam_image))
-    caption = "Results using YOLOv8 and Grad-CAM via torchcam"
-    
+    # Prepare input tensor for Grad-CAM
+    img_float = np.float32(image) / 255
+    transform = transforms.ToTensor()
+    tensor = transform(img_float).unsqueeze(0)
+
+    # Grad-CAM visualization
+    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 YOLOv8n"
     return Image.fromarray(final_image), caption