update app

app.py

@@ -1,307 +1,5 @@
-import numpy as np
 import gradio as gr
-import cv2 
-import matplotlib.pyplot as plt
-from PIL import Image
-from io import BytesIO
-from mpl_toolkits.axes_grid1 import make_axes_locatable
-
-from models.HybridGNet2IGSC import Hybrid 
-from utils.utils import scipy_to_torch_sparse, genMatrixesLungsHeart
-import scipy.sparse as sp
-import torch
-import pandas as pd
-from zipfile import ZipFile
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-hybrid = None
-
-def getDenseMask(landmarks, h, w):
-    
-    RL = landmarks[0:44]
-    LL = landmarks[44:94]
-    H = landmarks[94:]
-    
-    img = np.zeros([h, w], 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 getMasks(landmarks, h, w):
-    
-    RL = landmarks[0:44]
-    LL = landmarks[44:94]
-    H = landmarks[94:]
-    
-    RL = RL.reshape(-1, 1, 2).astype('int')
-    LL = LL.reshape(-1, 1, 2).astype('int')
-    H = H.reshape(-1, 1, 2).astype('int')
-    
-    RL_mask = np.zeros([h, w], dtype = 'uint8')
-    LL_mask = np.zeros([h, w], dtype = 'uint8')
-    H_mask = np.zeros([h, w], dtype = 'uint8')
-    
-    RL_mask = cv2.drawContours(RL_mask, [RL], -1, 255, -1)
-    LL_mask = cv2.drawContours(LL_mask, [LL], -1, 255, -1)
-    H_mask = cv2.drawContours(H_mask, [H], -1, 255, -1)
-
-    return RL_mask, LL_mask, H_mask
-
-def drawOnTop(img, landmarks, original_shape):
-    h, w = original_shape
-    output = getDenseMask(landmarks, h, w)
-    
-    image = np.zeros([h, w, 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)
-    
-    RL, LL, H = landmarks[0:44], landmarks[44:94], landmarks[94:]
-    
-    # Draw the landmarks as dots
-    
-    for l in RL:
-        image = cv2.circle(image, (int(l[0]), int(l[1])), 5, (1, 0, 1), -1)
-    for l in LL:
-        image = cv2.circle(image, (int(l[0]), int(l[1])), 5, (1, 0, 1), -1)
-    for l in H:
-        image = cv2.circle(image, (int(l[0]), int(l[1])), 5, (1, 1, 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()]
-    A_ = [A.copy(), A.copy(), A.copy(), AD.copy(), AD.copy(), AD.copy()]
-
-    config = {}
-    config['n_nodes'] = [N1, N1, N1, N2, N2, N2]
-    
-    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
-    config['eval_sampling'] = True
-
-    hybrid = Hybrid(config.copy(), D_t, U_t, A_t).to(device)
-    hybrid.load_state_dict(torch.load("weights/weights.pt", map_location=torch.device(device)))
-    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 removePreprocess(output, info):
-    """
-    output: np.array of shape (n_samples, N_landmarks, 2)
-    info: (h, w, padding)
-    """
-    h, w, padding = info
-    padh, padw, auxh, auxw = padding
-    
-    # Scale
-    if h != 1024 or w != 1024:
-        output = output * h
-    else:
-        output = output * 1024
-    
-    # Subtract padding
-    output[:, :, 0] = output[:, :, 0] - padw//2
-    output[:, :, 1] = output[:, :, 1] - padh//2
-    
-    return output
-
-
-def zip_files(files):
-    with ZipFile("complete_results.zip", "w") as zipObj:
-        for idx, file in enumerate(files):
-            zipObj.write(file, arcname=file.split("/")[-1])
-    return "complete_results.zip"
-
-
-def plot_landmarks_with_uncertainty(img, landmarks, uncertainty, figsize=(6,6)):
-    # Get dense mask as in drawOnTop
-    h, w = img.shape[:2]
-    dense_mask = getDenseMask(landmarks, h, w)
-
-    # Start with image overlay
-    overlay = np.zeros([h, w, 3])
-    overlay[:,:,0] = img + 0.3 * (dense_mask == 1).astype('float') - 0.1 * (dense_mask == 2).astype('float')
-    overlay[:,:,1] = img + 0.3 * (dense_mask == 2).astype('float') - 0.1 * (dense_mask == 1).astype('float')
-    overlay[:,:,2] = img - 0.1 * (dense_mask == 1).astype('float') - 0.2 * (dense_mask == 2).astype('float') 
-    overlay = np.clip(overlay, 0, 1)
-
-    # Plot
-    fig, ax = plt.subplots(figsize=figsize)
-    ax.imshow(overlay)
-
-    # Scatter landmarks colored by uncertainty
-    scatter = ax.scatter(
-        landmarks[:,0], landmarks[:,1], 
-        c=uncertainty, cmap='hot',
-        s=50, vmin=0, vmax=np.max(uncertainty)
-    )
-
-    # Colorbar
-    divider = make_axes_locatable(ax)
-    cax = divider.append_axes("right", size="5%", pad=0.05)
-    plt.colorbar(scatter, cax=cax, label='Node uncertainty')
-
-    ax.set_xlim(0, img.shape[1])
-    ax.set_ylim(img.shape[0], 0)
-    ax.axis('off')
-    fig.tight_layout()
-    return fig
-
-def segment(input_img, noise_std=0.0):
-    global hybrid, device
-    
-    if hybrid is None:
-        hybrid = loadModel(device)
-    
-    # ------------------ HANDLE SKETCH / INPAINT ------------------
-    if isinstance(input_img, dict):
-        original = input_img["image"].astype(np.float32)/255.0
-        mask = input_img["mask"]
-        if mask.ndim == 3:
-            mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
-        mask = mask.astype(np.float32)/255.0
-        mask = 1.0 - mask  # black strokes
-        input_img = np.minimum(original, mask)
-    else:
-        input_img = input_img.astype(np.float32)/255.0
-
-    # ------------------ ADD GAUSSIAN NOISE ------------------
-    if noise_std > 0:
-        noise = np.random.normal(0, noise_std, input_img.shape)
-        input_img = np.clip(input_img + noise, 0.0, 1.0)
-
-    # ------------------ PREPROCESS & PREDICT ------------------
-    original_shape = input_img.shape[:2]
-    img, (h, w, padding) = preprocess(input_img)    
-    data = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).to(device).float()
-    n_samples = 100
-
-    with torch.no_grad():
-        mu, log_var, conv6, conv5 = hybrid.encode(data)
-        latent_var = np.exp(log_var.cpu().numpy())
-
-         # Sample N latent vectors to decode
-        zs = [hybrid.sampling(mu, log_var) for _ in range(n_samples)]
-        z_exp = torch.stack(zs, dim=0)
-
-        # Expand skip connections to match batch size
-        conv6_exp = conv6.repeat(n_samples, 1, 1, 1)
-        conv5_exp = conv5.repeat(n_samples, 1, 1, 1)
-
-        # Decode in batch
-        output, _, _ = hybrid.decode(z_exp, conv6_exp, conv5_exp)
-        output = output.cpu().numpy().reshape(n_samples, -1, 2)
-        output = removePreprocess(output, (h, w, padding)).astype('int')
-
-        # Compute mean and std per node
-        means = np.mean(output, axis=0)
-        stds = np.std(output, axis=0)
-            
-    # ------------------ SAVE LANDMARKS & MASKS ------------------
-    RL = means[0:44]
-    LL = means[44:94]
-    H = means[94:]
-    
-    np.savetxt("tmp/RL_landmarks.txt", RL, delimiter=" ", fmt="%d")
-    np.savetxt("tmp/LL_landmarks.txt", LL, delimiter=" ", fmt="%d")
-    np.savetxt("tmp/H_landmarks.txt", H, delimiter=" ", fmt="%d")
-    
-    RL_mask, LL_mask, H_mask = getMasks(means, original_shape[0], original_shape[1])
-    cv2.imwrite("tmp/RL_mask.png", RL_mask)
-    cv2.imwrite("tmp/LL_mask.png", LL_mask)
-    cv2.imwrite("tmp/H_mask.png", H_mask)
-    
-    RL_std = stds[0:44]
-    LL_std = stds[44:94]
-    H_std  = stds[94:]
-
-    # Save as text files
-    np.savetxt("tmp/RL_std.txt", RL_std, delimiter=" ", fmt="%.4f")
-    np.savetxt("tmp/LL_std.txt", LL_std, delimiter=" ", fmt="%.4f")
-    np.savetxt("tmp/H_std.txt", H_std, delimiter=" ", fmt="%.4f")
-
-    zipf = zip_files([
-        "tmp/RL_landmarks.txt", "tmp/LL_landmarks.txt", "tmp/H_landmarks.txt",
-        "tmp/RL_mask.png", "tmp/LL_mask.png", "tmp/H_mask.png",
-        "tmp/RL_std.txt", "tmp/LL_std.txt", "tmp/H_std.txt"
-    ])
-    
-    # ------------------ RANDOM UNCERTAINTY ------------------
-    node_uncertainty = np.mean(stds, axis=1) 
-    fig = plot_landmarks_with_uncertainty(input_img, means, node_uncertainty)
-    
-    output_path = "tmp/segmentation_with_uncertainty.png"
-    fig.savefig(output_path, format="png", dpi=150)
-    plt.close(fig)
-    
-    return output_path, [
-        "tmp/RL_landmarks.txt", "tmp/LL_landmarks.txt", "tmp/H_landmarks.txt",
-        "tmp/RL_mask.png", "tmp/LL_mask.png", "tmp/H_mask.png",
-        "tmp/RL_std.txt", "tmp/LL_std.txt", "tmp/H_std.txt",
-        zipf
-    ]
-
-
+from utils.segmentation import segment
 
 # ------------------------- GRADIO -------------------------
 if __name__ == "__main__":    
@@ -310,19 +8,18 @@ if __name__ == "__main__":
 
         with gr.Tab("Segment Image"):
             with gr.Row():
-                with gr.Column():
+                with gr.Column(scale=1):
                     image_input = gr.Image(
                         type="numpy", 
                         tool="sketch", 
                         image_mode="L", 
-                        height=512, 
-                        shape=(512, 512)
+                        height=450, 
                     )
                     
                     noise_slider = gr.Slider(
                         label="Gaussian Noise Std Dev",
                         minimum=0.0,
-                        maximum=0.25,  # max = strong noise / blurry
+                        maximum=0.25,
                         step=0.01,
                         value=0.0
                     )
@@ -336,9 +33,8 @@ if __name__ == "__main__":
                         'utils/example3.png','utils/example4.jpg'
                     ])
                     
