Add files via upload

README.md

@@ -1 +1,9 @@
-# Alzheimer-Classifier-Demo
+---
+license: mit
+emoji: 🧠
+title: Alzheimer Classifier
+sdk: streamlit
+colorFrom: gray
+colorTo: purple
+---
+# MRI-classifier-streamlit

app.py

@@ -0,0 +1,169 @@
+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 download_pictures import download_images
+
+#SET PAGE TITLE
+st.set_page_config(page_title = "Alzheimer Classifier", page_icon = ":brain:", layout = "wide")
+
+#LOAD MODEL
+model = load_model()
+
+#LOAD IMAGES
+download_images()
+
+#SET NIFTI FILE LOADING AND PROCESSING CONFIGURATIONS
+transforms = Compose([
+    ScaleIntensity(),
+    EnsureChannelFirst(),
+    Resize((96, 96, 96)),
+])
+load_img = LoadImage(image_only=True)
+
+#SET CLASSES
+class_names = CLASSES
+
+#SET IMAGE PATH LIST FOR STREAMLIT'S SELECT BOX
+filelist=[""]
+for root, dirs, files in os.walk("images/raw"):
+      for file in files:
+             filename=file.split(".")[0]
+             filelist.append(filename)
+filelist = tuple(filelist)
+
+#SILENCE STREAMIT WARNING
+st.set_option('deprecation.showPyplotGlobalUse', False)
+
+#SET STREAMLIT SESSION STATES
+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
+
+###########################################################
+###################### STREAMLIT APP ######################
+###########################################################   
+
+with st.sidebar:
+    st.title("Alzheimer Classifier Demo")
+    img_path = st.selectbox(
+            "Select Image",
+            filelist,
+            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(os.path.join("images/preprocessed", 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(os.path.join("images/preprocessed", 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(os.path.join("images/preprocessed", 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(os.path.join("images/raw", 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()

constants.py

@@ -0,0 +1,15 @@
+import os
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Root dir
+ROOT_DIR = os.getcwd()
+
+# Model checkpoints and repo names
+MODEL_FILENAME = os.getenv("MODEL_FILENAME")
+HF_MODEL_REPO_NAME = os.getenv("HF_MODEL_REPO_NAME")
+
+# Other constants
+CLASSES = ["Alzheimer's Desease", "Mild Cognitive Impairment", "Control"]

download_pictures.py

@@ -0,0 +1,6 @@
+from huggingface_hub import snapshot_download
+from constants import HF_MODEL_REPO_NAME, ROOT_DIR
+import os
+
+def download_images():
+    snapshot_download(repo_id=HF_MODEL_REPO_NAME, repo_type="dataset", local_dir=os.path.join(ROOT_DIR, "images"))

model/download_model.py

@@ -0,0 +1,30 @@
+import os.path
+import os
+
+import monai.networks.nets as nets
+import torch
+
+from huggingface_hub import hf_hub_download
+
+from constants import ROOT_DIR, MODEL_FILENAME, HF_MODEL_REPO_NAME
+
+def load_model():
+    """
+    Load pretrained model
+    """
+
+    model_path = os.path.join(ROOT_DIR, "model", MODEL_FILENAME)
+
+    # If model doesnt exist download from huggingface
+    if not os.path.exists(model_path):
+        hf_hub_download(HF_MODEL_REPO_NAME, MODEL_FILENAME, local_dir=os.path.join(ROOT_DIR, "model"))
+
+    model = nets.DenseNet121(spatial_dims=3, in_channels=1, out_channels=3)
+    if torch.cuda.is_available():
+        checkpoint = torch.load(model_path)
+    else:
+        checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
+    model.load_state_dict(checkpoint)
+    model.eval()
+
+    return model

requirements.txt

@@ -0,0 +1,9 @@
+huggingface-hub==0.20.1
+matplotlib==3.8.2
+monai==1.3.0
+nibabel==5.2.0
+nilearn==0.10.2
+numpy==1.26.2
+python-dotenv==1.0.0
+streamlit==1.29.0
+torch==2.1.1