First push

app.py

@@ -0,0 +1,126 @@
+import numpy as np
+import gradio as gr
+import cv2 
+
+from models.HybridGNet2IGSC import Hybrid 
+from utils.utils import scipy_to_torch_sparse, genMatrixesLungsHeart
+import scipy.sparse as sp
+import torch
+
+
+def getDenseMask(landmarks):
+    RL = landmarks[0:44]
+    LL = landmarks[44:94]
+    H = landmarks[94:]
+    
+    img = np.zeros([1024,1024], dtype = 'uint8')
+    
+    RL = RL.reshape(-1, 1, 2).astype('int')
+    LL = LL.reshape(-1, 1, 2).astype('int')
+    H = H.reshape(-1, 1, 2).astype('int')
+
+    img = cv2.drawContours(img, [RL], -1, 1, -1)
+    img = cv2.drawContours(img, [LL], -1, 1, -1)
+    img = cv2.drawContours(img, [H], -1, 2, -1)
+    
+    return img
+
+
+def drawOnTop(img, landmarks):
+    output = getDenseMask(landmarks)
+    
+    image = np.zeros([1024, 1024, 3])
+    image[:,:,0] = img + 0.3 * (output == 1).astype('float') - 0.1 * (output == 2).astype('float')
+    image[:,:,1] = img + 0.3 * (output == 2).astype('float') - 0.1 * (output == 1).astype('float') 
+    image[:,:,2] = img - 0.1 * (output == 1).astype('float') - 0.2 * (output == 2).astype('float') 
+
+    image = np.clip(image, 0, 1)
+
+    return image
+    
+
+def loadModel(device):
+    A, AD, D, U = genMatrixesLungsHeart()
+    N1 = A.shape[0]
+    N2 = AD.shape[0]
+
+    A = sp.csc_matrix(A).tocoo()
+    AD = sp.csc_matrix(AD).tocoo()
+    D = sp.csc_matrix(D).tocoo()
+    U = sp.csc_matrix(U).tocoo()
+
+    D_ = [D.copy()]
+    U_ = [U.copy()]
+
+    config = {}
+
+    config['n_nodes'] = [N1, N1, N1, N2, N2, N2]
+    A_ = [A.copy(), A.copy(), A.copy(), AD.copy(), AD.copy(), AD.copy()]
+    
+    A_t, D_t, U_t = ([scipy_to_torch_sparse(x).to(device) for x in X] for X in (A_, D_, U_))
+
+    config['latents'] = 64
+    config['inputsize'] = 1024
+
+    f = 32
+    config['filters'] = [2, f, f, f, f//2, f//2, f//2]
+    config['skip_features'] = f
+
+    hybrid = Hybrid(config.copy(), D_t, U_t, A_t).to(device)
+    hybrid.load_state_dict(torch.load("weights/bestMSE.pt"))
+    hybrid.eval()
+    
+    return hybrid
+
+
+def pad_to_square(img):
+    h, w = img.shape[:2]
+    
+    if h > w:
+        padw = (h - w) 
+        auxw = padw % 2
+        img = np.pad(img, ((0, 0), (padw//2, padw//2 + auxw)), 'constant')
+        
+        padh = 0
+        auxh = 0
+        
+    else:
+        padh = (w - h) 
+        auxh = padh % 2
+        img = np.pad(img, ((padh//2, padh//2 + auxh), (0, 0)), 'constant')
+
+        padw = 0
+        auxw = 0
+        
+    return img, (padh, padw, auxh, auxw)
+    
+
+def preprocess(input_img):
+    img, padding = pad_to_square(input_img)
+    
+    h, w = img.shape[:2]
+    if h != 1024 or w != 1024:
+        img = cv2.resize(img, (1024, 1024), interpolation = cv2.INTER_CUBIC)
+        
+    return img, (h, w, padding)
+
+    
+def segment(input_img):
+    input_img = cv2.imread(input_img, 0) / 255.0
+    
+    img, (h, w, padding) = preprocess(input_img)    
+    
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    hybrid = loadModel(device)
+    
+    data = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).to(device).float()
+    
+    with torch.no_grad():
+        output = hybrid(data)[0].cpu().numpy().reshape(-1, 2) * 1024
+       
+    return drawOnTop(img, output)
+
+
+if __name__ == "__main__":
+    demo = gr.Interface(segment, gr.Image(type="filepath"), "image")
+    demo.launch()

models/HybridGNet2IGSC.py

@@ -0,0 +1,200 @@
+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.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)
+        
+        # Genero la capa fully connected del 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]  
+                     
+        # Guardo las capas de convoluciones en grafo
+        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 forward(self, x):
+        self.mu, self.log_var, conv6, conv5 = self.encoder(x)
+
+        if self.training:
+            z = self.sampling(self.mu, self.log_var)
+        else:
+            z = self.mu
+            
+        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()) # Sin relu y sin bias
+        
+        return x, pos1, pos2

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,5 @@
+pytorch==2.0.1
+numpy==1.25.0
+opencv-python==4.8.0.74
+scipy==1.10.1
+pyg=2.3.0

utils/utils.py

@@ -0,0 +1,103 @@
+import numpy as np
+import scipy.sparse as sp
+import torch
+
+def scipy_to_torch_sparse(scp_matrix):
+    values = scp_matrix.data
+    indices = np.vstack((scp_matrix.row, scp_matrix.col))
+    i = torch.LongTensor(indices)
+    v = torch.FloatTensor(values)
+    shape = scp_matrix.shape
+
+    sparse_tensor = torch.sparse.FloatTensor(i, v, torch.Size(shape))
+    return sparse_tensor
+
+## Adjacency Matrix
+def mOrgan(N):
+    sub = np.zeros([N, N])
+    for i in range(0, N):
+        sub[i, i-1] = 1
+        sub[i, (i+1)%N] = 1
+    return sub
+
+## Downsampling Matrix
+def mOrganD(N):
+    N2 = int(np.ceil(N/2))
+    sub = np.zeros([N2, N])
+    
+    for i in range(0, N2):
+        if (2*i+1) == N:
+            sub[i, 2*i] = 1
+        else:
+            sub[i, 2*i] = 1/2
+            sub[i, 2*i+1] = 1/2
+            
+    return sub
+
+def mOrganU(N):
+    N2 = int(np.ceil(N/2))
+    sub = np.zeros([N, N2])
+    
+    for i in range(0, N):
+        if i % 2 == 0:
+            sub[i, i//2] = 1
+        else:
+            sub[i, i//2] = 1/2
+            sub[i, (i//2 + 1) % N2] = 1/2
+            
+    return sub
+
+def genMatrixesLungsHeart():       
+    RLUNG = 44
+    LLUNG = 50
+    HEART = 26
+    
+    Asub1 = mOrgan(RLUNG)
+    Asub2 = mOrgan(LLUNG)
+    Asub3 = mOrgan(HEART)
+    
+    ADsub1 = mOrgan(int(np.ceil(RLUNG / 2)))
+    ADsub2 = mOrgan(int(np.ceil(LLUNG / 2)))
+    ADsub3 = mOrgan(int(np.ceil(HEART / 2)))
+                    
+    Dsub1 = mOrganD(RLUNG)
+    Dsub2 = mOrganD(LLUNG)
+    Dsub3 = mOrganD(HEART)
+    
+    Usub1 = mOrganU(RLUNG)
+    Usub2 = mOrganU(LLUNG)
+    Usub3 = mOrganU(HEART)
+        
+    p1 = RLUNG
+    p2 = p1 + LLUNG
+    p3 = p2 + HEART
+    
+    p1_ = int(np.ceil(RLUNG / 2))
+    p2_ = p1_ + int(np.ceil(LLUNG / 2))
+    p3_ = p2_ + int(np.ceil(HEART / 2))
+    
+    A = np.zeros([p3, p3])
+    
+    A[:p1, :p1] = Asub1
+    A[p1:p2, p1:p2] = Asub2
+    A[p2:p3, p2:p3] = Asub3
+    
+    AD = np.zeros([p3_, p3_])
+    
+    AD[:p1_, :p1_] = ADsub1
+    AD[p1_:p2_, p1_:p2_] = ADsub2
+    AD[p2_:p3_, p2_:p3_] = ADsub3
+   
+    D = np.zeros([p3_, p3])
+    
+    D[:p1_, :p1] = Dsub1
+    D[p1_:p2_, p1:p2] = Dsub2
+    D[p2_:p3_, p2:p3] = Dsub3
+    
+    U = np.zeros([p3, p3_])
+    
+    U[:p1, :p1_] = Usub1
+    U[p1:p2, p1_:p2_] = Usub2
+    U[p2:p3, p2_:p3_] = Usub3
+    
+    return A, AD, D, U

weights/weights.pt

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