app

README.md

@@ -1,12 +1,15 @@
 ---
-title: CheXmask U
-emoji: 🏆
-colorFrom: gray
-colorTo: purple
+title: Chest x-ray HybridGNet Segmentation
+emoji: ⚡
+colorFrom: yellow
+colorTo: blue
 sdk: gradio
-sdk_version: 6.0.0
+sdk_version: 3.50.2
 app_file: app.py
 pinned: false
+license: gpl-3.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Demo of the HybridGNet model with 2 image-to-graph skip connections from: arxiv.org/abs/2203.10977
+Original HybridGNet model: arxiv.org/abs/2106.09832
+The training procedure was taken from: arxiv.org/abs/2211.07395

app.py

@@ -0,0 +1,71 @@
+import gradio as gr
+from utils.segmentation import segment
+
+# ------------------------- GRADIO -------------------------
+if __name__ == "__main__":    
+    with gr.Blocks() as demo:
+        gr.Markdown("""
+                    # CheXmask-U: Uncertainty in Landmark-based Anatomical Segmentation
+                    
+                    Demo of the **uncertainty estimation framework** proposed in the paper "CheXmask-U: Quantifying uncertainty in landmark-based anatomical segmentation for X-ray images".
+                    The demonstration performs landmark-based segmentation (lungs and heart) and quantifies the uncertainty in the predicted position of each anatomical landmark.
+                    
+                    ### 📝 Instructions
+                    
+                    1.  **Upload** a chest X-ray image (PA or AP view) in PNG or JPEG format, or select an example image.
+                    2.  **Explore Prediction Variability**: You can either use your mouse to draw on the input image and perform inpainting in different regions or adjust the Gaussian Noise Std Dev slider to simulate image corruption. 
+                    3.  Click on **"Segment Image"**.
+                        * The output image will display the segmentation overlay where the color gradient indicates the node-wise predictive uncertainty.
+                        * The Results file output will contain the coordinates and per-node uncertainty estimates.
+                    
+                    Note: Pre-processing is not needed, it will be done automatically and removed after the segmentation.
+                    """)
+
+        with gr.Tab("Segment Image"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    image_input = gr.Image(
+                        type="numpy", 
+                        tool="sketch", 
+                        image_mode="L", 
+                        height=450, 
+                    )
+                    
+                    noise_slider = gr.Slider(
+                        label="Gaussian Noise Std Dev",
+                        minimum=0.0,
+                        maximum=0.25,
+                        step=0.01,
+                        value=0.0
+                    )
+                    
+                    with gr.Row():
+                        clear_button = gr.Button("Clear")
+                        image_button = gr.Button("Segment Image")
+
+                    gr.Examples(inputs=image_input, examples=[
+                        'utils/example1.jpg','utils/example2.jpg',
+                        'utils/example3.png','utils/example4.jpg'
+                    ])
+                    
+                with gr.Column(scale=2):
+                    image_output = gr.Image(type="filepath", height=450)
+                    results = gr.File()       
+                    
+        gr.Markdown("""
+                    Example images extracted from Wikipedia, released under:
+                    1. CC0 Universial Public Domain. Source: https://commons.wikimedia.org/wiki/File:Normal_posteroanterior_(PA)_chest_radiograph_(X-ray).jpg
+                    2. Creative Commons Attribution-Share Alike 4.0 International. Source: https://commons.wikimedia.org/wiki/File:Chest_X-ray.jpg
+                    3. Creative Commons Attribution 3.0 Unported. Source https://commons.wikimedia.org/wiki/File:Implantable_cardioverter_defibrillator_chest_X-ray.jpg
+                    4. Creative Commons Attribution-Share Alike 3.0 Unported. Source: https://commons.wikimedia.org/wiki/File:Medical_X-Ray_imaging_PRD06_nevit.jpg
+                    """)
+                    
+        clear_button.click(lambda: None, None, image_input, queue=False)
+        clear_button.click(lambda: None, None, image_output, queue=False)
+        image_button.click(
+            segment,
+            inputs=[image_input, noise_slider],
+            outputs=[image_output, results],
+            queue=False
+        )
+    demo.launch(share=True)

models/HybridGNet2IGSC.py

@@ -0,0 +1,208 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models.modelUtils import ChebConv, Pool, residualBlock
+import torchvision.ops.roi_align as roi_align
+import numpy as np
+
+class EncoderConv(nn.Module):
+    def __init__(self, latents = 64, hw = 32):
+        super(EncoderConv, self).__init__()
+        
+        self.latents = latents
+        self.c = 4
+        
+        self.size = self.c * np.array([2,4,8,16,32], dtype = np.intc)
+        
+        self.maxpool = nn.MaxPool2d(2)
+        
+        self.dconv_down1 = residualBlock(1, self.size[0])
+        self.dconv_down2 = residualBlock(self.size[0], self.size[1])
+        self.dconv_down3 = residualBlock(self.size[1], self.size[2])
+        self.dconv_down4 = residualBlock(self.size[2], self.size[3])
+        self.dconv_down5 = residualBlock(self.size[3], self.size[4])
+        self.dconv_down6 = residualBlock(self.size[4], self.size[4])
+        
+        self.fc_mu = nn.Linear(in_features=self.size[4]*hw*hw, out_features=self.latents)
+        self.fc_logvar = nn.Linear(in_features=self.size[4]*hw*hw, out_features=self.latents)
+
+    def forward(self, x):
+        x = self.dconv_down1(x)
+        x = self.maxpool(x)
+
+        x = self.dconv_down2(x)
+        x = self.maxpool(x)
+        
+        conv3 = self.dconv_down3(x)
+        x = self.maxpool(conv3)
+        
+        conv4 = self.dconv_down4(x)
+        x = self.maxpool(conv4)
+        
+        conv5 = self.dconv_down5(x)
+        x = self.maxpool(conv5)
+        
+        conv6 = self.dconv_down6(x)
+        
+        x = conv6.view(conv6.size(0), -1) # flatten batch of multi-channel feature maps to a batch of feature vectors
+        
+        x_mu = self.fc_mu(x)
+        x_logvar = self.fc_logvar(x)
+                
+        return x_mu, x_logvar, conv6, conv5
+
+
+class SkipBlock(nn.Module):
+    def __init__(self, in_filters, window):
+        super(SkipBlock, self).__init__()
+        
+        self.window = window
+        self.graphConv_pre = ChebConv(in_filters, 2, 1, bias = False) 
+    
+    def lookup(self, pos, layer, salida = (1,1)):
+        B = pos.shape[0]
+        N = pos.shape[1]
+        F = layer.shape[1]
+        h = layer.shape[-1]
+        
+        ## Scale from [0,1] to [0, h]
+        pos = pos * h
+        
+        _x1 = (self.window[0] // 2) * 1.0
+        _x2 = (self.window[0] // 2 + 1) * 1.0
+        _y1 = (self.window[1] // 2) * 1.0
+        _y2 = (self.window[1] // 2 + 1) * 1.0
+        
+        boxes = []
+        for batch in range(0, B):
+            x1 = pos[batch,:,0].reshape(-1, 1) - _x1
+            x2 = pos[batch,:,0].reshape(-1, 1) + _x2
+            y1 = pos[batch,:,1].reshape(-1, 1) - _y1
+            y2 = pos[batch,:,1].reshape(-1, 1) + _y2
+            
+            aux = torch.cat([x1, y1, x2, y2], axis = 1)            
+            boxes.append(aux)
+                    
+        skip = roi_align(layer, boxes, output_size = salida, aligned=True)
+        vista = skip.view([B, N, -1])
+
+        return vista
+    
+    def forward(self, x, adj, conv_layer):
+        pos = self.graphConv_pre(x, adj)
+        skip = self.lookup(pos, conv_layer)
+        
+        return torch.cat((x, skip, pos), axis = 2), pos
+    
+
+class Hybrid(nn.Module):
+    def __init__(self, config, downsample_matrices, upsample_matrices, adjacency_matrices):
+        super(Hybrid, self).__init__()
+        
+        self.config = config
+        hw = config['inputsize'] // 32
+        self.z = config['latents']
+        self.encoder = EncoderConv(latents = self.z, hw = hw)
+        self.eval_sampling = config['eval_sampling']
+        
+        self.downsample_matrices = downsample_matrices
+        self.upsample_matrices = upsample_matrices
+        self.adjacency_matrices = adjacency_matrices
+        self.kld_weight = 1e-5
+                
+        n_nodes = config['n_nodes']
+        self.filters = config['filters']
+        self.K = 6
+        self.window = (3,3)
+        
+        # Generate the fully connected layer for the decoder
+        outshape = self.filters[-1] * n_nodes[-1]          
+        self.dec_lin = torch.nn.Linear(self.z, outshape)
+                
+        self.normalization2u = torch.nn.InstanceNorm1d(self.filters[1])
+        self.normalization3u = torch.nn.InstanceNorm1d(self.filters[2])
+        self.normalization4u = torch.nn.InstanceNorm1d(self.filters[3])
+        self.normalization5u = torch.nn.InstanceNorm1d(self.filters[4])
+        self.normalization6u = torch.nn.InstanceNorm1d(self.filters[5])
+        
+        outsize1 = self.encoder.size[4]
+        outsize2 = self.encoder.size[4]  
+                     
+        # Store graph convolution layers
+        self.graphConv_up6 = ChebConv(self.filters[6], self.filters[5], self.K)
+        self.graphConv_up5 = ChebConv(self.filters[5], self.filters[4], self.K)       
+        
+        self.SC_1 = SkipBlock(self.filters[4], self.window)
+        
+        self.graphConv_up4 = ChebConv(self.filters[4] + outsize1 + 2, self.filters[3], self.K)        
+        self.graphConv_up3 = ChebConv(self.filters[3], self.filters[2], self.K)
+        
+        self.SC_2 = SkipBlock(self.filters[2], self.window)
+        
+        self.graphConv_up2 = ChebConv(self.filters[2] + outsize2 + 2, self.filters[1], self.K)
+        self.graphConv_up1 = ChebConv(self.filters[1], self.filters[0], 1, bias = False)
+                
+        self.pool = Pool()
+        
+        self.reset_parameters()
+        
+    def reset_parameters(self):
+        torch.nn.init.normal_(self.dec_lin.weight, 0, 0.1)
+
+    def sampling(self, mu, log_var):
+        std = torch.exp(0.5*log_var)
+        eps = torch.randn_like(std)
+        return eps.mul(std).add_(mu) 
+    
+    def encode(self, x):
+        """Encode the input and return latent representations and skip connections"""
+        mu, log_var, conv6, conv5 = self.encoder(x)
+        return mu, log_var, conv6, conv5
+    
+    def decode(self, z, conv6, conv5):
+        """Decode from latent space using skip connections"""
+        x = self.dec_lin(z)
+        x = F.relu(x)
+        
+        x = x.reshape(x.shape[0], -1, self.filters[-1])
+        
+        x = self.graphConv_up6(x, self.adjacency_matrices[5]._indices())
+        x = self.normalization6u(x)
+        x = F.relu(x)
+        
+        x = self.graphConv_up5(x, self.adjacency_matrices[4]._indices())
+        x = self.normalization5u(x)
+        x = F.relu(x)
+        
+        x, pos1 = self.SC_1(x, self.adjacency_matrices[3]._indices(), conv6)
+        
+        x = self.graphConv_up4(x, self.adjacency_matrices[3]._indices())
+        x = self.normalization4u(x)
+        x = F.relu(x)
+        
+        x = self.pool(x, self.upsample_matrices[0])
+        
+        x = self.graphConv_up3(x, self.adjacency_matrices[2]._indices())
+        x = self.normalization3u(x)
+        x = F.relu(x)
+        
+        x, pos2 = self.SC_2(x, self.adjacency_matrices[1]._indices(), conv5)
+        
+        x = self.graphConv_up2(x, self.adjacency_matrices[1]._indices())
+        x = self.normalization2u(x)
+        x = F.relu(x)
+        
+        x = self.graphConv_up1(x, self.adjacency_matrices[0]._indices()) # No relu and no bias
+        
+        return x, pos1, pos2
+        
+    def forward(self, x):
+        """Full forward pass (both encoding and decoding)"""
+        self.mu, self.log_var, conv6, conv5 = self.encode(x)
+        
+        if self.training or self.eval_sampling:
+            z = self.sampling(self.mu, self.log_var)
+        else:
+            z = self.mu
+            
+        return self.decode(z, conv6, conv5) 

models/modelUtils.py

@@ -0,0 +1,69 @@
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.conv.cheb_conv import ChebConv
+from torch_geometric.nn.inits import zeros, normal
+
+# We change the default initialization from zeros to a normal distribution
+class ChebConv(ChebConv):
+    def reset_parameters(self):
+        for lin in self.lins:
+            normal(lin, mean = 0, std = 0.1)
+            #lin.reset_parameters()
+        normal(self.bias, mean = 0, std = 0.1)
+        #zeros(self.bias)
+
+# Pooling from COMA: https://github.com/pixelite1201/pytorch_coma/blob/master/layers.py
+class Pool(MessagePassing):
+    def __init__(self):
+        # source_to_target is the default value for flow, but is specified here for explicitness
+        super(Pool, self).__init__(flow='source_to_target')
+
+    def forward(self, x, pool_mat,  dtype=None):
+        pool_mat = pool_mat.transpose(0, 1)
+        out = self.propagate(edge_index=pool_mat._indices(), x=x, norm=pool_mat._values(), size=pool_mat.size())
+        return out
+
+    def message(self, x_j, norm):
+        return norm.view(1, -1, 1) * x_j
+    
+    
+import torch.nn as nn
+import torch.nn.functional as F
+
+class residualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1):
+        """
+        Args:
+          in_channels (int):  Number of input channels.
+          out_channels (int): Number of output channels.
+          stride (int):       Controls the stride.
+        """
+        super(residualBlock, self).__init__()
+
+        self.skip = nn.Sequential()
+
+        if stride != 1 or in_channels != out_channels:
+          self.skip = nn.Sequential(
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, bias=False),
+            nn.BatchNorm2d(out_channels, track_running_stats=False))
+        else:
+          self.skip = None
+
+        self.block = nn.Sequential(nn.BatchNorm2d(in_channels, track_running_stats=False),
+                                   nn.ReLU(inplace=True),
+                                   nn.Conv2d(in_channels, out_channels, 3, padding=1),
+                                   nn.BatchNorm2d(out_channels, track_running_stats=False),
+                                   nn.ReLU(inplace=True),
+                                   nn.Conv2d(out_channels, out_channels, 3, padding=1)
+                                   )   
+
+    def forward(self, x):
+        identity = x
+        out = self.block(x)
+
+        if self.skip is not None:
+            identity = self.skip(x)
+
+        out += identity
+        out = F.relu(out)
+
+        return out

requirements.txt

@@ -0,0 +1,6 @@
+torch==2.0.1
+numpy==1.25.0
+opencv-python==4.8.0.74
+scipy==1.10.1
+torch_geometric==2.3.0
+torchvision==0.15.2

weights/weights.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c5658bf5ff990fbcccde968a198bb48b57db22ffeaca32a837d6031a10395e2
+size 70083499