added dff

yolov5.py

@@ -82,90 +82,166 @@ def xai_yolov5(image):
     return Image.fromarray(final_image), caption
 
 
-"""
-import yaml
-import torch
-import warnings
-warnings.filterwarnings('ignore')
-from PIL import Image
+
 import numpy as np
-import requests
-import cv2
+from PIL import Image
 import torch
-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
+import cv2
+from typing import Callable, List, Tuple, Optional
+from sklearn.decomposition import NMF
+from pytorch_grad_cam.activations_and_gradients import ActivationsAndGradients
+from pytorch_grad_cam.utils.image import scale_cam_image, create_labels_legend, show_factorization_on_image
+import matplotlib.pyplot as plt
+from pytorch_grad_cam.utils.image import show_factorization_on_image
+import requests    
+import yaml
+import matplotlib.patches as patches
 
-# 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):
-    
-    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):
- 
-
-    img = np.array(Image.open(os.path.join(os.getcwd(), "data/xai/sample1.jpeg")))
+def dff_l(activations,  model,  n_components):
+    batch_size, channels, h, w = activations.shape
+    print('activation', activations.shape)
+    target_layer_index = 4        
+    reshaped_activations = activations.transpose((1, 0, 2, 3))
+    reshaped_activations[np.isnan(reshaped_activations)] = 0
+    reshaped_activations = reshaped_activations.reshape(
+        reshaped_activations.shape[0], -1)
+    offset = reshaped_activations.min(axis=-1)
+    reshaped_activations = reshaped_activations - offset[:, None]
+    model = NMF(n_components=n_components, init='random', random_state=0)
+    W = model.fit_transform(reshaped_activations)
+    H = model.components_
+    concepts = W + offset[:, None]
+    explanations = H.reshape(n_components, batch_size, h, w)
+    explanations = explanations.transpose((1, 0, 2, 3))
+    return concepts, explanations
+
+class DeepFeatureFactorization:
+    def __init__(self,
+                 model: torch.nn.Module,
+                 target_layer: torch.nn.Module,
+                 reshape_transform: Callable = None,
+                 computation_on_concepts=None
+                 ):
+        self.model = model
+        self.computation_on_concepts = computation_on_concepts
+        self.activations_and_grads = ActivationsAndGradients(
+            self.model, [target_layer], reshape_transform)
+
+    def __call__(self,
+                 input_tensor: torch.Tensor,
+                 model: torch.nn.Module,
+                 n_components: int = 16):
+        if isinstance(input_tensor, np.ndarray):
+            input_tensor = torch.from_numpy(input_tensor)  # Convert NumPy array 
+
+        batch_size, channels, h, w = input_tensor.size()
+        _ = self.activations_and_grads(input_tensor)
+
+        with torch.no_grad():
+            activations = self.activations_and_grads.activations[0].cpu(
+            ).numpy()
+
+        concepts, explanations = dff_l(activations, model,  n_components=n_components)
+        processed_explanations = []
+
+        for batch in explanations:
+            processed_explanations.append(scale_cam_image(batch, (w, h)))
+
+        if self.computation_on_concepts:
+            with torch.no_grad():
+                concept_tensors = torch.from_numpy(
+                    np.float32(concepts).transpose((1, 0)))
+                concept_outputs = self.computation_on_concepts(
+                    concept_tensors).cpu().numpy()
+            return concepts, processed_explanations, concept_outputs
+        else:
+            return concepts, processed_explanations,  explanations
+
+    def __del__(self):
+        self.activations_and_grads.release()
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        self.activations_and_grads.release()
+        if isinstance(exc_value, IndexError):
+            # Handle IndexError here...
+            print(
+                f"An exception occurred in ActivationSummary with block: {exc_type}. Message: {exc_value}")
+            return True
+
+
+def dff_nmf(image, target_lyr, n_components):
+    mean = [0.485, 0.456, 0.406]  # Mean for RGB channels
+    std = [0.229, 0.224, 0.225]   # Standard deviation for RGB channels
+    img, rgb_img_float, input_tensor = image.to(device)
     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
+    rgb_img_float = np.float32(img) / 255.0
+    input_tensor = torch.from_numpy(rgb_img_float).permute(2, 0, 1).unsqueeze(0).to(device) 
 
-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)
+    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device)
+    dff= DeepFeatureFactorization(model=model,    
+                               target_layer=model.model.model.model[int(target_lyr)],    
+                               computation_on_concepts=None)
     
-    import matplotlib.pyplot as plt
-    plt.imshow(visualization)
-    plt.savefig("test" + str(lyr_idx) + ".png")
-    result = np.hstack((img, visualization))
+    concepts, batch_explanations, explanations = dff(input_tensor, model, n_components)
 
+
+    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
+    num_classes = model.model.model.model[-1].nc 
+
+    for indx in range( explanations[0].shape[0]):
+        upsampled_input =  explanations[0][indx]
+        upsampled_input = torch.tensor(upsampled_input)
+        device = next(model.parameters()).device
+        input_tensor = upsampled_input.unsqueeze(0)
+        input_tensor = input_tensor.unsqueeze(1).repeat(1, 128, 1, 1)    
+        detection_lyr = model.model.model.model[-1]
+        output1 = detection_lyr.m[0](input_tensor.to(device))
+        objectness = output1[..., 4]  # Objectness score (index 4)
+        class_scores = output1[..., 5:]  # Class scores (from index 5 onwards, representing 80 classes)
+        objectness = torch.sigmoid(objectness)
+        class_scores = torch.sigmoid(class_scores)
+        confidence_mask = objectness > 0.5
+        objectness = objectness[confidence_mask]
+        class_scores = class_scores[confidence_mask]
+        scores, class_ids = class_scores.max(dim=-1)  # Get max class score per cell
+        scores = scores * objectness  # Adjust scores by objectness
+        boxes = output1[..., :4]  # First 4 values are x1, y1, x2, y2
+        boxes = boxes[confidence_mask]  # Filter boxes by confidence mask
+        fig, ax = plt.subplots(1, figsize=(10, 10))
+        ax.imshow(torch.tensor(batch_explanations[0][indx]).cpu().numpy(), cmap="gray")  # Display image
+        top_score_idx = scores.argmax(dim=0)  # Get the index of the max score
+        top_score = scores[top_score_idx].item()
+        top_class_id = class_ids[top_score_idx].item()
+        top_box = boxes[top_score_idx].cpu().numpy()
+        scale_factor = 16
+        x1, y1, x2, y2 = top_box
+        x1, y1, x2, y2 = x1 * scale_factor, y1 * scale_factor, x2 * scale_factor, y2 * scale_factor
+        rect = patches.Rectangle(
+                (x1, y1), x2 - x1, y2 - y1,
+                linewidth=2, edgecolor='r', facecolor='none')
+        ax.add_patch(rect)
     
-    # Resize for visualization
-    if result.shape[0] > 500:
-        result = cv2.resize(result, (result.shape[1]//4, result.shape[0]//4))
-    
-    return result
+        predicted_label = labels[top_class_id]  # Map ID to label
+        ax.text(x1, y1, f"{predicted_label}: {top_score:.2f}",
+            color='r', fontsize=12, verticalalignment='top')
+        plt.show()
+        plt.savefig("test_" + str(indx) + ".png" )
+        plt.clf()
+        return rgb_img_float,  explanations
+
+
+def visualize_batch_explanations(rgb_img_float, batch_explanations, image_weight=0.7):
+    for i, explanation in enumerate(batch_explanations):
+        # Create visualization for each explanation
+        visualization = show_factorization_on_image(rgb_img_float, explanation, image_weight=image_weight)
+        plt.figure()
+        plt.imshow(visualization)  # Correctly pass the visualization data
+        plt.title(f'Explanation {i + 1}')  # Set the title for each plot
+        plt.axis('off')  # Hide axes
+        plt.show()  # Show the plot
+    plt.savefig("test_w.png")
 
-# 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))
-"""