|
import streamlit as st |
|
import nibabel as nib |
|
import os.path |
|
import os |
|
from nilearn import plotting |
|
|
|
import torch |
|
from monai.transforms import ( |
|
EnsureChannelFirst, |
|
Compose, |
|
Resize, |
|
ScaleIntensity, |
|
LoadImage, |
|
) |
|
import torch.nn.functional as F |
|
import numpy as np |
|
from statistics import mean |
|
|
|
from constants import CLASSES |
|
from model.download_model import load_model |
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
st.set_page_config(page_title = "Alzheimer Classifier", page_icon = ":brain:", layout = "wide") |
|
|
|
|
|
model = load_model() |
|
|
|
|
|
transforms = Compose([ |
|
ScaleIntensity(), |
|
EnsureChannelFirst(), |
|
Resize((96, 96, 96)), |
|
]) |
|
load_img = LoadImage(image_only=True) |
|
|
|
|
|
class_names = CLASSES |
|
|
|
|
|
if 'clicked_pp' not in st.session_state: |
|
st.session_state.clicked_pp = False |
|
|
|
if 'clicked_pred' not in st.session_state: |
|
st.session_state.clicked_pred = False |
|
|
|
def click_pp_true(): |
|
st.session_state.clicked_pp = True |
|
|
|
def click_pred_true(): |
|
st.session_state.clicked_pred = True |
|
|
|
def click_false(): |
|
st.session_state.clicked_pp = False |
|
st.session_state.clicked_pred = False |
|
|
|
|
|
|
|
|
|
|
|
with st.sidebar: |
|
st.title("Alzheimer Classifier Demo") |
|
img_path = st.selectbox( |
|
"Select Image", |
|
tuple(class_names), |
|
on_change= click_false, |
|
) |
|
col1, col2 = st.columns((1,1)) |
|
with col1: |
|
run_preprocess = st.button("Preprocess Image", on_click=click_pp_true) |
|
if st.session_state.clicked_pp: |
|
with col2: |
|
run_pred = st.button("Run Prediction", on_click= click_pred_true) |
|
|
|
with st.container(): |
|
if img_path != "": |
|
if st.session_state.clicked_pp: |
|
if st.session_state.clicked_pred == False: |
|
with st.container(): |
|
pred_image = nib.load(hf_hub_download(repo_id= "rootstrap-org/Alzheimer-Classifier-Demo", repo_type="dataset", subfolder="preprocessed", filename = img_path + ".nii.gz")) |
|
|
|
bounds_pred = plotting.find_cuts._get_auto_mask_bounds(pred_image) |
|
|
|
st.sidebar.write("#") |
|
y_value_pred = st.sidebar.slider('Move the slider to adjust the coronal cut ', bounds_pred[1][0], bounds_pred[1][1], mean([bounds_pred[1][0], bounds_pred[1][1]])) |
|
x_value_pred = st.sidebar.slider('Move the slider to adjust the sagittal cut ', bounds_pred[0][0], bounds_pred[0][1], mean([bounds_pred[0][0], bounds_pred[0][1]])) |
|
z_value_pred = st.sidebar.slider('Move the slider to adjust the axial cut ', bounds_pred[2][0], bounds_pred[2][1], mean([bounds_pred[2][0], bounds_pred[2][1]])) |
|
|
|
plotting.plot_img(pred_image, cmap="grey", cut_coords=(x_value_pred,y_value_pred,z_value_pred), black_bg=True) |
|
st.pyplot() |
|
|
|
else: |
|
with st.container(): |
|
pred_image = nib.load(hf_hub_download(repo_id= "rootstrap-org/Alzheimer-Classifier-Demo", repo_type="dataset", subfolder="preprocessed", filename = img_path + ".nii.gz")) |
|
|
|
bounds_pred = plotting.find_cuts._get_auto_mask_bounds(pred_image) |
|
|
|
st.sidebar.write("#") |
|
y_value_pred = st.sidebar.slider('Move the slider to adjust the coronal cut ', bounds_pred[1][0], bounds_pred[1][1], mean([bounds_pred[1][0], bounds_pred[1][1]])) |
|
x_value_pred = st.sidebar.slider('Move the slider to adjust the sagittal cut ', bounds_pred[0][0], bounds_pred[0][1], mean([bounds_pred[0][0], bounds_pred[0][1]])) |
|
z_value_pred = st.sidebar.slider('Move the slider to adjust the axial cut ', bounds_pred[2][0], bounds_pred[2][1], mean([bounds_pred[2][0], bounds_pred[2][1]])) |
|
|
|
img_array = load_img(hf_hub_download(repo_id= "rootstrap-org/Alzheimer-Classifier-Demo", repo_type="dataset", subfolder="preprocessed", filename = img_path + ".nii.gz")) |
|
new_data = transforms(img_array) |
|
new_data_tensor = torch.from_numpy(np.array([new_data])) |
|
|
|
with torch.no_grad(): |
|
output = model(new_data_tensor) |
|
|
|
probabilities = F.softmax(output, dim=1) |
|
probabilities_np = probabilities.numpy() |
|
probabilities_item = probabilities_np[0] |
|
probabilities_percentage = probabilities_item * 100 |
|
predicted_class_index = np.argmax(probabilities_np[0]) |
|
predicted_class_name = class_names[predicted_class_index] |
|
predicted_probability = probabilities_percentage[predicted_class_index] |
|
|
|
st.sidebar.write("#") |
|
if predicted_class_index == 0: |
|
color_name = "red" |
|
elif predicted_class_index == 1: |
|
color_name = "blue" |
|
elif predicted_class_index == 2: |
|
color_name = "green" |
|
|
|
if predicted_probability > 80: |
|
color_prob = "green" |
|
elif predicted_probability > 60: |
|
color_prob = "yellow" |
|
else: |
|
color_prob = "red" |
|
|
|
class_col, pred_col = st.columns((1,1)) |
|
|
|
with class_col: |
|
st.write(f"### Predicted Class: :{color_name}[{predicted_class_name}]") |
|
|
|
with pred_col: |
|
st.write(f"### Probability: :{color_prob}[{predicted_probability:.2f}%]") |
|
|
|
plotting.plot_img(pred_image, cmap="grey", cut_coords=(x_value_pred,y_value_pred,z_value_pred), black_bg=True) |
|
st.pyplot() |
|
|
|
else: |
|
raw_image = nib.load(hf_hub_download(repo_id= "rootstrap-org/Alzheimer-Classifier-Demo", repo_type="dataset", subfolder="raw", filename = img_path + ".nii")) |
|
|
|
bounds_raw = plotting.find_cuts._get_auto_mask_bounds(raw_image) |
|
|
|
st.sidebar.write("#") |
|
y_value_raw = st.sidebar.slider('Move the slider to adjust the coronal cut', bounds_raw[1][0], bounds_raw[1][1], mean([bounds_raw[1][0], bounds_raw[1][1]])) |
|
x_value_raw = st.sidebar.slider('Move the slider to adjust the sagittal cut', bounds_raw[0][0], bounds_raw[0][1], mean([bounds_raw[0][0], bounds_raw[0][1]])) |
|
z_value_raw = st.sidebar.slider('Move the slider to adjust the axial cut', bounds_raw[2][0], bounds_raw[2][1], mean([bounds_raw[2][0], bounds_raw[2][1]])) |
|
|
|
plotting.plot_img(raw_image, cmap = "grey", cut_coords=(x_value_raw,y_value_raw,z_value_raw), black_bg=True) |
|
st.pyplot() |
|
|
