first commit

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.py[cod]
+*.so

LICENSE

@@ -0,0 +1,9 @@
+Non Commercial License Notice: 
+
+Copyright (c) 2025 Sanofi.
+
+Permission is hereby granted, free of charge, for academic research purposes only and for non-commercial uses only, to any person from academic research or non-profit organizations obtaining a copy of this software and associated documentation files (the "Software"), to use, copy, modify, or merge the Software, subject to the following conditions: this permission notice shall be included in all copies of the Software or of substantial portions of the Software. 
+
+For purposes of this license, “non-commercial use” excludes uses foreseeably resulting in a commercial benefit. To use this software for other purposes (such as the development of a commercial product, including but not limited to software, service, or pharmaceuticals, or in a collaboration with a private company), please contact SANOFI at [email protected]. 
+
+All other rights are reserved, including those for text and data mining, AI training and similar technologies. The Software is provided “as is”, without warranty of any kind, express or implied, including the warranties of noninfringement. 

app.py

@@ -0,0 +1,59 @@
+import gradio as gr
+from huggingface_hub import hf_hub_download
+import torch
+import json
+from PIL import Image
+import numpy as np
+from model import MIPHEIViT
+
+# Load model once
+repo_id = "Estabousi/MIPHEI-vit"
+model = MIPHEIViT.from_pretrained_hf(repo_id=repo_id)
+config_path = hf_hub_download(repo_id=repo_id, filename="config_hf.json")
+
+model.eval()
+mean = torch.Tensor([0.485, 0.456, 0.406]).to(torch.float32).reshape((-1, 1, 1))
+std = torch.Tensor([0.229, 0.224, 0.225]).to(torch.float32).reshape((-1, 1, 1))
+with open(config_path, "r") as f:
+    config = json.load(f)
+channel_names = config["targ_channel_names"]
+
+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 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)
+
+    # 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)
+
+    # Return list: input + 16 channels
+    return channel_imgs
+
+# Prepare Gradio UI
+demo = 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="MIPHEI-ViT: Full mIF Prediction",
+    description=(
+        "Upload an H&E image tile (colorectal 256×256 pxs at 0.5 µm/px recommended). "
+        "The image will be resized to (256×256) if needed.\n"
+        "The model predicts 16-channel multiplex immunofluorescence, "
+        "with each marker shown as a grayscale image."
+)
+)
+
+if __name__ == "__main__":
+    demo.launch()

model.py

