File size: 1,920 Bytes
4f93ba9
 
71b8b5d
 
 
 
4f93ba9
71b8b5d
4f93ba9
71b8b5d
 
 
 
 
4f93ba9
71b8b5d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4f93ba9
 
 
71b8b5d
 
 
 
 
4f93ba9
 
71b8b5d
4f93ba9
71b8b5d
4f93ba9
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
from yolov8_explainer import YOLOv8Explainer
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image

# Set random colors for detection bounding boxes
COLORS = np.random.uniform(0, 255, size=(80, 3))

def parse_detections_yolov8(results):
    boxes, colors, names = [], [], []
    detections = results.boxes
    for box in detections:
        confidence = box.conf[0].item()
        if confidence < 0.2:  # Filter out low-confidence detections
            continue
        xmin, ymin, xmax, ymax = map(int, box.xyxy[0].tolist())
        category = int(box.cls[0].item())
        name = results.names[category]
        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
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, 2)
        cv2.putText(img, name, (xmin, ymin - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2,
                    lineType=cv2.LINE_AA)
    return img

def xai_yolov8(image):
    # Load YOLOv8 model
    model = YOLO('yolov8n.pt')  # Load YOLOv8 nano model
    model.to('cpu')
    model.eval()

    # Initialize YOLOv8 Explainer
    explainer = YOLOv8Explainer(model)

    # Run YOLO detection
    results = model(image)
    boxes, colors, names = parse_detections_yolov8(results[0])
    detections_img = draw_detections(boxes, colors, names, image.copy())

    # Generate Grad-CAM visualization
    cam_image, renormalized_cam_image = explainer.visualize(image, results)

    # Combine original, Grad-CAM, and renormalized Grad-CAM images
    final_image = np.hstack((image, cam_image, renormalized_cam_image))
    caption = "Results using YOLOv8 and YOLOv8Explainer"
    return Image.fromarray(final_image), caption