fix: scaling based on preprocessor instead of grid estimation and padding

similarity.py

@@ -189,7 +189,7 @@ class JinaV4SimilarityMapper:
         print(f"Token map: {token_map}")
         return tokens, query_embeddings, token_map
 
-    def process_image(self, image: Union[str, bytes, Image.Image]) -> Tuple[Image.Image, torch.Tensor, Tuple[int, int]]:
+    def process_image(self, image: Union[str, bytes, Image.Image]) -> Tuple[Image.Image, torch.Tensor, Tuple[int, int], Tuple[int, int]]:
         """
         Process image to get patch embeddings in multivector format.
         
@@ -200,34 +200,34 @@ class JinaV4SimilarityMapper:
             pil_image: Original PIL image.
             patch_embeddings: Image patch embeddings [num_patches/num_vectors, embed_dim].
             size: Original image size (width, height).
+            grid_size: Patch grid dimensions (height, width) after merge.
         """
         pil_image = self._load_image(image)
-
         proc_out = self.preprocessor.process_images(images=[pil_image])
-        for key, value in proc_out.items():
-            if isinstance(value, torch.Tensor):
-                print(f"proc out {key} shape: {value.shape}")
-                
+        
+        # Get the grid dimensions from preprocessor
+        image_grid_thw = proc_out["image_grid_thw"]
+        _, height, width = image_grid_thw[0].tolist()
+        # Account for 2x2 merge
+        grid_height = height // 2
+        grid_width = width // 2
+        
         size = pil_image.size
         image_embeddings = self.model.encode_image(
             images=[pil_image],
             task="retrieval",
             return_multivector=True,
-            max_pixels=1024*1024,  # Max image resolution
+            max_pixels=1024*1024,
             truncate_dim=self.num_vectors
         )
-        image_embeddings = image_embeddings[0] # [num_patches/num_vectors, embed_dim]
-        non_zero_mask = (image_embeddings.abs().sum(dim=1) > 0)
-        image_embeddings = image_embeddings[non_zero_mask]
-
-        # <|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n
-        vision_start_position_from_start = 3 + 1
-        vision_end_position_from_end = 6 + 1
+        image_embeddings = image_embeddings[0]
+        
         # Remove special tokens
+        vision_start_position_from_start = 5
+        vision_end_position_from_end = 6
         image_embeddings = image_embeddings[vision_start_position_from_start:-vision_end_position_from_end]
-        print(f"Image embeddings shape: {image_embeddings.shape}")
-        print(f"Image size: {size}")
-        return pil_image, image_embeddings, size
+        
+        return pil_image, image_embeddings, size, (grid_height, grid_width)
 
     def _load_image(self, image: Union[str, bytes, Image.Image]) -> Image.Image:
         """Load image from various formats (URL, path, bytes, PIL Image)."""
@@ -273,45 +273,37 @@ class JinaV4SimilarityMapper:
         similarity_scores = torch.cosine_similarity(token_expanded, patch_embeddings, dim=1)
         return similarity_scores
 
-    def generate_heatmap(self, image: Image.Image, similarity_map: torch.Tensor, size: Tuple[int, int]) -> str:
+    def generate_heatmap(self, image: Image.Image, similarity_map: torch.Tensor, size: Tuple[int, int], grid_size: Tuple[int, int]) -> str:
         """
         Generate a heatmap overlay on the image and return as base64.
         
         Args:
             image: Original PIL image.
-            similarity_map: Similarity scores [num_vectors/num_patches].
+            similarity_map: Similarity scores [num_patches].
             size: Original image size (width, height).
-            
-        Returns:
-            Base64-encoded PNG image with heatmap.
+            grid_size: Patch grid dimensions (height, width).
         """
-        num_patches = similarity_map.shape[0]
+        # num_patches = similarity_map.shape[0]
+        grid_height, grid_width = grid_size
+        
         # Normalize to [0, 1]
         similarity_map = (similarity_map - similarity_map.min()) / (
             similarity_map.max() - similarity_map.min() + 1e-8
         )
-        # Calculate grid dimensions from image aspect ratio
-        width, height = size
-        aspect_ratio = width / height
-        grid_width = int(np.ceil(np.sqrt(num_patches * aspect_ratio)))
-        grid_height = int(np.ceil(num_patches / grid_width))
-        total_patches = grid_width * grid_height
-        # Ensure similarity map fits grid (padding/truncation)
-        if num_patches < total_patches:
-            padding = torch.zeros(total_patches - num_patches, device=similarity_map.device)
-            similarity_map = torch.cat([similarity_map, padding])
-        else:
-            similarity_map = similarity_map[:total_patches]
-        # Reshape to 2D grid [grid_height, grid_width]
+        
+        # Reshape to 2D grid
         similarity_2d = similarity_map.reshape(grid_height, grid_width).cpu().numpy()
+        
         # Create & resize heatmap
         heatmap = (self.colormap(similarity_2d) * 255).astype(np.uint8)
         heatmap = Image.fromarray(heatmap[..., :3], mode="RGB")
         heatmap = heatmap.resize(size, resample=Image.BICUBIC)
+        
         # Blend with original image
         original_rgba = image.convert("RGBA")
         heatmap_rgba = heatmap.convert("RGBA")
         blended = Image.blend(original_rgba, heatmap_rgba, alpha=self.heatmap_alpha)
+        
         # Encode to base64
         buffer = BytesIO()
         blended.save(buffer, format="PNG")
@@ -325,30 +317,22 @@ class JinaV4SimilarityMapper:
     ) -> Tuple[List[str], Dict[str, str]]:
         """
         Main method to generate similarity maps for all query tokens.
-        
-        Args:
-            query: Input query text.
-            image: Image to analyze.
-            aggregation: How to aggregate multivector similarities.
-            
-        Returns:
-            tokens: List of query tokens.
-            heatmaps: Dictionary of {token: base64_heatmap}.
         """
-        tokens, query_embeddings, token_map = self.process_query(query)
-        pil_image, patch_embeddings, size = self.process_image(image)
+        _, query_embeddings, token_map = self.process_query(query)
+        pil_image, patch_embeddings, size, grid_size = self.process_image(image)
+        
         heatmaps = {}
         tokens_for_ui = []
+        
         for idx, token in token_map.items():
-            print(f"Processing token: {token} (index {idx})")
             if self._should_filter_token(token):
                 continue   
             tokens_for_ui.append(token)
-            token_embedding = query_embeddings[idx]  # [embed_dim]
+            token_embedding = query_embeddings[idx]
             sim_map = self.compute_similarity_map(
                 token_embedding, patch_embeddings, aggregation
             )
-            heatmap_b64 = self.generate_heatmap(pil_image, sim_map, size)
+            heatmap_b64 = self.generate_heatmap(pil_image, sim_map, size, grid_size)
             heatmaps[token] = heatmap_b64
         
         return tokens_for_ui, heatmaps