Spaces:
Sleeping
Sleeping
File size: 2,714 Bytes
978b355 71b8b5d 978b355 672c2bf 37b71af 4f93ba9 383e8f6 71b8b5d 383e8f6 6c34a8c 383e8f6 6c34a8c 383e8f6 6c34a8c 83c6e0c 978b355 71b8b5d 978b355 6c34a8c 978b355 6c34a8c 978b355 6c34a8c 383e8f6 71b8b5d 383e8f6 71b8b5d 978b355 6c34a8c 20700c3 383e8f6 978b355 6c34a8c a83e547 |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import torch
import cv2
import numpy as np
from PIL import Image
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
from ultralytics import YOLO
COLORS = np.random.uniform(0, 255, size=(80, 3))
def parse_detections(detections, model):
boxes, colors, names = [], [], []
for detection in detections.boxes:
xmin, ymin, xmax, ymax = map(int, detection.xyxy[0].tolist())
confidence = detection.conf.item()
if confidence < 0.2:
continue
class_id = int(detection.cls.item())
name = model.names[class_id]
boxes.append((xmin, ymin, xmax, ymax))
colors.append(COLORS[class_id])
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 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_yolov8s(image):
model = YOLO('yolov8s.pt') # Ensure the model weights are available
model.eval()
results = model(image)
detections = results[0]
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)
final_image = np.hstack((image, cam_image, renormalized_cam_image))
caption = "Results using YOLOv8"
return Image.fromarray(final_image), caption
|