updated calibration logic

app.py

@@ -160,43 +160,49 @@ async def concept_ensemble(file: UploadFile = File(...), user_prompt: str = Quer
     image = Image.open(file.file).convert("RGB")
     blip = MODELS["blip"]
     
-    # Get model's caption
     inputs_gen = blip["processor"](images=image, return_tensors="pt").to(DEVICE)
     with torch.no_grad():
         generated_ids = blip["model"].generate(**inputs_gen, max_length=40)
         model_caption = blip["processor"].decode(generated_ids[0], skip_special_tokens=True)
 
-    # 1. NEW: Localized Keyword Embedding
-    # We focus on the core nouns and adjectives to prevent 'template bias'
-    def get_focused_embedding(text):
+    def get_clean_embedding(text):
         inputs = blip["processor"](text=text, return_tensors="pt", padding=True).to(DEVICE)
         with torch.no_grad():
-            # Get output from the BERT-based text decoder
             outputs = blip["model"].text_decoder.bert(**inputs)
-            # Average hidden states of ALL tokens to capture keyword specifics
             return F.normalize(outputs.last_hidden_state.mean(dim=1), p=2, dim=-1)
 
-    user_embed = get_focused_embedding(user_prompt)
-    model_embed = get_focused_embedding(model_caption)
+    user_embed = get_clean_embedding(user_prompt)
+    model_embed = get_clean_embedding(model_caption)
     
+    # --- MLE TRICK: Word-Level Calibration ---
+    # This prevents 'Pink Cafe' and 'Yellow Sofa' from being 0.99
+    user_words = set(user_prompt.lower().split())
+    model_words = set(model_caption.lower().split())
+    intersection = user_words.intersection(model_words)
+    union = user_words.union(model_words)
+    jaccard_sim = len(intersection) / len(union) if len(union) > 0 else 0
+    
+    # Calculate raw embedding similarity
+    raw_sim = torch.matmul(user_embed, model_embed.T).item()
+    
+    # Weighted Similarity: Combine vector meaning with actual word overlap
+    # This will pull the 0.99 score down if the keywords don't match
+    calibrated_overlap = (raw_sim * 0.4) + (jaccard_sim * 0.6)
+
     # Visual alignment
     with torch.no_grad():
         vision_outputs = blip["model"].vision_model(inputs_gen["pixel_values"])
         image_embed = F.normalize(vision_outputs.last_hidden_state[:, 0, :], p=2, dim=-1)
-
-    # 2. Calculate Corrected Scores
-    sim_image_user = torch.matmul(image_embed, user_embed.T).item()
-    sim_image_model = torch.matmul(image_embed, model_embed.T).item()
-    sim_user_model = torch.matmul(user_embed, model_embed.T).item()
+        sim_image_user = torch.matmul(image_embed, user_embed.T).item()
 
     return {
         "captions": {"user": user_prompt, "model": model_caption},
         "similarity_scores": {
-            "visual_alignment_user": round(sim_image_user, 4),
-            "visual_alignment_model": round(sim_image_model, 4),
-            "semantic_overlap": round(sim_user_model, 4)
+            "semantic_overlap": round(calibrated_overlap, 4),
+            "visual_alignment": round(sim_image_user, 4),
+            "word_match_penalty": round(1 - jaccard_sim, 2)
         },
-        "interpretation": "Strong Agreement" if sim_user_model > 0.8 else "Perspective Divergence"
+        "interpretation": "Perspective Divergence" if calibrated_overlap < 0.6 else "Strong Agreement"
     }
 
     
@@ -206,32 +212,31 @@ async def get_saliency_heatmap(file: UploadFile = File(...), query_text: str = Q
     orig_img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
     
     blip = MODELS["blip"]
-    # We enable 'output_attentions' to grab the internal map directly
+    # We must explicitly call the vision_model to get the attentions cleanly
     inputs = blip["processor"](images=orig_img, text=query_text, return_tensors="pt").to(DEVICE)
     
     with torch.no_grad():
-        outputs = blip["model"](**inputs, output_attentions=True)
-        # Use the last layer of vision encoder self-attention
-        # Shape: (batch, heads, patches, patches)
-        attentions = outputs.vision_model_output.attentions[-1]
+        # Get vision outputs specifically to access the self-attention maps
+        vision_outputs = blip["model"].vision_model(
+            pixel_values=inputs.pixel_values, 
+            output_attentions=True
+        )
+        
+        # Access attentions from the vision model output
+        # Shape: (layers, batch, heads, patches, patches)
+        attentions = vision_outputs.attentions[-1] 
         
-        # Average across heads and take the attention from the [CLS] token to all patches
-        # Patch size for BLIP is typically 14x14 or 16x16
+        # Grid size (usually 16x16 for BLIP)
         grid_size = int(np.sqrt(attentions.shape[-1] - 1))
-        # Remove [CLS] token and reshape to grid
+        # Take attention from the [CLS] token (index 0) to all other patches
         mask = attentions[0, :, 0, 1:].mean(0).view(grid_size, grid_size).cpu().numpy()
 
-    # 1. Normalize and Upscale
+    # Normalize, upscale, and blur for that "Pinterest-chic" glow
     mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
     mask_pill = Image.fromarray((mask * 255).astype('uint8')).resize(orig_img.size, resample=Image.BICUBIC)
+    mask_pill = mask_pill.filter(ImageFilter.GaussianBlur(radius=12))
     
-    # 2. Apply Gaussian Glow for XAI Aesthetic
-    mask_pill = mask_pill.filter(ImageFilter.GaussianBlur(radius=15))
-    mask_final = np.array(mask_pill) / 255.0
-
-    # 3. Apply Colormap and Blend
-    cm = plt.get_cmap('jet')
-    heatmap_rgba = cm(mask_final)
+    heatmap_rgba = plt.get_cmap('jet')(np.array(mask_pill)/255.0)
     heatmap_img = Image.fromarray((heatmap_rgba[:, :, :3] * 255).astype('uint8')).convert("RGB")
     
     blended_img = Image.blend(orig_img, heatmap_img, alpha=0.5)