update: dff_nmf

yolov8.py

@@ -162,7 +162,6 @@ class DeepFeatureFactorization:
 
 
 
-
 def dff_nmf(image, target_lyr, n_components):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     mean = [0.485, 0.456, 0.406]  # Mean for RGB channels
@@ -171,77 +170,52 @@ def dff_nmf(image, target_lyr, n_components):
     rgb_img_float = np.float32(img) / 255.0
     input_tensor = torch.from_numpy(rgb_img_float).permute(2, 0, 1).unsqueeze(0).to(device) 
 
-    model = YOLO('yolov8s.pt')
-    dff= DeepFeatureFactorization(model=model,    
-                               target_layer=model.model.model[int(target_lyr)],    
-                               computation_on_concepts=None)
+    model = YOLO('yolov8s.pt')  # Ensure the model is loaded correctly
+    dff = DeepFeatureFactorization(model=model,    
+                                   target_layer=model.model.model[int(target_lyr)],    
+                                   computation_on_concepts=None)
     
     concepts, batch_explanations, explanations = dff(input_tensor, model, n_components)
 
+    # Getting predictions directly from YOLO
+    with torch.no_grad():
+        results = model(input_tensor)
+
+    # Post-processing to extract detections
+    boxes, scores, classes = results.xywh[0][:, :4], results.xywh[0][:, 4], results.xywh[0][:, 5]
+    boxes = boxes.cpu().numpy()
+    scores = scores.cpu().numpy()
+    classes = classes.cpu().numpy()
 
-    #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[-1].nc 
-    results = []
-    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[-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=(8, 8))
-        ax.axis("off")
-        ax.imshow(torch.tensor(batch_explanations[0][indx]).cpu().numpy(), cmap="plasma")  # 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')
+    # Filter detections with confidence score > threshold (e.g., 0.5)
+    high_conf_boxes = boxes[scores > 0.5]
+    high_conf_classes = classes[scores > 0.5]
+    
+    # Use the processed detections for visualization and further tasks
+    # Example visualization and output processing
+    fig, ax = plt.subplots(1, figsize=(8, 8))
+    ax.axis("off")
+    ax.imshow(rgb_img_float)
+
+    for box, cls in zip(high_conf_boxes, high_conf_classes):
+        x1, y1, x2, y2 = box
+        rect = patches.Rectangle((x1, y1), x2 - x1, y2 - y1,
+                                 linewidth=2, edgecolor='r', facecolor='none')
         ax.add_patch(rect)
+        ax.text(x1, y1, f"Class {cls}", color='r', fontsize=12, verticalalignment='top')
+
+    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
+    fig.canvas.draw()
+    image_array = np.array(fig.canvas.renderer.buffer_rgba())
+    image_resized = cv2.resize(image_array, (640, 640))
+    rgba_channels = cv2.split(image_resized)
+    alpha_channel = rgba_channels[3] 
+    rgb_channels = np.stack(rgba_channels[:3], axis=-1)
     
-        #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.subplots_adjust(left=0, right=1, top=1, bottom=0)
-
-        fig.canvas.draw()  # Draw the canvas to make sure the image is rendered
-        image_array = np.array(fig.canvas.renderer.buffer_rgba())  # Convert to numpy array
-        print("____________image_arrya", image_array.shape)
-        image_resized = cv2.resize(image_array, (640, 640))
-        rgba_channels = cv2.split(image_resized)
-        alpha_channel = rgba_channels[3] 
-        rgb_channels = np.stack(rgba_channels[:3], axis=-1)
-        #overlay_img = (alpha_channel[..., None] * image) + ((1 - alpha_channel[..., None]) * rgb_channels)
-        
-        #temp = image_array.reshape((rgb_img_float.shape[0],rgb_img_float.shape[1]) )
-        #visualization = show_factorization_on_image(rgb_img_float, image_array.resize((rgb_img_float.shape)) , image_weight=0.3)
-        visualization = show_factorization_on_image(rgb_img_float, np.transpose(rgb_channels, (2, 0, 1)), image_weight=0.3)
-        results.append(visualization)
-        plt.clf()  
-        #return image_array
-
+    visualization = show_factorization_on_image(rgb_img_float, np.transpose(rgb_channels, (2, 0, 1)), image_weight=0.3)
+    
+    return rgb_img_float, batch_explanations, visualization
 
-    return rgb_img_float, batch_explanations, results
 
 
 def visualize_batch_explanations(rgb_img_float, batch_explanations, image_weight=0.7):