-                with gr.Column():
-                    #image_output = gr.Image(type="numpy", height=750)
-                    image_output = gr.Image(type="filepath", height=512)
+                with gr.Column(scale=2):
+                    image_output = gr.Image(type="filepath", height=450)
                     results = gr.File()       
                     
         gr.Markdown("""

utils/plotting.py

@@ -0,0 +1,77 @@
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+def getDenseMask(landmarks, h, w):
+    RL, LL, H = landmarks[:44], landmarks[44:94], landmarks[94:]
+    img = np.zeros([h, w], 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, original_shape):
+    h, w = original_shape
+    output = getDenseMask(landmarks, h, w)
+    image = np.zeros([h,w,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)
+    RL, LL, H = landmarks[:44], landmarks[44:94], landmarks[94:]
+    for l in RL: image = cv2.circle(image,(int(l[0]),int(l[1])),5,(1,0,1),-1)
+    for l in LL: image = cv2.circle(image,(int(l[0]),int(l[1])),5,(1,0,1),-1)
+    for l in H: image = cv2.circle(image,(int(l[0]),int(l[1])),5,(1,1,0),-1)
+    return image
+
+def create_overlay(img, landmarks):
+    h, w = img.shape[:2]
+    dense_mask = getDenseMask(landmarks, h, w)
+    overlay = np.zeros([h, w, 3])
+    
+    overlay[:,:,0] = img + 0.3 * (dense_mask == 1).astype('float') - 0.1 * (dense_mask == 2).astype('float')
+    overlay[:,:,1] = img + 0.3 * (dense_mask == 2).astype('float') - 0.1 * (dense_mask == 1).astype('float')
+    overlay[:,:,2] = img - 0.1 * (dense_mask == 1).astype('float') - 0.2 * (dense_mask == 2).astype('float')
+    overlay = np.clip(overlay, 0, 1)
+    
+    return overlay
+
+def plot_side_by_side_comparison(img_orig, means_orig, uncertainty_orig, img_corr, means_corr, uncertainty_corr):
+
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 7))
+    
+    fig.set_constrained_layout(True)
+
+    vmax = max(np.max(np.mean(uncertainty_orig, axis=1)), np.max(np.mean(uncertainty_corr, axis=1)))
+
+    # --- Original ---
+    overlay_orig = create_overlay(img_orig, means_orig)
+    ax1.imshow(overlay_orig)
+    scatter1 = ax1.scatter(
+        means_orig[:, 0], means_orig[:, 1],
+        c=np.mean(uncertainty_orig, axis=1),
+        cmap='hot', s=50, vmin=0, vmax=vmax
+    )
+    ax1.set_title("Original", fontsize=16, pad=10)
+    ax1.axis('off')
+
+    # --- Corrupted ---
+    overlay_corr = create_overlay(img_corr, means_corr)
+    ax2.imshow(overlay_corr)
+    scatter2 = ax2.scatter(
+        means_corr[:, 0], means_corr[:, 1],
+        c=np.mean(uncertainty_corr, axis=1),
+        cmap='hot', s=50, vmin=0, vmax=vmax
+    )
+    ax2.set_title("Corrupted", fontsize=16, pad=10)
+    ax2.axis('off')
+
+    # Shared colorbar
+    cbar = fig.colorbar(scatter2, ax=[ax1, ax2], fraction=0.046, pad=0.01, shrink=0.85)
+    cbar.ax.tick_params(labelsize=10)
+
+    return fig

utils/segmentation.py

@@ -0,0 +1,259 @@
+import numpy as np
+import cv2
+import torch
+import scipy.sparse as sp
+import sys
+import os
+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
+
+hybrid = None
+
+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:
+            zipObj.write(file, arcname=file.split("/")[-1])
+    return output_name
+
+def getMasks(landmarks, h, w):
+    RL, LL, H = landmarks[:44], landmarks[44:94], landmarks[94:]
+    RL_mask, LL_mask, H_mask = [np.zeros([h, w], dtype='uint8') for _ in range(3)]
+    RL_mask = cv2.drawContours(RL_mask, [RL.reshape(-1,1,2).astype('int')], -1, 255, -1)
+    LL_mask = cv2.drawContours(LL_mask, [LL.reshape(-1,1,2).astype('int')], -1, 255, -1)
+    H_mask = cv2.drawContours(H_mask, [H.reshape(-1,1,2).astype('int')], -1, 255, -1)
+    return RL_mask, LL_mask, H_mask
+
+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')
+        return img, (0, padw, 0, auxw)
+    else:
+        padh = w - h
+        auxh = padh % 2
+        img = np.pad(img, ((padh//2, padh//2+auxh),(0,0)), 'constant')
+        return img, (padh, 0, auxh, 0)
+
+def preprocess(img):
+    img, padding = pad_to_square(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 removePreprocess(output, info):
+    h, w, padding = info
+    padh, padw, auxh, auxw = padding
+    if h != 1024 or w != 1024:
+        output = output * h
+    else:
+        output = output * 1024
+    output[:,:,0] -= padw//2
+    output[:,:,1] -= padh//2
+    return output
+
+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.eval()
+    return hybrid
+
+def predict_landmarks(img, n_samples=100):
+    global hybrid
+    img_proc, (h, w, padding) = preprocess(img)
+    data = torch.from_numpy(img_proc).unsqueeze(0).unsqueeze(0).to(next(hybrid.parameters()).device).float()
+    with torch.no_grad():
+        mu, log_var, conv6, conv5 = hybrid.encode(data)
+        zs = [hybrid.sampling(mu, log_var) for _ in range(n_samples)]
+        z_exp = torch.stack(zs, dim=0)
+        conv6_exp, conv5_exp = conv6.repeat(n_samples,1,1,1), conv5.repeat(n_samples,1,1,1)
+        output, _, _ = hybrid.decode(z_exp, conv6_exp, conv5_exp)
+        output = output.cpu().numpy().reshape(n_samples,-1,2)
+        output = removePreprocess(output, (h,w,padding)).astype('int')
+        means, stds = np.mean(output,axis=0), np.std(output,axis=0)
+    return means, stds
+
+
+def segment(input_img, noise_std=0.0):
+    """
+    input_img: dict with keys "image" (numpy array) and optionally "mask"
+    noise_std: standard deviation of Gaussian noise to add for robustness
+    Returns: path to comparison figure, list of saved files
+    """
+    global hybrid
+
+    if hybrid is None:
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        hybrid = loadModel(device)
+
+    # Original image and corrupted version
+    img_orig = input_img["image"].astype(np.float32) / 255.0
+    mask = input_img.get("mask", None)
+    if mask is not None:
+        mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY).astype(np.float32) / 255.0
+        mask = 1.0 - mask
+        img_corr = np.minimum(img_orig, mask)
+    else:
+        img_corr = img_orig.copy()
+
+    if noise_std > 0:
+        noise = np.random.normal(0, noise_std, img_corr.shape)
+        img_corr = np.clip(img_corr + noise, 0.0, 1.0)
+
+    # Predict landmarks
+    means_orig, stds_orig = predict_landmarks(img_orig)
+    means_corr, stds_corr = predict_landmarks(img_corr)
+
+    # Save landmarks and masks
+    os.makedirs("tmp", exist_ok=True)
+
+    RL, LL, H = means_orig[:44], means_orig[44:94], means_orig[94:]
+    np.savetxt("tmp/RL_landmarks.txt", RL, delimiter=" ", fmt="%d")
+    np.savetxt("tmp/LL_landmarks.txt", LL, delimiter=" ", fmt="%d")
+    np.savetxt("tmp/H_landmarks.txt", H, delimiter=" ", fmt="%d")
+
+    RL_mask, LL_mask, H_mask = getMasks(means_orig, img_orig.shape[0], img_orig.shape[1])
+    cv2.imwrite("tmp/RL_mask.png", RL_mask)
+    cv2.imwrite("tmp/LL_mask.png", LL_mask)
+    cv2.imwrite("tmp/H_mask.png", H_mask)
+
+    RL_std, LL_std, H_std = stds_orig[:44], stds_orig[44:94], stds_orig[94:]
+    np.savetxt("tmp/RL_std.txt", RL_std, delimiter=" ", fmt="%.4f")
+    np.savetxt("tmp/LL_std.txt", LL_std, delimiter=" ", fmt="%.4f")
+    np.savetxt("tmp/H_std.txt", H_std, delimiter=" ", fmt="%.4f")
+
+    zipf = zip_files([
+        "tmp/RL_landmarks.txt","tmp/LL_landmarks.txt","tmp/H_landmarks.txt",
+        "tmp/RL_mask.png","tmp/LL_mask.png","tmp/H_mask.png",
+        "tmp/RL_std.txt","tmp/LL_std.txt","tmp/H_std.txt"
+    ])
+
+    # Optional: plot side-by-side comparison
+    fig = plot_side_by_side_comparison(img_orig, means_orig, stds_orig, img_corr, means_corr, stds_corr)
+    output_path = "tmp/segmentation_comparison.png"
+    fig.savefig(output_path, dpi=300)
+    import matplotlib.pyplot as plt
+    plt.close(fig)
+
+    saved_files = [
+        "tmp/RL_landmarks.txt","tmp/LL_landmarks.txt","tmp/H_landmarks.txt",
+        "tmp/RL_mask.png","tmp/LL_mask.png","tmp/H_mask.png",
+        "tmp/RL_std.txt","tmp/LL_std.txt","tmp/H_std.txt",
+        zipf
+    ]
+
+    return output_path, saved_files

utils/utils.py

@@ -1,103 +0,0 @@
-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