Add: DFF support

yolov5.py

@@ -0,0 +1,163 @@
+
+import torch
+import cv2
+import warnings
+warnings.filterwarnings('ignore')
+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
+import yaml
+import requests
+from pytorch_grad_cam import DeepFeatureFactorization
+from pytorch_grad_cam.utils.image import show_cam_on_image, preprocess_image
+from pytorch_grad_cam.utils.image import deprocess_image, show_factorization_on_image
+COLORS = np.random.uniform(0, 255, size=(80, 3))
+
+
+def parse_detections(results):
+    detections = results.pandas().xyxy[0].to_dict()
+    boxes, colors, names, classes = [], [], [], []
+    for i in range(len(detections["xmin"])):
+        confidence = detections["confidence"][i]
+        if confidence < 0.2:
+            continue
+        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)
+        classes.append(category) 
+    return boxes, colors, names, classes
+
+
+def draw_detections(boxes, colors, names, classes, img):
+    for box, color, name, cls in zip(boxes, colors, names, classes):
+        xmin, ymin, xmax, ymax = box
+        label = f"{cls}: {name}"  # Combine class ID and name
+        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, 2)
+        cv2.putText(
+            img, label, (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
+    cam_image = show_cam_on_image(img_float, grayscale_cam, use_rgb=True)
+    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_yolov5(image):
+    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
+    model.eval()
+    model.cpu()
+
+    target_layers = [model.model.model.model[-2]]  # Grad-CAM target layer
+
+    # Run YOLO detection
+    results = model([image])
+    boxes, colors, names, classes = parse_detections(results) 
+    detections_img = draw_detections(boxes, colors, names,classes,  image.copy())
+
+    # 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, detections_img, renormalized_cam_image))
+    caption = "Results using YOLOv5"
+    return Image.fromarray(final_image), caption
+
+
+# Check if CUDA is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+mean = [0.485, 0.456, 0.406]  # Mean for RGB channels
+std = [0.229, 0.224, 0.225]   # Standard deviation for RGB channels
+# Load YOLOv5 model and move it to the appropriate device
+model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device)
+print(f"Loaded YOLOv5 model on {device}")
+def create_labels(concept_scores, top_k=2):
+    """Create a list with the category names of the top scoring categories."""
+    yolov5_categories_url = \
+        "https://github.com/ultralytics/yolov5/raw/master/data/coco128.yaml"  # URL to the YOLOv5 categories file
+    yaml_data = requests.get(yolov5_categories_url).text
+    labels = yaml.safe_load(yaml_data)['names']  # Parse the YAML file to get class names
+    
+    concept_categories = np.argsort(concept_scores, axis=1)[:, ::-1][:, :top_k]
+    concept_labels_topk = []
+    for concept_index in range(concept_categories.shape[0]):
+        categories = concept_categories[concept_index, :]    
+        concept_labels = []
+        for category in categories:
+            score = concept_scores[concept_index, category]
+            label = f"{labels[category]}:{score:.2f}"
+            concept_labels.append(label)
+        concept_labels_topk.append("\n".join(concept_labels))
+    return concept_labels_topk
+
+def get_image_from_url(url, device):
+    """A function that gets a URL of an image,
+    and returns a numpy image and a preprocessed
+    torch tensor ready to pass to the model"""
+
+    img = np.array(Image.open("/home/drovco/Bhumika/NeuralVista/data/xai/sample1.jpeg"))
+    img = cv2.resize(img, (640, 640))
+    rgb_img_float = np.float32(img) /255.0
+    input_tensor = torch.from_numpy(rgb_img_float).permute(2, 0, 1).unsqueeze(0).to(device)
+    return img, rgb_img_float, input_tensor
+
+def visualize_image(model, img_url, n_components=20, top_k=1, lyr_idx = 2):
+    img, rgb_img_float, input_tensor = get_image_from_url(img_url, device)
+    
+    # Specify the target layer for DeepFeatureFactorization (e.g., YOLO's backbone)
+    target_layer = model.model.model.model[-lyr_idx]  # Select a feature extraction layer
+    
+    dff = DeepFeatureFactorization(model=model.model, target_layer=target_layer)
+    
+    # Run DFF on the input tensor
+    concepts, batch_explanations = dff(input_tensor, n_components)
+    
+    # Softmax normalization
+    concept_outputs = torch.softmax(torch.from_numpy(concepts), axis=-1).numpy()    
+    concept_label_strings = create_labels(concept_outputs, top_k=top_k)
+    
+    # Visualize explanations
+    visualization = show_factorization_on_image(rgb_img_float, 
+                                                batch_explanations[0],
+                                                image_weight=0.2,
+                                                concept_labels=concept_label_strings)
+    
+    import matplotlib.pyplot as plt
+    plt.imshow(visualization)
+    plt.savefig("test" + str(lyr_idx) + ".png")
+    result = np.hstack((img, visualization))
+
+    
+    # Resize for visualization
+    if result.shape[0] > 500:
+        result = cv2.resize(result, (result.shape[1]//4, result.shape[0]//4))
+    
+    return result
+
+# Test with images
+for indx in range(2,12):
+    Image.fromarray(visualize_image(model, 
+                                        "https://github.com/jacobgil/pytorch-grad-cam/blob/master/examples/both.png?raw=true", lyr_idx = indx))