example, color, contrast

HEADER.md

@@ -1,6 +1,9 @@
 # MIPHEI-ViT Demo: 16-channel mIF Prediction
 
 <p align="center">
+  <a title="arXiv" href="https://arxiv.org/abs/2505.10294" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+        <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
+    </a>
   <a href="https://huggingface.co/Estabousi/MIPHEI-vit" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
     <img src="https://img.shields.io/badge/🤗 Model-MIPHEI--ViT-lightgrey?logo=huggingface" height="25">
   </a>
@@ -17,7 +20,7 @@ The model returns **16 grayscale images**, each representing a predicted mIF mar
 
 ---
 
-Try it with low-zoom screenshots from public datasets:
+Try it with **provided examples** or low-zoom screenshots from public datasets:
 
 **ORION (in-domain test set):**
 - [CRC2](https://labsyspharm.github.io/orion-crc/minerva/P37_S30-CRC02/index.html#s=0&w=0&g=5&m=-1&a=-100_-100&v=1.0673_0.6057_0.5&o=-100_-100_1_1&p=Q)

app.py

@@ -18,25 +18,63 @@ with open(config_path, "r") as f:
     config = json.load(f)
 channel_names = config["targ_channel_names"]
 
+
+channel_colors = {
+    "Hoechst": (0, 0, 255),         # Blue (DAPI, nuclear stain)
+    "CD31": (0, 255, 255),          # Cyan (endothelial)
+    "CD45": (255, 255, 0),          # Yellow (leukocyte common antigen)
+    "CD68": (255, 165, 0),          # Orange (macrophages)
+    "CD4": (255, 0, 0),             # Red (helper T cells)
+    "FOXP3": (138, 43, 226),        # Purple/Blue-Violet (regulatory T cells)
+    "CD8a": (303, 100, 100),            # Green (cytotoxic T cells)
+    "CD45RO": (255, 105, 180),      # Hot Pink (memory T cells)
+    "CD20": (0, 191, 255),          # Deep Sky Blue (B cells)
+    "PD-L1": (255, 0, 255),         # Magenta
+    "CD3e": (95, 95, 94),          # Crimson (T cells)
+    "CD163": (184, 134, 11),        # Dark Goldenrod (M2 macrophages)
+    "E-cadherin": (242, 12, 43),    # Spring Green (epithelial marker)
+    "Ki67": (255, 20, 147),         # Deep Pink (proliferation marker)
+    "Pan-CK": (255, 0, 0),          # Red (epithelial/carcinoma)
+    "SMA": (0, 255, 0),             # Green (smooth muscle, myofibroblasts)
+}
+
+# Contrast correction factors per channel (255 for Hoechst, 150 otherwise)
+default_contrast = 150.0
+correction_map = {"Hoechst": 255.0, "CD8a": 100, "CD31": 100, "CD4": 100, "CD68": 100, "FOXP3": 100}
+max_contrast_correction_value = torch.tensor([
+    correction_map.get(name, default_contrast) / 255 for name in channel_names
+]).reshape(len(channel_names), 1, 1)
+
+
 def preprocess(image):
     image = image.convert("RGB").resize((256, 256))
     tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255
     tensor = (tensor - mean) / std
     return tensor.unsqueeze(0)  # [1, 3, H, W]
 
+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
+    rgb = np.stack([gray * rgb_color[i] for i in range(3)], axis=-1).astype(np.uint8)
+    return Image.fromarray(rgb, mode='RGB')
+
 def predict(image):
     input_tensor = preprocess(image)
     with torch.inference_mode():
         output = model(input_tensor)[0]  # [16, H, W]
         output = (output.clamp(-0.9, 0.9) + 0.9) / 1.8
-        output = np.uint8(output.cpu().numpy() * 255)
+        output_vis = output / max_contrast_correction_value.to(output.device).clamp(min=1e-6)
+        output_vis = output_vis.clamp(0, 1) * 255
+        output_vis = np.uint8(output_vis.cpu().numpy())
+        output = output.cpu().numpy()
 
     # Convert each mIF channel to grayscale PIL image
     channel_imgs = []
-    for i in range(output.shape[0]):
-        ch_img = output[i]
-        pil_ch = Image.fromarray(ch_img, mode='L')
-        channel_imgs.append(pil_ch)
+    for i in range(output_vis.shape[0]):
+        ch_name = channel_names[i]
+        ch_gray = Image.fromarray(output_vis[i], mode='L')
+        ch_colored = apply_color_map(ch_gray, channel_colors[ch_name])
+        channel_imgs.append(ch_colored)
 
     # Return predicted 16 channels
     return channel_imgs
@@ -47,14 +85,37 @@ with open("HEADER.md", "r", encoding="utf-8") as f:
 
 # Build interface using Blocks
 with gr.Blocks() as demo:
+
     gr.Markdown(HEADER_MD)
-    gr.Interface(
-        fn=predict,
-        inputs=gr.Image(type="pil", label="Input H&E"),
-        outputs=[gr.Image(type="pil", label=f"mIF Channel {channel_names[i]}") for i in range(16)],
-        title=None,
-        description=None
-    )
+
+    with gr.Row():
+        # LEFT: input + examples + button
+        with gr.Column(scale=0.5):
+            input_image = gr.Image(type="pil", label="Input H&E")
+            run_btn = gr.Button("Run Prediction")
+            gr.Examples(
+                examples=[
+                    ["examples/crc100k_val.jpg"],
+                    ["examples/orion_test_1.jpg"],
+                    ["examples/orion_test_2.jpg"],
+                    ["examples/orion_test_3.jpg"],
+                    ["examples/orion_test_4.jpg"],
+                    ["examples/orion_test_5.jpg"],
+                    ["examples/tcga.jpg"],
+                    ["examples/hemit.jpg"],
+                ],
+                inputs=[input_image],
+                label="Example H&E tile (TCGA, ORION Test, CRC100K, HEMIT)"
+            )
+
+        # RIGHT: outputs
+        with gr.Column():
+            output_images = [
+                gr.Image(type="pil", label=f"mIF Channel {channel_names[i]}")
+                for i in range(16)
+            ]
+
+    run_btn.click(fn=predict, inputs=input_image, outputs=output_images)
 
 if __name__ == "__main__":
     demo.launch()