demo / app.py
amazinghaha's picture
Update app.py
6d40932 verified
raw
history blame
22.3 kB
import gradio as gr
import torch
import os
import numpy as np
import SimpleITK as sitk
from scipy.ndimage import zoom
import pickle
from model.Vision_Transformer_with_mask import vit_base_patch16_224,Attention,CrossAttention,Attention_ori
from model.CoordAttention import *
from typing import Tuple, Type
from torch import Tensor, nn
from xgboost import XGBClassifier
from sklearn.ensemble import AdaBoostClassifier
import umap
def load_from_pkl(load_path):
data_input = open(load_path, 'rb')
read_data = pickle.load(data_input)
data_input.close()
return read_data
class MLP_att_out(nn.Module):
def __init__(self, input_dim, inter_dim=None, output_dim=None, activation="relu", drop=0.0):
super().__init__()
self.input_dim = input_dim
self.inter_dim = inter_dim
self.output_dim = output_dim
if inter_dim is None: self.inter_dim=input_dim
if output_dim is None: self.output_dim=input_dim
self.linear1 = nn.Linear(self.input_dim, self.inter_dim)
self.activation = self._get_activation_fn(activation)
self.dropout3 = nn.Dropout(drop)
self.linear2 = nn.Linear(self.inter_dim, self.output_dim)
self.dropout4 = nn.Dropout(drop)
self.norm3 = nn.LayerNorm(self.output_dim)
def forward(self, x):
x = self.linear2(self.dropout3(self.activation(self.linear1(x))))
x = x + self.dropout4(x)
x = self.norm3(x)
return x
def _get_activation_fn(self, activation):
"""Return an activation function given a string"""
if activation == "relu":
return F.relu
if activation == "gelu":
return F.gelu
if activation == "glu":
return F.glu
raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
class MLPBlock(nn.Module):
def __init__(
self,
embedding_dim: int,
mlp_dim: int,
act: Type[nn.Module] = nn.GELU,
) -> None:
super().__init__()
self.lin1 = nn.Linear(embedding_dim, mlp_dim)
self.lin2 = nn.Linear(mlp_dim, embedding_dim)
self.act = act()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.lin2(self.act(self.lin1(x)))
class FusionAttentionBlock(nn.Module):
def __init__(
self,
embedding_dim: int,
num_heads: int,
mlp_dim: int = 2048,
activation: Type[nn.Module] = nn.ReLU,
) -> None:
"""
A transformer block with four layers: (1) self-attention of sparse
inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp
block on sparse inputs, and (4) cross attention of dense inputs to sparse
inputs.
Arguments:
embedding_dim (int): the channel dimension of the embeddings
num_heads (int): the number of heads in the attention layers
mlp_dim (int): the hidden dimension of the mlp block
activation (nn.Module): the activation of the mlp block
"""
super().__init__()
self.self_attn = Attention_ori(embedding_dim, num_heads)
self.norm1 = nn.LayerNorm(embedding_dim)
self.cross_attn_mask_to_image = CrossAttention(dim=embedding_dim, num_heads=num_heads)
self.norm2 = nn.LayerNorm(embedding_dim)
self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
self.norm3 = nn.LayerNorm(embedding_dim)
self.norm4 = nn.LayerNorm(embedding_dim)
self.cross_attn_image_to_mask = CrossAttention(dim=embedding_dim, num_heads=num_heads)
def forward(self, img_emb: Tensor, mask_emb: Tensor, atten_mask: Tensor) -> Tuple[ Tensor]:
# Self attention block #最开始的时候 queries=query_pe
#queries: Tensor, keys: Tensor
queries = mask_emb
attn_out = self.self_attn(queries) #小图
queries = attn_out
#queries = queries + attn_out
queries = self.norm1(queries)
# Cross attention block, mask attending to image embedding
q = queries #1,5,256
k = img_emb # v是值,因此用keys?
input_x = torch.cat((q, k), dim=1) # 2 50 768
attn_out = self.cross_attn_mask_to_image(input_x) #TODO 要不要mask呢 交叉的时候 先不用试试
queries = queries + attn_out
queries = self.norm2(queries)
# MLP block
mlp_out = self.mlp(queries)
queries = queries + mlp_out
queries = self.norm3(queries)
# Cross attention block, image embedding attending to tokens
q = img_emb
k = queries
input_x = torch.cat((q, k), dim=1)
attn_out = self.cross_attn_image_to_mask(input_x)
img_emb = img_emb + attn_out
img_emb = self.norm4(img_emb)
return img_emb
class my_model7(nn.Module):
'''不用mask的版本
concate 部分 加了nor 加 attention
attention 用不一样的方法
'''
def __init__(self, pretrained=False,num_classes=3,in_chans=1,img_size=224, **kwargs):
super().__init__()
self.backboon1 = vit_base_patch16_224(pretrained=False,in_chans=in_chans, as_backbone=True,img_size=img_size)
if pretrained:
pre_train_model = timm.create_model('vit_base_patch16_224', pretrained=True, in_chans=in_chans, num_classes=3)
self.backboon1 = load_weights(self.backboon1, pre_train_model.state_dict())
#self.backboon2 = vit_base_patch32_224(pretrained=False,as_backbone=True) #TODO 同一个网络共享参数/不共享参数/patch不同网络
self.self_atten_img = Attention_ori(dim= self.backboon1.embed_dim, num_heads=self.backboon1.num_heads)
#self.self_atten_mask = Attention(dim=self.backboon1.embed_dim, num_heads=self.backboon1.num_heads)
self.self_atten_mask = Attention_ori(dim=self.backboon1.embed_dim, num_heads=self.backboon1.num_heads)
self.cross_atten = FusionAttentionBlock(embedding_dim=self.backboon1.embed_dim, num_heads=self.backboon1.num_heads)
#self.external_attention = ExternalAttention(d_model=2304,S=8)
self.mlp = MLP_att_out(input_dim=self.backboon1.embed_dim * 3, output_dim=self.backboon1.embed_dim)
self.attention = CoordAtt(1,1,1)
self.norm1 = nn.LayerNorm(self.backboon1.embed_dim)
self.norm2 = nn.LayerNorm(self.backboon1.embed_dim)
self.norm3 = nn.LayerNorm(self.backboon1.embed_dim)
self.avgpool = nn.AdaptiveAvgPool1d(1)
self.head = nn.Linear(self.backboon1.embed_dim*3, num_classes) if num_classes > 0 else nn.Identity()
#self.head = nn.Linear(196, num_classes) if num_classes > 0 else nn.Identity()
def forward(self, img, mask):
x1 = self.backboon1(torch.cat((img, torch.zeros_like(img)), dim=1)) #TODO 是否用同一模型 还是不同 中间是否融合多尺度
x2 = self.backboon1(torch.cat((img*mask, torch.zeros_like(img)), dim=1)) #输出经过了归一化层 #小图
#自注意力+残差
x2_atten_mask = self.backboon1.atten_mask
x1_atten = self.self_atten_img(x1)
x2_atten = self.self_atten_mask(x2)
x1_out = self.norm1((x1 + x1_atten))
x2_out = self.norm2((x2 + x2_atten))
#交叉注意力
corss_out = self.norm3(self.cross_atten(x1, x2, x2_atten_mask))
#得到输出特征
out = torch.concat((x1_out, corss_out, x2_out), dim=2).permute(0, 2, 1)#12 2304 196
out = self.attention(out) #12 2304 196
#out_ = out.permute(0, 2, 1)
#out = self.mlp(out) # mlp #特征融合 2 196 768
# out = self.norm1(out) #这个好像不用 好像可以删掉
out = self.avgpool(out) # B C 1
out = torch.flatten(out, 1)
out = self.head(out)
return out
Image_3D = None
Current_name = None
ALL_message = load_from_pkl('./label0601.pkl')
ALL_message2 = load_from_pkl('./all_data_label.pkl')
a = ALL_message2['train']
a.update(ALL_message2['val'])
a.update(ALL_message2['test'])
ALL_message2 = a
LC_model_Paht = './train_ADA_1.pkl'
LC_model = load_from_pkl(LC_model_Paht)['model'][0]
TF_model_Paht = './tf_model.pkl'
TF_model = load_from_pkl(TF_model_Paht)['model']
DR_model = load_from_pkl(TF_model_Paht)['dr']
Model_Paht = './model_epoch120.pth.tar'
checkpoint = torch.load(Model_Paht, map_location='cpu')
classnet = my_model7(pretrained=False,num_classes=3,in_chans=1, img_size=224)
classnet.load_state_dict(checkpoint['model_dict'])
def resize3D(img, aimsize, order=3):
"""
:param img: 3D array
:param aimsize: list, one or three elements, like [256], or [256,56,56]
:return:
"""
_shape = img.shape
if len(aimsize) == 1:
aimsize = [aimsize[0] for _ in range(3)]
if aimsize[0] is None:
return zoom(img, (1, aimsize[1] / _shape[1], aimsize[2] / _shape[2]), order=order) # resample for cube_size
if aimsize[1] is None:
return zoom(img, (aimsize[0] / _shape[0], 1, aimsize[2] / _shape[2]), order=order) # resample for cube_size
if aimsize[2] is None:
return zoom(img, (aimsize[0] / _shape[0], aimsize[1] / _shape[1], 1), order=order) # resample for cube_size
return zoom(img, (aimsize[0] / _shape[0], aimsize[1] / _shape[1], aimsize[2] / _shape[2]),
order=order) # resample for cube_size
def get_lc():
global Current_name
lc_min = np.array([17,1,0,1,1,1,1,1 , 1 , 1])
lc_max = np.array([96 ,2, 3 ,2, 2,2 , 2 ,2 ,2 ,4])
lc_key = ['age', 'sex', 'time', 'postpartum', 'traumatism', 'diabetes', 'high_blood_pressure', 'cerebral_infarction', 'postoperation']
lc_all = [ALL_message2[Current_name][ii] for ii in lc_key]
site_ = Current_name.split('_',1)[-1]
if site_ == 'A_L': lc_all.append(1)
elif site_ == 'A_R': lc_all.append(2)
elif site_ == 'B_L': lc_all.append(3)
elif site_ == 'B_R': lc_all.append(4)
else: pass
lc_all = (np.array(lc_all)-lc_min)/(lc_max-lc_min+ 1e-12)
a = 5
return lc_all
def inference():
global Image_small_3D
global ROI_small_3D
model = classnet
data_3d = Image_small_3D
lc_data = get_lc()
lc_data = np.expand_dims(lc_data, axis=0)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.eval()
try:
#影像模型
with torch.no_grad():
all_probs = np.empty((0, 3))
for ii in range(0, data_3d.shape[1]):
data = torch.from_numpy(data_3d[:,ii,:])
roi = torch.from_numpy(ROI_small_3D[:,ii,:].astype(np.int8))
image = torch.unsqueeze(data, 0)
roi = torch.unsqueeze(torch.unsqueeze(roi, 0),0).to(device).float()
patch_data = torch.unsqueeze(image, 0).to(device).float() # (N, C_{in}, D_{in}, H_{in}, W_{in})
# Pre : Prediction Result
pre_probs = model(patch_data,roi)
pre_probs = torch.nn.functional.softmax(pre_probs, dim=1)
all_probs = np.concatenate((all_probs, pre_probs.cpu().numpy()), axis=0)
dl_prob = np.mean(all_probs, axis=0)
dl_prob = np.expand_dims(dl_prob, axis=0)
lc_prob = LC_model.predict_proba(lc_data)
feature = DR_model.transform(np.concatenate([dl_prob, lc_prob], axis=1))
final_p = TF_model.predict_proba(feature)
final_p = np.round(final_p[0], decimals=2)
return {'急性期': final_p[0], '亚急性期': final_p[1], '慢性期': final_p[2]}
except:
return ' '
def get_Image_reslice(input_file):
'''得到图像 返回随即层'''
global Image_3D
global Current_name
global Input_File
if isinstance(input_file, str):
input_file = input_file
else:
input_file = input_file.name
Input_File = input_file
print(input_file)
Image_3D = sitk.GetArrayFromImage(sitk.ReadImage(input_file))
Current_name = input_file.split(os.sep)[-1].split('.')[0].rsplit('_', 1)[0]
Image_3D = (np.max(Image_3D) - Image_3D) / (np.max(Image_3D) - np.min(Image_3D))
random_z = np.random.randint(0, Image_3D.shape[0])
image_slice_z = Image_3D[random_z, :, :]
random_y = np.random.randint(0, Image_3D.shape[1])
image_slice_y = Image_3D[:, random_y, :]
random_x = np.random.randint(0, Image_3D.shape[2])
image_slice_x = Image_3D[:, :, random_x]
# return zoom(image_slice_z, (10 / image_slice_z.shape[0], 10 / image_slice_z.shape[1]), order=3) , \
# zoom(image_slice_y, (10 / image_slice_y.shape[0], 10 / image_slice_y.shape[1]), order=3), \
# zoom(image_slice_x, (10 / image_slice_x.shape[0], 10 / image_slice_x.shape[1]), order=3)
return image_slice_z, \
image_slice_y, \
image_slice_x, random_z, random_y, random_x, '影像数据加载成功'
def get_ROI(input_file):
'''得到图像 返回随即层'''
global ROI_3D
if isinstance(input_file, str):
input_file = input_file
else:
input_file = input_file.name
Image_3D = sitk.GetArrayFromImage(sitk.ReadImage(input_file))
ROI_3D = Image_3D
unique_elements = np.unique(ROI_3D)
a = 5
if np.where(unique_elements>1)[0]:
return '这个数据没有经过二值化'
else:
return '感兴趣区域加载成功'
def change_image_slice_x(slice):
image_slice = Image_3D[:, :, slice - 1]
cut_thre = np.percentile(image_slice, 99.9) # 直方图99.9%右侧值不要
image_slice[image_slice >= cut_thre] = cut_thre
image_slice = (((np.max(image_slice) -image_slice)/(np.max(image_slice) - np.min(image_slice)))*255).astype(np.int16)
a = 5
return image_slice
def change_image_slice_y(slice):
image_slice = Image_3D[:, slice - 1, :]
cut_thre = np.percentile(image_slice, 99.9) # 直方图99.9%右侧值不要
image_slice[image_slice >= cut_thre] = cut_thre
image_slice = (((np.max(image_slice) - image_slice) / (np.max(image_slice) - np.min(image_slice))) * 255).astype(
np.int16)
return image_slice
def change_image_slice_z(slice):
image_slice = Image_3D[slice - 1, :, :]
cut_thre = np.percentile(image_slice, 99.9) # 直方图99.9%右侧值不要
image_slice[image_slice >= cut_thre] = cut_thre
image_slice = (((np.max(image_slice) - image_slice) / (np.max(image_slice) - np.min(image_slice))) * 255).astype(np.int16)
return image_slice
def get_medical_message():
global Current_name
if Current_name == None:
return '请先加载数据', ' '
else:
past = ALL_message[Current_name]['past']
now = ALL_message[Current_name]['now']
return past, now
def clear_all():
global Image_3D
global Current_name
Current_name = None
Image_3D = None
return np.ones((10, 10)), np.ones((10, 10)), np.ones((10, 10)), '', '', ' ',"尚未进行预处理 请先预处理再按“分期结果”按钮","尚未加载影像数据","尚未加载感兴趣区域"
def get_box(mask):
"""
:param mask: array,输入金标准图像
:return:
"""
# 得到boxx坐标
# 计算得到bbox,形式为[dim0min, dim0max, dim1min, dim1max, dim2min, dim2max]
indexx = np.where(mask > 0.) # 返回坐标,几维就是几组坐标,坐标纵向看
dim0min, dim0max, dim1min, dim1max, dim2min, dim2max = [np.min(indexx[0]), np.max(indexx[0]),
np.min(indexx[1]), np.max(indexx[1]),
np.min(indexx[2]), np.max(indexx[2])]
bbox = [dim0min, dim0max, dim1min, dim1max, dim2min, dim2max]
return bbox
def arry_crop_3D(img,mask,ex_pix):
'''
得到小图,并外扩
:param img array 3D
:param mask array
:param ex_pix: list [a,b,c] 向两侧各自外扩多少 维度顺序与输入一致
:param z_waikuo:z轴是否外扩,默认第一维 务必提前确认 !!
'''
if len(ex_pix)==1:
ex_pix=[ex_pix[0] for _ in range(3)]
elif len(ex_pix) == 2:
print('如果z轴不外扩,第一维请输入0')
sys.exit()
[dim0min, dim0max, dim1min, dim1max, dim2min, dim2max] = get_box(mask)
#判断能否外扩
dim0,dim1,dim2 = img.shape
dim1_l_index = np.clip(dim1min-ex_pix[1],0 ,dim1) #dim1外扩后左边的坐标,若触碰边界,则尽量外扩至边界
dim1_r_index = np.clip(dim1max + ex_pix[1], 0, dim1)
dim2_l_index = np.clip(dim2min - ex_pix[2], 0, dim2)
dim2_r_index = np.clip(dim2max + ex_pix[2], 0, dim2)
fina_img = img[:, dim1_l_index:dim1_r_index+1, dim2_l_index:dim2_r_index+1]
fina_mask = mask[:, dim1_l_index:dim1_r_index+1, dim2_l_index:dim2_r_index+1]
if ex_pix[0]:
dim0_l_index = np.clip(dim0min - ex_pix[0], 0, dim0)
dim0_r_index = np.clip(dim0max + ex_pix[0], 0, dim0)
fina_img = fina_img[dim0_l_index:dim0_r_index+1, :, :]
fina_mask = fina_mask[dim0_l_index:dim0_r_index+1, :, :]
else: #不外扩
print('dim0 不外扩')
dim0_l_index = dim0min
dim0_r_index = dim0max
fina_img = fina_img[dim0_l_index:dim0_r_index+1, :, :]
fina_mask = fina_mask[dim0_l_index:dim0_r_index+1, :, :]
return fina_img, fina_mask
def data_pretreatment():
global Image_3D
global ROI_3D
global ROI_small_3D
global Image_small_3D
global Current_name
global Input_File
if Image_3D.all() ==None:
return '没有数据'
else:
roi = ROI_3D
# waikuo = [4, 4, 4]
# fina_img, fina_mask = arry_crop_3D(Image_3D,roi,waikuo)
cut_thre = np.percentile(fina_img, 99.9) # 直方图99.9%右侧值不要
fina_img[fina_img >= cut_thre] = cut_thre
z, y, x = fina_img.shape
fina_img = resize3D(fina_img, [224,y,224], order=3)
fina_roi = resize3D(roi, [224, y, 224], order=3)
fina_img = (np.max(fina_img)-fina_img)/(np.max(fina_img)-np.min(fina_img))
Image_small_3D = fina_img
ROI_small_3D = fina_roi
return '预处理结束'
class App:
def __init__(self):
self.demo = None
self.main()
def main(self):
# get_name = gr.Interface(lambda name: name, inputs="textbox", outputs="textbox")
# prepend_hello = gr.Interface(lambda name: f"Hello {name}!", inputs="textbox", outputs="textbox")
# append_nice = gr.Interface(lambda greeting: f"{greeting} Nice to meet you!",
# inputs="textbox", outputs=gr.Textbox(label="Greeting"))
# iface_1 = gr.Interface(fn=get_Image_reslice, inputs=gr.inputs.File(label="Upload NIfTI file"), outputs=[,gr.Image(shape=(5, 5)),gr.Image(shape=(5, 5))])
with gr.Blocks() as demo:
with gr.Row():
with gr.Column(scale=1):
inp = gr.inputs.File(label="Upload MRI file")
inp2 = gr.inputs.File(label="Upload ROI file")
with gr.Column(scale=1):
out8 = gr.Textbox(placeholder="尚未加载影像数据")
out9 = gr.Textbox(placeholder="尚未加载感兴趣区域")
with gr.Row():
btn1 = gr.Button("Upload MRI")
btn5 = gr.Button("Upload ROI")
clear = gr.Button(" Clear All")
with gr.Tab("Image"):
with gr.Row():
with gr.Column(scale=1):
out1 = gr.Image(shape=(10, 10))
slider1 = gr.Slider(1, 128, label='z轴层数', step=1, interactive=True)
with gr.Column(scale=1):
out2 = gr.Image(shape=(10, 10))
slider2 = gr.Slider(1, 256, label='y轴层数', step=1, interactive=True)
with gr.Column(scale=1):
out3 = gr.Image(shape=(10, 10))
slider3 = gr.Slider(1, 128, label='x轴层数', step=1, interactive=True)
with gr.Tab("Medical Information"):
with gr.Row():
with gr.Column(scale=1):
btn2 = gr.Button(value="临床信息")
out4 = gr.Textbox(label="患病史")
out6 = gr.Textbox(label="现病史")
with gr.Column(scale=1):
btn4 = gr.Button("预处理")
out7 = gr.Textbox(placeholder="尚未进行预处理 请先预处理再按“分期结果”按钮", )
btn3 = gr.Button("分期结果")
out5 = gr.Label(num_top_classes=2, label='分期结果')
btn3.click(inference, inputs=None, outputs=out5)
btn4.click(data_pretreatment, inputs=None, outputs=out7)
btn2.click(get_medical_message, inputs=None, outputs=[out4, out6])
# demo = gr.Series(get_name, prepend_hello, append_nice)
btn1.click(get_Image_reslice, inp, [out1, out2, out3, slider1, slider2, slider3,out8])
btn5.click(get_ROI, inputs=inp2, outputs=out9)
slider3.change(change_image_slice_x, inputs=slider3, outputs=out3)
slider2.change(change_image_slice_y, inputs=slider2, outputs=out2)
slider1.change(change_image_slice_z, inputs=slider1, outputs=out1)
clear.click(clear_all, None, [out1, out2, out3, out4, out6, out5, out7,out8,out9], queue=True)
gr.Markdown('''# Examples''')
gr.Examples(
examples=[["./2239561_B_R_MRI.nii.gz"],
["./2239561_B_R_MRI.nii.gz"]],
inputs=inp,
outputs=[out1, out2, out3, slider1, slider2, slider3,out8],
fn=get_Image_reslice,
cache_examples=True,
)
gr.Examples(
examples=[["./2239561_B_R_ROI.nii.gz"],
["./2239561_B_R_ROI.nii.gz"]],
inputs=inp2,
outputs=out9,
fn=get_ROI,
cache_examples=True,
)
demo.queue(concurrency_count=6)
demo.launch(share=False)
app = App()