Add legend to prediction overlay with task color boxes and labels

inference_utils/model.py

@@ -14,7 +14,7 @@ from dataclasses import dataclass
 import os
 from typing import Tuple
 
-from PIL import Image
+from PIL import Image, ImageDraw
 from huggingface_hub import hf_hub_download
 import matplotlib.pyplot as plt
 import numpy as np
@@ -113,6 +113,43 @@ def predict(
     masks = [1 * (pred_mask[i] > 0.5) for i in range(len(pred_tasks_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))
+
     return pred_overlay, pred_tasks_not_found