Add legend to prediction overlay with dynamic font sizing and color boxes

inference_utils/model.py

@@ -14,7 +14,7 @@ from dataclasses import dataclass
 import os
 from typing import Tuple
 
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageFont
 from huggingface_hub import hf_hub_download
 import matplotlib.pyplot as plt
 import numpy as np
@@ -50,7 +50,10 @@ class Model:
             self._model, image, modality_type, targets
         )
         targets_not_found_str = (
-            "\n".join(t.target for t in prediction_targets_not_found)
+            "\n".join(
+                f"{t.target} ({t.adjusted_p_value, 2:.3f})"
+                for t in prediction_targets_not_found
+            )
             if prediction_targets_not_found
             else "All targets were found!"
         )
@@ -104,51 +107,16 @@ def predict(
         )
         pt.adjusted_p_value = float(adj_p_value)
 
-    pred_tasks_found, pred_tasks_not_found = segregate_prediction_tasks(
+    pred_targets_found, pred_tasks_not_found = segregate_prediction_tasks(
         prediction_tasks, 0.05
     )
 
     # Generate visualization
-    colors = generate_colors(len(pred_tasks_found))
-    masks = [1 * (pred_mask[i] > 0.5) for i in range(len(pred_tasks_found))]
+    colors = generate_colors(len(pred_targets_found))
+    masks = [1 * (pred_mask[i] > 0.5) for i in range(len(pred_targets_found))]
     pred_overlay = overlay_masks(image, masks, colors)
 
-    # Add legend
-    if len(pred_tasks_found) > 0:
-        # Convert to numpy for manipulation
-        pred_overlay = np.array(pred_overlay)
-
-        # Calculate legend dimensions
-        legend_height = 30 * len(pred_tasks_found)  # 30 pixels per entry
-        legend_padding = 10  # padding around legend
-        total_height = pred_overlay.shape[0] + legend_height + 2 * legend_padding
-
-        # Create new image with space for legend
-        new_image = np.zeros((total_height, pred_overlay.shape[1], 3), dtype=np.uint8)
-        new_image[: pred_overlay.shape[0], :] = pred_overlay
-        new_image[pred_overlay.shape[0] :] = 255  # White background for legend
-
-        # Convert to PIL once for all legend entries
-        img_pil = Image.fromarray(new_image)
-        draw = ImageDraw.Draw(img_pil)
-
-        # Draw legend entries
-        start_y = pred_overlay.shape[0] + legend_padding
-        for i, task in enumerate(pred_tasks_found):
-            # Draw color box
-            box_x = 10
-            box_y = start_y + i * 30
-            box_size = 20
-            box_coords = (box_x, box_y, box_x + box_size, box_y + box_size)
-            draw.rectangle(box_coords, fill=colors[i])
-
-            # Draw text (vertically centered with color box)
-            text_y = box_y + (box_size - 12) // 2  # Assuming ~12px text height
-            draw.text((box_x + box_size + 10, text_y), task.target, fill=(0, 0, 0))
-
-        pred_overlay = img_pil
-    else:
-        pred_overlay = Image.fromarray(np.array(pred_overlay))
+    pred_overlay = add_legend(pred_overlay, pred_targets_found, colors)
 
     return pred_overlay, pred_tasks_not_found
 
@@ -194,3 +162,80 @@ def overlay_masks(
             np.uint8
         )
     return Image.fromarray(overlay)
+
+
+def add_legend(
+    image: Image.Image,
+    pred_targets_found: list[PredictionTarget],
+    colors: list[Tuple[int, int, int]],
+) -> Image.Image:
+    if len(pred_targets_found) == 0:
+        return image
+
+    # Convert to numpy for manipulation
+    pred_overlay = np.array(image)
+
+    # Calculate dimensions based on image resolution
+    image_width = pred_overlay.shape[1]
+    font_size = max(16, int(image_width * 0.02))  # Scale with image width, minimum 16px
+    box_size = int(font_size * 1.5)  # Color box proportional to font
+    entry_height = int(box_size * 1.5)  # Space between entries
+    legend_padding = int(font_size * 0.75)  # Padding scales with font
+
+    # Calculate total legend height
+    legend_height = entry_height * len(pred_targets_found)
+    total_height = pred_overlay.shape[0] + legend_height + 2 * legend_padding
+
+    # Create new image with space for legend
+    new_image = np.zeros((total_height, pred_overlay.shape[1], 3), dtype=np.uint8)
+    new_image[: pred_overlay.shape[0], :] = pred_overlay
+    new_image[pred_overlay.shape[0] :] = 255  # White background for legend
+
+    # Convert to PIL once for all legend entries
+    img_pil = Image.fromarray(new_image)
+    draw = ImageDraw.Draw(img_pil)
+
+    # Try to load a system font with proper scaling
+    font = None
+    system_fonts = [
+        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",  # Linux
+        "/System/Library/Fonts/Helvetica.ttc",  # macOS
+        "C:\\Windows\\Fonts\\arial.ttf",  # Windows
+    ]
+    for font_path in system_fonts:
+        try:
+            font = ImageFont.truetype(font_path, font_size)
+            break
+        except (OSError, IOError):
+            continue
+
+    if font is None:
+        # Fallback to default font if no system fonts are available
+        font = ImageFont.load_default()
+
+    # Get font metrics for proper vertical centering
+    bbox = font.getbbox("Aj")  # Use tall characters to get true height
+    font_height = bbox[3] - bbox[1]  # bottom - top
+
+    # Draw legend entries
+    start_y = pred_overlay.shape[0] + legend_padding
+    for i, task in enumerate(pred_targets_found):
+        # Draw color box
+        box_x = legend_padding
+        box_y = start_y + i * entry_height
+        box_coords = (box_x, box_y, box_x + box_size, box_y + box_size)
+        draw.rectangle(box_coords, fill=colors[i])
+
+        # Draw text (vertically centered with color box)
+        text_y = box_y + (box_size - font_height) // 2  # Center text with box
+        # Format text with truncated p-value
+        p_value_truncated = "{:.2f}".format(task.adjusted_p_value)
+        legend_text = f"{task.target} ({p_value_truncated})"
+        draw.text(
+            (box_x + box_size + legend_padding, text_y),
+            legend_text,
+            fill=(0, 0, 0),
+            font=font,
+        )
+
+    return img_pil