@@ -0,0 +1,410 @@
+"""
+This script defines the MIPHEI-ViT architecture for image-to-image translation
+Some modules in this file are adapted from: https://github.com/hustvl/ViTMatte/
+"""
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import timm
+from timm.models import VisionTransformer, SwinTransformer
+from timm.models import load_state_dict_from_hf
+
+
+class Basic_Conv3x3(nn.Module):
+    """
+    Basic convolution layers including: Conv3x3, BatchNorm2d, ReLU layers.
+    https://github.com/hustvl/ViTMatte/blob/main/modeling/decoder/detail_capture.py#L5
+    """
+    def __init__(
+        self,
+        in_chans,
+        out_chans,
+        stride=2,
+        padding=1,
+    ):
+        super().__init__()
+        self.conv = nn.Conv2d(in_chans, out_chans, 3, stride, padding, bias=False)
+        self.bn = nn.BatchNorm2d(out_chans)
+        self.relu = nn.ReLU(inplace=False)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+
+        return x
+
+
+class ConvStream(nn.Module):
+    """
+    Simple ConvStream containing a series of basic conv3x3 layers to extract detail features.
+    """
+    def __init__(
+        self,
+        in_chans = 4,
+        out_chans = [48, 96, 192],
+    ):
+        super().__init__()
+        self.convs = nn.ModuleList()
+        
+        self.conv_chans = out_chans.copy()
+        self.conv_chans.insert(0, in_chans)
+        
+        for i in range(len(self.conv_chans)-1):
+            in_chan_ = self.conv_chans[i]
+            out_chan_ = self.conv_chans[i+1]
+            self.convs.append(
+                Basic_Conv3x3(in_chan_, out_chan_)
+            )
+    
+    def forward(self, x):
+        out_dict = {'D0': x}
+        for i in range(len(self.convs)):
+            x = self.convs[i](x)
+            name_ = 'D'+str(i+1)
+            out_dict[name_] = x
+        
+        return out_dict
+
+
+class SegmentationHead(nn.Sequential):
+    # https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/segmentation_models_pytorch/base/heads.py#L5
+    def __init__(
+        self, in_channels, out_channels, kernel_size=3, activation=None, use_attention=False,
+    ):
+        if use_attention:
+            attention = AttentionBlock(in_channels)
+        else:
+            attention = nn.Identity()
+        conv2d = nn.Conv2d(
+            in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2
+        )
+        activation = activation
+        super().__init__(attention, conv2d, activation)
+
+
+class AttentionBlock(nn.Module):
+    """
+    Attention gate
+
+    Parameters:
+    -----------
+    in_chns : int
+        Number of input channels.
+
+    Forward Input:
+    --------------
+    x : torch.Tensor
+        Input tensor of shape [B, C, H, W].
+
+    Returns:
+    --------
+    torch.Tensor
+        Reweighted tensor of the same shape as input.
+    """
+    def __init__(self, in_chns):
+        super(AttentionBlock, self).__init__()
+        # Attention generation
+        self.psi = nn.Sequential(
+            nn.Conv2d(in_chns, in_chns // 2, kernel_size=1, stride=1, padding=0, bias=True),
+            nn.BatchNorm2d(in_chns // 2),
+            nn.ReLU(),
+            nn.Conv2d(in_chns // 2, 1, kernel_size=1, stride=1, padding=0, bias=True),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        # Project decoder output to intermediate space
+        g = self.psi(x)
+        return x * g
+
+
+class Fusion_Block(nn.Module):
+    """
+    Simple fusion block to fuse feature from ConvStream and Plain Vision Transformer.
+    """
+    def __init__(
+        self,
+        in_chans,
+        out_chans,
+    ):
+        super().__init__()
+        self.conv = Basic_Conv3x3(in_chans, out_chans, stride=1, padding=1)
+
+    def forward(self, x, D):
+        F_up = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False) ## Nearest ?
+        out = torch.cat([D, F_up], dim=1)
+        out = self.conv(out)
+
+        return out
+
+
+class MIPHEIViT(nn.Module):
+    """
+    U-Net-style architecture inspired by ViTMatte, using a Vision Transformer (ViT or Swin)
+    as encoder and a convolutional decoder. Designed for dense image prediction tasks,
+    such as image-to-image translation.
+
+    Parameters:
+    -----------
+    encoder : nn.Module
+        A ViT- or Swin-based encoder that outputs spatial feature maps.
+    decoder : nn.Module
+        A decoder module that maps encoder features (and optionally the original image)
+        to the output prediction.
+
+    Example:
+    --------
+    model = MIPHEIViT(encoder=Encoder(vit), decoder=UNetDecoder())
+    output = model(input_tensor)
+    """
+    def __init__(self,
+                 encoder,
+                 decoder,
+                 ):
+        super(MIPHEIViT, self).__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.initialize()
+
+    def forward(self, x):
+
+        features = self.encoder(x)
+        outputs = self.decoder(features, x)
+        return outputs
+
+    def initialize(self):
+        pass
+
+    @classmethod
+    def from_pretrained_hf(cls, repo_path=None, repo_id=None):
+        from safetensors.torch import load_file
+        import json
+        if repo_path:
+            weights_path = os.path.join(repo_path, "model.safetensors")
+            config_path = os.path.join(repo_path, "config_hf.json")
+        else:
+            from huggingface_hub import hf_hub_download
+            weights_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors")
+            config_path = hf_hub_download(repo_id=repo_id, filename="config_hf.json")
+        
+        # Load config values
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        img_size = config["img_size"]
+        nc_out = len(config["targ_channel_names"])
+        use_attention = config["use_attention"]
+        hoptimus_hf_id = config["hoptimus_hf_id"]
+
+        vit = get_hoptimus0_hf(hoptimus_hf_id)
+        vit.set_input_size(img_size=(img_size, img_size))
+        encoder = Encoder(vit)
+        decoder = Detail_Capture(emb_chans=encoder.embed_dim, out_chans=nc_out, use_attention=use_attention, activation=nn.Tanh())
+        model = cls(encoder=encoder, decoder=decoder)
+        state_dict = load_file(weights_path)
+        state_dict = merge_lora_weights(model, state_dict)
+        load_info = model.load_state_dict(state_dict, strict=False)
+        validate_load_info(load_info)
+        model.eval()
+        return model
+
+    def set_input_size(self, img_size):
+        if any((s & (s - 1)) != 0 or s == 0 for s in img_size):
+            raise ValueError("Both height and width in img_size must be powers of 2")
+        if any(s < 128 for s in img_size):
+            raise ValueError("Height and width must be greater or equal to 128")
+        self.encoder.vit.set_input_size(img_size=img_size)
+        self.encoder.grid_size = self.encoder.vit.patch_embed.grid_size
+
+
+class Encoder(nn.Module):
+    """
+    Wraps a Vision Transformer (ViT or Swin) to produce feature maps compatible
+    with U-Net-like architectures. It reshapes and resizes transformer outputs
+    into spatial feature maps.
+
+    Parameters:
+    -----------
+    vit : VisionTransformer or SwinTransformer
+        A pretrained transformer model from `timm` that outputs patch embeddings.
+    """
+    def __init__(self, vit):
+        super().__init__()
+        if not isinstance(vit, (VisionTransformer, SwinTransformer)):
+            raise ValueError(f"Expected a VisionTransformer or SwinTransformer, got {type(vit)}")
+        self.vit = vit
+
+        self.is_swint = isinstance(vit, SwinTransformer)
+        self.grid_size = self.vit.patch_embed.grid_size
+        if self.is_swint:
+            self.num_prefix_tokens = 0
+            self.embed_dim = self.vit.embed_dim * 2 ** (self.vit.num_layers -1)
+        else:
+            self.num_prefix_tokens = self.vit.num_prefix_tokens
+            self.embed_dim = self.vit.embed_dim
+        patch_size = self.vit.patch_embed.patch_size
+        img_size = self.vit.patch_embed.img_size
+        assert img_size[0] % 16 == 0
+        assert img_size[1] % 16 == 0
+
+        if self.is_swint:
+            self.scale_factor = (2., 2.)
+        else:
+            if patch_size != (16, 16):
+                target_grid_size = (img_size[0] / 16, img_size[1] / 16)
+                self.scale_factor = (target_grid_size[0] / self.grid_size[0], target_grid_size[1] / self.grid_size[1])
+            else:
+                self.scale_factor = None
+
+    def forward(self, x):
+        features = self.vit(x)
+        if self.is_swint:
+            features = features.permute(0, 3, 1, 2)
+        else:
+            features = features[:, self.num_prefix_tokens:]
+            features = features.permute(0, 2, 1)
+            features = features.view((-1, self.embed_dim, *self.grid_size))
+        if self.scale_factor is not None:
+            features = F.interpolate(features, scale_factor=self.scale_factor, mode="bicubic")
+        return features
+
+
+class Detail_Capture(nn.Module):
+    """
+    Simple and Lightweight Detail Capture Module for ViT Matting.
+    """
+    def __init__(
+        self,
+        emb_chans,
+        in_chans=3,
+        out_chans=1,
+        convstream_out = [48, 96, 192],
+        fusion_out = [256, 128, 64, 32],
+        use_attention=True,
+        activation=torch.nn.Identity()
+    ):
+        super().__init__()
+        assert len(fusion_out) == len(convstream_out) + 1
+
+        self.convstream = ConvStream(in_chans=in_chans)
+        self.conv_chans = self.convstream.conv_chans
+        self.num_heads = out_chans
+
+        self.fusion_blks = nn.ModuleList()
+        self.fus_channs = fusion_out.copy()
+        self.fus_channs.insert(0, emb_chans)
+        for i in range(len(self.fus_channs)-1):
+            self.fusion_blks.append(
+                Fusion_Block(
+                    in_chans = self.fus_channs[i] + self.conv_chans[-(i+1)],
+                    out_chans = self.fus_channs[i+1],
+                )
+            )
+
+        for idx in range(self.num_heads):
+            setattr(self, f'segmentation_head_{idx}', SegmentationHead(
+                in_channels=fusion_out[-1],
+                out_channels=1,
+                activation=activation,
+                kernel_size=3,
+                use_attention=use_attention
+            ))
+
+    def forward(self, features, images):
+        detail_features = self.convstream(images)
+        for i in range(len(self.fusion_blks)):
+            d_name_ = 'D'+str(len(self.fusion_blks)-i-1)
+            features = self.fusion_blks[i](features, detail_features[d_name_])
+        
+        outputs = []
+        for idx_head in range(self.num_heads):
+            segmentation_head = getattr(self, f'segmentation_head_{idx_head}')
+            output = segmentation_head(features)
+            outputs.append(output)
+        outputs = torch.cat(outputs, dim=1)
+
+        return outputs
+
+
+def merge_lora_weights(model, state_dict, alpha=1.0, block_prefix="encoder.vit.blocks"):
+    """
+    Merges LoRA weights into the base attention Q and V projection weights for each transformer block.
+    We keep LoRA weights in the model.safetensors to avoid having the original foundation model weights in the repo.
+
+    Parameters:
+    -----------
+    model : torch.nn.Module
+        The model containing the transformer blocks to modify (e.g., ViT backbone).
+    state_dict : dict
+        The state_dict containing LoRA matrices with keys formatted as 
+        '{block_prefix}.{idx}.attn.qkv.lora_q.A', etc.
+        This dict is modified in-place to remove LoRA weights after merging.
+    alpha : float, optional
+        Scaling factor for the LoRA update. Defaults to 1.0.
+    block_prefix : str, optional
+        Prefix to locate transformer blocks in the model. Defaults to "encoder.vit.blocks".
+
+    Returns:
+    --------
+    dict
+        The modified state_dict with LoRA weights removed after merging.
+    """
+    with torch.no_grad():
+        for idx in range(len(model.encoder.vit.blocks)):
+            prefix = f"{block_prefix}.{idx}.attn.qkv"
+
+            # Extract LoRA matrices
+            A_q = state_dict.pop(f"{prefix}.lora_q.A")
+            B_q = state_dict.pop(f"{prefix}.lora_q.B")
+            A_v = state_dict.pop(f"{prefix}.lora_v.A")
+            B_v = state_dict.pop(f"{prefix}.lora_v.B")
+
+            # Compute low-rank updates (transposed to match weight shape)
+            delta_q = (alpha * A_q @ B_q).T
+            delta_v = (alpha * A_v @ B_v).T
+
+            # Get original QKV weight matrix (shape: [3*dim, dim])
+            W = model.get_parameter(f"{prefix}.weight")
+            dim = delta_q.shape[0]
+            assert W.shape[0] == 3 * dim, f"Unexpected QKV shape: {W.shape}"
+
+            # Apply LoRA deltas to Q and V projections
+            W[:dim, :] += delta_q           # Q projection
+            W[2 * dim:, :] += delta_v       # V projection
+
+    return state_dict
+
+
+def get_hoptimus0_hf(repo_id):
+    """ Hoptimus foundation model from hugginface repo id
+    """
+    model = timm.create_model(
+        "vit_giant_patch14_reg4_dinov2", img_size=224,
+        drop_path_rate=0., num_classes=0,
+        global_pool="", pretrained=False, init_values=1e-5,
+        dynamic_img_size=False)
+    state_dict = load_state_dict_from_hf(repo_id, weights_only=True)
+    model.load_state_dict(state_dict)
+    return model
+
+
+def validate_load_info(load_info):
+    """
+    Validates the result of model.load_state_dict(..., strict=False).
+
+    Raises:
+        ValueError if unexpected keys are found,
+        or if missing keys are not related to the allowed encoder modules.
+    """
+    # 1. Raise if any unexpected keys
+    if load_info.unexpected_keys:
+        raise ValueError(f"Unexpected keys in state_dict: {load_info.unexpected_keys}")
+
+    # 2. Raise if any missing keys are not part of allowed encoder modules
+    for key in load_info.missing_keys:
+        if ".lora" in key:
+            raise ValueError(f"Missing LoRA checkpoint in state_dict: {key}")
+        elif not any(part in key for part in ["encoder.vit.", "encoder.model."]):
+            raise ValueError(f"Missing key in state_dict: {key}")

requirements.txt

@@ -0,0 +1,7 @@
+torch
+timm
+gradio
+safetensors
+numpy
+Pillow
+huggingface_hub