Spaces:
Sleeping
Sleeping
| 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 | |