gradio_code2.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
import plotly.graph_objs as go

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 view_3d_dynamic(nft):
    # 0이 아닌 값의 인덱스 추출
    nonzero_indices = np.nonzero(nft)

    # 인덱스 추출
    x_indices = nonzero_indices[0]
    y_indices = nonzero_indices[1]
    z_indices = nonzero_indices[2]

    # 인덱스에 해당하는 값을 추출
    values = nft[x_indices, y_indices, z_indices]

    # 시각화를 위한 Scatter3d 객체 생성
    scatter = go.Scatter3d(
        x=x_indices,
        y=y_indices,
        z=z_indices,
        mode='markers',
        marker=dict(
        size=2,
        color=values,
        colorscale='Viridis',
        opacity=0.5
        )
    )

    # 레이아웃 설정
    layout = go.Layout(
        scene=dict(
            xaxis=dict(
                range=[0, nft.shape[0]]
            ),
            yaxis=dict(
                range=[0, nft.shape[1]] 
            ),
            zaxis=dict(
                range=[0, nft.shape[2]]
            )
        )
    )

    # 그래프 생성
    fig = go.Figure(data=[scatter], layout=layout)
    fig.show()
    return 


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):

    non_zero_indices = np.nonzero(arr)

    x = non_zero_indices[0]
    y = non_zero_indices[1]
    z = non_zero_indices[2]

    x_size, y_size, z_size = arr.shape

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    normalized_values = arr[x, y, z] / np.max(arr)

    colormap = plt.cm.get_cmap('viridis')

    ax.scatter(x, y, z, c=colormap(normalized_values), marker='o')

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

    ax.set_xlim([0, x_size])
    ax.set_ylim([0, y_size])
    ax.set_zlim([0, z_size])

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

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

    view_3d_dynamic(arr) # 움직이는 창

    return img

def show_val_hist(arr):

    plt.figure()
    data = []
    data_count = {}

    for d in arr.flatten():
        if d != 0 :
            data.append(d)
            if d in data_count:
                data_count[d] += 1
            else:
                data_count[d] = 1

    plt.hist(data, bins=np.max(arr), edgecolor='black')

    buf = BytesIO()
    buf.truncate(0) # 버퍼 초기화
    plt.savefig(buf, format='png')
    buf.seek(0)

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

    sorted_dict = sorted(data_count.items(), key=lambda x: x[1], reverse=True)

    key1 , vol1 = sorted_dict[0]
    key2, vol2 = sorted_dict[1]

    for i in range(len(data)):
        if key1 == data[i]:
            left, right = vol1, vol2
        elif key2 == data[i]:
            left, right = vol2, vol1

    if left > right:
        rl_sub = left - right
    else:
        rl_sub = right - left

    return img, 'left: ' + str(left) + ' | right: ' + str(right) + ' | sub: ' + str(rl_sub)


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


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

    return visualize_3d_array(L0) , hist, volumes ,predicted_volume


demo = gd.Interface(sepia,["file"], [gd.outputs.Image('pil'), gd.outputs.Image('pil'), "text","text"], live=False,title="Correlation between Substantia Nigra and Schizophrenia") 

demo.launch(share=True)