update

models/HybridGNet2IGSC.py

@@ -4,6 +4,110 @@ 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
+from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
+import json
+import scipy.sparse as sp
+
+
+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
 
 class EncoderConv(nn.Module):
     def __init__(self, latents = 64, hw = 32):
@@ -205,4 +309,198 @@ class Hybrid(nn.Module):
         else:
             z = self.mu
             
-        return self.decode(z, conv6, conv5) 
+        return self.decode(z, conv6, conv5) 
+    
+
+
+class HybridNoSkip(nn.Module):
+    def __init__(self, config, downsample_matrices, upsample_matrices, adjacency_matrices):
+        super(HybridNoSkip, self).__init__()
+               
+        hw = config['inputsize'] // 32
+        self.eval_sampling = config['eval_sampling']
+        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
+        
+        # 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])
+        
+        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.graphConv_up4 = ChebConv(self.filters[4], self.filters[3], self.K)
+        self.graphConv_up3 = ChebConv(self.filters[3], self.filters[2], self.K)
+        self.graphConv_up2 = ChebConv(self.filters[2], self.filters[1], self.K)
+        
+        ## Out layer: Sin bias, normalization ni relu
+        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):
+        mu, log_var, conv6, conv5 = self.encoder(x)
+        return mu, log_var, conv6, conv5
+    
+    def decode(self, z, conv6, conv5):
+        # Decode from latent space z to reconstruct x
+        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 = 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 = 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, None, None
+
+    def forward(self, x):
+        # Full forward pass: encode, sample (if training), then decode.
+        self.mu, self.log_var, conv6, conv5 = self.encode(x)
+        
+        if self.training:
+            z = self.sampling(self.mu, self.log_var)
+        else:
+            z = self.mu
+            
+        return self.decode(z, conv6, conv5)
+    
+
+class HybridGNetHF(nn.Module, PyTorchModelHubMixin):
+    repo_url = "https://github.com/mcosarinsky/CheXmask-U"
+    license = "mit"
+    pipeline_tag = "image-segmentation"
+
+    def __init__(self, latents=64, inputsize=1024, K=6, filters=None,
+                 skip_features=32, eval_sampling=True, use_skip=True,
+                 n_nodes=None, device="cpu", **kwargs):
+        super().__init__()
+        
+        self.device = device 
+
+        # Defaults
+        if filters is None:
+            filters = [2, 32, 32, 32, 16, 16, 16]
+        
+        # Save config
+        self.config = {
+            'latents': latents,
+            'inputsize': inputsize,
+            'K': K,
+            'filters': filters,
+            'skip_features': skip_features,
+            'eval_sampling': eval_sampling,
+            'use_skip': use_skip
+        }
+        self.config.update(kwargs)
+        self.use_skip = use_skip
+        
+        # ---- generate matrices ----
+        A, AD, D, U = genMatrixesLungsHeart()
+        N1, N2 = A.shape[0], AD.shape[0]
+        self.config['n_nodes'] = [N1, N1, N1, N2, N2, N2]
+
+        # ---- convert to sparse tensors and move to device ----
+        A_ = [sp.csc_matrix(A).tocoo() for _ in range(3)] + [sp.csc_matrix(AD).tocoo() for _ in range(3)]
+        D_ = [sp.csc_matrix(D).tocoo()]
+        U_ = [sp.csc_matrix(U).tocoo()]
+
+        self.A_t = [scipy_to_torch_sparse(x).to(self.device) for x in A_]
+        self.D_t = [scipy_to_torch_sparse(x).to(self.device) for x in D_]
+        self.U_t = [scipy_to_torch_sparse(x).to(self.device) for x in U_]
+
+        # ---- build model ----
+        if self.use_skip:
+            self.model = Hybrid(self.config, self.D_t, self.U_t, self.A_t)
+        else:
+            self.model = HybridNoSkip(self.config, self.D_t, self.U_t, self.A_t)
+
+        # move model parameters to device
+        self.model.to(self.device)
+
+    def forward(self, x):
+        return self.model(x)
+
+    # -----------------------------
+    # Dynamic from_pretrained from Hugging Face Hub ONLY
+    # -----------------------------
+    @classmethod
+    def from_pretrained(cls, repo_id, subfolder=None, device="cpu", **kwargs):
+        """
+        Loads model directly from Hugging Face Hub. Does NOT use local paths.
+        """
+        # Download config from Hub
+        config_file = hf_hub_download(
+            repo_id=repo_id,
+            filename="config.json",
+            subfolder=subfolder
+        )
+        with open(config_file, "r") as f:
+            config = json.load(f)
+        
+        # Merge any additional kwargs
+        config.update(kwargs)
+
+        # Dynamically compute n_nodes
+        A, AD, D, U = genMatrixesLungsHeart()
+        N1, N2 = A.shape[0], AD.shape[0]
+        config['n_nodes'] = [N1, N1, N1, N2, N2, N2]
+
+        # Instantiate model on desired device
+        model = cls(device=device, **config)
+
+        # Download weights from Hub
+        weights_path = hf_hub_download(
+            repo_id=repo_id,
+            filename="pytorch_model.bin",
+            subfolder=subfolder
+        )
+        state_dict = torch.load(weights_path, map_location=device)
+        if not next(iter(state_dict.keys())).startswith("model."):
+            state_dict = {f"model.{k}": v for k, v in state_dict.items()}
+        model.load_state_dict(state_dict)
+
+        return model

requirements.txt

@@ -4,3 +4,4 @@ opencv-python==4.8.0.74
 scipy==1.10.1
 torch_geometric==2.3.0
 torchvision==0.15.2
+huggingface_hub==1.2.3

utils/segmentation.py

@@ -9,7 +9,7 @@ from zipfile import ZipFile
 from .plotting import plot_side_by_side_comparison
 
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from models.HybridGNet2IGSC import Hybrid
+from models.HybridGNet2IGSC import HybridGNetHF
 
 hybrid = None
 
@@ -20,106 +20,6 @@ def seed_everything(seed=42):
     if torch.cuda.is_available():
         torch.cuda.manual_seed_all(seed)
 
-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
-
 def zip_files(files, output_name="complete_results.zip"):
     with ZipFile(output_name, "w") as zipObj:
         for file in files:
@@ -167,16 +67,11 @@ def removePreprocess(output, info):
 
 def loadModel(device):    
     global hybrid
-    A, AD, D, U = genMatrixesLungsHeart()
-    N1, N2 = A.shape[0], AD.shape[0]
-    A, AD, D, U = [sp.csc_matrix(x).tocoo() for x in [A, AD, D, U]]
-    D_, U_ = [D.copy()], [U.copy()]
-    A_ = [A.copy(), A.copy(), A.copy(), AD.copy(), AD.copy(), AD.copy()]
-    config = {'n_nodes':[N1,N1,N1,N2,N2,N2], 'latents':64, 'inputsize':1024,
-              'filters':[2,32,32,32,16,16,16], 'skip_features':32, 'eval_sampling':True}
-    A_t, D_t, U_t = ([scipy_to_torch_sparse(x).to(device) for x in X] for X in (A_,D_,U_))
-    hybrid = Hybrid(config.copy(), D_t, U_t, A_t).to(device)
-    hybrid.load_state_dict(torch.load("weights/weights.pt", map_location=device))
+    hybrid = HybridGNetHF.from_pretrained(
+            repo_id="mcosarinsky/CheXmask-U",
+            subfolder="v1_skip",
+            device=device
+            )
     hybrid.eval()
     return hybrid