overlay version

app.py

@@ -1,8 +1,10 @@
 import gradio as gr
+from datetime import datetime
 from huggingface_hub import hf_hub_download
 import torch
 import json
 from PIL import Image
+from PIL import ImageDraw, ImageFont
 import numpy as np
 from model import MIPHEIViT
 
@@ -44,6 +46,7 @@ correction_map = {"Hoechst": 255.0, "CD8a": 100, "CD31": 100, "CD4": 100, "CD68"
 max_contrast_correction_value = torch.tensor([
     correction_map.get(name, default_contrast) / 255 for name in channel_names
 ]).reshape(len(channel_names), 1, 1)
+overlay_markers = ["Hoechst", "Pan-CK", "SMA", "CD45"]
 
 
 def preprocess(image):
@@ -52,6 +55,47 @@ def preprocess(image):
     tensor = (tensor - mean) / std
     return tensor.unsqueeze(0)  # [1, 3, H, W]
 
+
+def draw_legend_on_image(image, channel_names, channel_colors, indices, box_size=18, spacing=5, top_margin=5):
+    """Draw a semi-transparent legend on the bottom-right of the image."""
+    overlay = image.convert("RGBA")  # to allow alpha
+    legend_layer = Image.new("RGBA", overlay.size, (255, 255, 255, 0))
+    draw = ImageDraw.Draw(legend_layer)
+    font = ImageFont.load_default()
+
+    legend_height = top_margin + box_size * len(indices) + spacing * (len(indices) - 1)
+    legend_width = 60  # adjust as needed
+    x_start = overlay.width - legend_width - 10
+    y_start = overlay.height - legend_height - 10
+
+    # Semi-transparent background
+    draw.rectangle(
+        [x_start - 5, y_start - 5, x_start + legend_width + 5, y_start + legend_height + 5],
+        fill=(255, 255, 255, 180)  # semi-transparent white
+    )
+
+    for i, idx in enumerate(indices):
+        name = channel_names[idx]
+        color = channel_colors[name]
+        y = y_start + i * (box_size + spacing)
+        draw.rectangle([x_start, y, x_start + box_size, y + box_size], fill=color + (255,))
+        draw.text((x_start + box_size + 5, y), name, fill=(0, 0, 0, 255), font=font)
+
+    # Composite legend onto overlay
+    combined = Image.alpha_composite(overlay, legend_layer)
+    return combined.convert("RGB")  # back to RGB for display
+
+
+def merge_colored_images(color_imgs, top4_idx):
+    # Convert images to float32 NumPy arrays
+    accum = np.zeros_like(np.array(color_imgs[0]), dtype=np.float32)
+    for idx in top4_idx:
+        img = np.array(color_imgs[idx]).astype(np.float32)
+        accum += img  # additive blending
+
+    accum = np.clip(accum, 0, 255).astype(np.uint8)
+    return Image.fromarray(accum, mode='RGB')
+
 def apply_color_map(gray_img, rgb_color):
     """Map a grayscale image to RGB using a fixed pseudocolor."""
     gray = np.asarray(gray_img).astype(np.float32) / 255.0
@@ -59,6 +103,7 @@ def apply_color_map(gray_img, rgb_color):
     return Image.fromarray(rgb, mode='RGB')
 
 def predict(image):
+    print(f"[{datetime.now().isoformat()}] Inference run")
     input_tensor = preprocess(image)
     with torch.inference_mode():
         output = model(input_tensor)[0]  # [16, H, W]
@@ -76,8 +121,10 @@ def predict(image):
         ch_colored = apply_color_map(ch_gray, channel_colors[ch_name])
         channel_imgs.append(ch_colored)
 
-    # Return predicted 16 channels
-    return channel_imgs
+    fixed_idx = [channel_names.index(name) for name in overlay_markers]
+    overlay = merge_colored_images(channel_imgs, fixed_idx)
+    overlay_with_legend = draw_legend_on_image(overlay, channel_names, channel_colors, fixed_idx)
+    return [overlay_with_legend] + channel_imgs
 
 # Markdown header
 with open("HEADER.md", "r", encoding="utf-8") as f:
@@ -110,10 +157,12 @@ with gr.Blocks() as demo:
 
         # RIGHT: outputs
         with gr.Column(scale=2):
-            output_images = [
+            overlay_image = gr.Image(type="pil", label="mIF Overlay")
+            channel_images = [
                 gr.Image(type="pil", label=f"mIF Channel {channel_names[i]}")
                 for i in range(16)
             ]
+            output_images = [overlay_image] + channel_images
 
     run_btn.click(fn=predict, inputs=input_image, outputs=output_images)