gradio_code.py

from monai.utils import set_determinism, first
from monai.transforms import (
    AsDiscrete,
    AddChanneld,
    Compose,
    CropForegroundd,
    CenterSpatialCropd,
    EnsureChannelFirstd,
    LoadImaged,
    Orientationd,
    RandFlipd,
    RandCropByPosNegLabeld,
    RandCropByLabelClassesd,
    RandSpatialCropSamplesd,
    RandShiftIntensityd,
    RandZoomd,
    ScaleIntensityd,
    Spacingd,
    SpatialPadd,
    GaussianSmoothd,
    RandRotate90d,
    ToTensord,
    RandSpatialCropd,
    RandGaussianSmoothd,    
    RandGaussianSharpend,
    RandGaussianNoised,
)

from monai.config import print_config
from monai.losses import DiceCELoss
from monai.metrics import DiceMetric
from monai.utils import set_determinism, first

from monai.data import (
    DataLoader,
    Dataset,
    CacheDataset,
    load_decathlon_datalist,
    decollate_batch,
)

from sklearn.metrics import roc_auc_score

import monai
from monai.inferers import sliding_window_inference
from monai.transforms import Resize
from scipy.stats import mode

from skimage import measure

import torch
import pandas as pd
import os

import numpy as np
import nibabel as nib
import gradio as gd
import glob
from PIL import Image

import matplotlib.pyplot as plt
import base64
from io import BytesIO

device = "cpu"
model = monai.networks.nets.SwinUNETR(img_size=(128,128,64), 
                            in_channels=1, out_channels=2, 
                            feature_size=24).to(device)
model.load_state_dict(torch.load("./swinunetr_tmp.pth",map_location=torch.device('cpu')))


def save_as_obj(filename, vertices):
    with open(filename, 'w') as f:
        for vertex in vertices:
            f.write(f"v {vertex[0]} {vertex[1]} {vertex[2]}\n")


def visualize_3d_array(arr):
    # Get the indices where the pixel values are non-zero
    non_zero_indices = np.nonzero(arr)

    # Extract the x, y, and z coordinates
    x = non_zero_indices[0]
    y = non_zero_indices[1]
    z = non_zero_indices[2]

    # Get the shape of the input array
    x_size, y_size, z_size = arr.shape

    # Create a figure and axis for the 3D graph
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Normalize the values to range between 0 and 1
    normalized_values = arr[x, y, z] / np.max(arr)

    # Create a colormap based on the normalized values
    colormap = plt.cm.get_cmap('viridis')

    # Plot the non-zero points with colors based on the normalized values
    ax.scatter(x, y, z, c=colormap(normalized_values), marker='o')

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    # Set the axis limits to match the input array size
    ax.set_xlim([0, x_size])
    ax.set_ylim([0, y_size])
    ax.set_zlim([0, z_size])

    plt.show()

    buf = BytesIO()
    fig.savefig(buf, format='png')
    buf.seek(0)
    img_bytes = buf.read()

    img = Image.open(buf)
    img = img.convert("RGB")

    return img


def sepia(input_img, dropdown):

    img = nib.load(input_img.name)

    nib.save(img, "./output.nii.gz")


    plist = glob.glob("./output.nii.gz")

    valid_idx = np.arange(0,len(plist))
    data_dicts = [
        {
            "image1": os.path.join(plist[idx]),
        }
        for idx in valid_idx
    ]
    valid_Data =  data_dicts

    test_transforms = Compose(
        [
            LoadImaged(keys=["image1"]),
            EnsureChannelFirstd(keys=["image1"]),   
            ToTensord(keys=["image1"]),       
            ScaleIntensityd(
                    keys=["image1"],
                    minv=0.0,
                    maxv=1.0,
                ), 
        ]
    )

    test_ds = Dataset(
        data=valid_Data,
        transform=test_transforms,
    )

    test_loader = DataLoader(
        test_ds, batch_size=1, shuffle=False, 
    )

    with torch.no_grad():
        for step, batch in enumerate(test_loader):  
            val_inputs  = (batch["image1"]).to(device)
            sz = batch["image1"].shape[2:]
            R = Resize(spatial_size=(sz[0]*2,sz[1]*2,sz[2]*4))
            R0 = Resize(spatial_size=(sz[0],sz[1],sz[2]),mode="nearest")
            
            K = R(val_inputs[0])
            kz = K.shape[1:]
            x1 = int(0.5*(kz[0]-192))
            y1 = int(0.5*(kz[1]-192))
            z1 = int(0.5*(kz[2]-96))
            val_inputs = K[:,x1:x1+192,y1:y1+192,z1:z1+96].unsqueeze(0)

            val_outputs1 = sliding_window_inference(
                val_inputs, [128,128,64], 1, model, overlap=0.25,mode='gaussian')
            
            val_outputs = val_outputs1.softmax(1)
            
            r0 = val_outputs[0,1,:,:,:].detach().cpu()>.05        
            r = val_outputs[0,1,:,:,:].detach().cpu()>.25
            
            r0[val_inputs[0,0]==0]=0
            r[val_inputs[0,0]==0]=0

            SN = np.zeros_like(r)        
            L0, N0 = measure.label(r0, connectivity=3,return_num=True)        
            L, N1 = measure.label(r, connectivity=3,return_num=True)

            if N1>1:
                a,b=(mode(L[L>0]))
                SN[L==a[0]]=1
                L[L==a[0]]=0
                a,b=(mode(L[L>0]))
                SN[L==a[0]]=1 
            for sn in range(1,N0+1):
                if np.sum(SN*(L0==sn))==0:
                    L0[L0==sn]=0


    predicted_volume = str(np.rot90(np.sum(L0>0,axis=2)).sum())

    return visualize_3d_array(L0) , predicted_volume


demo = gd.Interface(sepia,["file"], [gd.outputs.Image('pil'), "text"], live=False,title="Experience", css="body {background-image: url('file=C:/Users/user/Downloads/background.jpg')}") 

demo.launch(share=True)