Testing

SkinGPT.py

@@ -41,6 +41,7 @@ class Blip2QFormer(nn.Module):
         )
         self.vision_proj = nn.Linear(vision_width, self.bert_config.hidden_size)
         self._init_weights()
+        self.last_attention = None
 
     def _init_weights(self):
         nn.init.normal_(self.query_tokens, std=0.02)
@@ -71,6 +72,7 @@ class Blip2QFormer(nn.Module):
             output_attentions=True,
             return_dict=True
         )
+        self.last_attention = outputs.attentions[-1]
         return outputs.last_hidden_state[:, :self.num_query_tokens]
 
 

test.py

@@ -27,37 +27,46 @@ class SkinGPTTester:
         with torch.no_grad():
             # Get attention maps
             _ = self.classifier.model.encode_image(image_tensor)
+            
+            # Get attention from Q-Former
+            if self.classifier.model.q_former.last_attention is None:
+                print("Warning: No attention maps available. Make sure output_attentions=True in BERT config.")
+                return
+            
+            # Get the last layer's attention
             attention = self.classifier.model.q_former.last_attention[0].mean(dim=0)
-        
-        # Reshape attention to image size
-        h = w = int(math.sqrt(attention.shape[1]))
-        attention = attention.reshape(h, w)
-        
-        # Plot
-        plt.figure(figsize=(15, 5))
-        
-        # Original image
-        plt.subplot(1, 3, 1)
-        plt.imshow(image)
-        plt.title('Original Image')
-        plt.axis('off')
-        
-        # Attention map
-        plt.subplot(1, 3, 2)
-        plt.imshow(attention, cmap='hot')
-        plt.title('Attention Map')
-        plt.axis('off')
-        
-        # Overlay
-        plt.subplot(1, 3, 3)
-        plt.imshow(image)
-        plt.imshow(attention, alpha=0.5, cmap='hot')
-        plt.title('Attention Overlay')
-        plt.axis('off')
-        
-        plt.tight_layout()
-        plt.savefig('attention_visualization.png')
-        plt.close()
+            
+            # Reshape attention to image size
+            h = w = int(math.sqrt(attention.shape[1]))
+            attention = attention.reshape(h, w)
+            
+            # Plot
+            plt.figure(figsize=(15, 5))
+            
+            # Original image
+            plt.subplot(1, 3, 1)
+            plt.imshow(image)
+            plt.title('Original Image')
+            plt.axis('off')
+            
+            # Attention map
+            plt.subplot(1, 3, 2)
+            plt.imshow(attention, cmap='hot')
+            plt.title('Attention Map')
+            plt.axis('off')
+            
+            # Overlay
+            plt.subplot(1, 3, 3)
+            plt.imshow(image)
+            plt.imshow(attention, alpha=0.5, cmap='hot')
+            plt.title('Attention Overlay')
+            plt.axis('off')
+            
+            plt.tight_layout()
+            plt.savefig('attention_visualization.png')
+            plt.close()
+            
+            print(f"Attention visualization saved as 'attention_visualization.png'")
 
     def check_feature_similarity(self, image_path1, image_path2):
         """Compare embeddings of two images"""