Upload 2 files

app.py

@@ -0,0 +1,306 @@
+# Import required libraries
+import os
+import io
+import torch
+import tempfile
+import numpy as np
+import streamlit as st
+
+# Import utility and custom functions
+from PIL import Image
+from Util.DICOM import DICOM_Utils
+from Util.Custom_Model import Build_Custom_Model, reshape_transform
+
+# Import additional MONAI and PyTorch Grad-CAM utilities
+from monai.config import print_config
+from monai.utils import set_determinism
+from monai.networks.nets import SEResNet50
+from monai.transforms import (
+    Activations,
+    EnsureChannelFirst,
+    AsDiscrete,
+    Compose,
+    LoadImage,
+    RandFlip,
+    RandRotate,
+    RandZoom,
+    ScaleIntensity,
+    AsChannelFirst,
+    AddChannel,
+    RandSpatialCrop,
+    ScaleIntensityRangePercentiles,
+    Resize,
+)
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+
+
+# (Int) Random seed
+SEED = 0
+
+# (Int) Model parameters
+NUM_CLASSES = 1
+
+# (String) CT Model directory
+CT_MODEL_DIRECTORY = "C:\\Src\\GitHub\\AI_UI\\models\\CLOTS\\CT"
+
+# (String) MRI Model directory
+MRI_MODEL_DIRECTORY = "C:\\Src\\GitHub\\AI_UI\\models\\CLOTS\\MRI"
+
+# (Boolean) Use custom model
+CUSTOM_MODEL_FLAG = True
+
+# (List[int]) Image size
+SPATIAL_SIZE = [224, 224]
+
+# (String) CT Model file name
+CT_MODEL_FILE_NAME = "best_metric_model.pth"
+
+# (String) MRI Model file name
+MRI_MODEL_FILE_NAME = "best_metric_model.pth"
+
+# (Boolean) List model modules
+LIST_MODEL_MODULES = False
+
+# (String) Model name
+CT_MODEL_NAME = "swin_base_patch4_window7_224"
+
+# (String) Model name
+MRI_MODEL_NAME = "swin_base_patch4_window7_224"
+
+# (Float) Model inference threshold
+CT_INFERENCE_THRESHOLD = 0.5
+
+# (Float) Model inference threshold
+MRI_INFERENCE_THRESHOLD = 0.5
+
+# (Int) Display CAM Class ID
+CAM_CLASS_ID = 0
+
+# (Int) Window Center for image display
+DEFAULT_CT_WINDOW_CENTER = 40
+
+# (Int) Window Width for image display
+DEFAULT_CT_WINDOW_WIDTH = 100
+
+# (Int) Window Center for image display
+DEFAULT_MRI_WINDOW_CENTER = 400
+
+# (Int) Window Width for image display
+DEFAULT_MRI_WINDOW_WIDTH = 1000
+
+# (Int) Minimum value for Window Center
+WINDOW_CENTER_MIN = -600
+
+# (Int) Maximum value for Window Center
+WINDOW_CENTER_MAX = 1000
+
+# (Int) Minimum value for Window Width
+WINDOW_WIDTH_MIN = 1
+
+# (Int) Maximum value for Window Width
+WINDOW_WIDTH_MAX = 3000
+
+# Evaluation Transforms
+eval_transforms = Compose(
+    [
+        LoadImage(image_only=True), 
+        AsChannelFirst(), 
+        ScaleIntensityRangePercentiles(lower=20, upper=80, b_min=0.0, b_max=1.0, clip=False, relative=True), 
+        Resize(spatial_size=SPATIAL_SIZE)
+    ]
+)
+
+# CAM Original Transforms
+cam_original_transforms = Compose(
+    [
+        LoadImage(image_only=True), 
+        AsChannelFirst(), 
+        Resize(spatial_size=SPATIAL_SIZE)
+    ]
+)
+
+# CAM Original Transforms
+original_transforms = Compose(
+    [
+        LoadImage(image_only=True), 
+        AsChannelFirst()
+    ]
+)
+
+# Function to convert PIL Image to byte stream in PNG format for downloading
+def image_to_bytes(image):
+    byte_stream = io.BytesIO()
+    image.save(byte_stream, format='PNG')
+    return byte_stream.getvalue()
+
+set_determinism(seed=SEED)
+torch.manual_seed(SEED)
+
+# Parameters
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ct_root_dir = tempfile.mkdtemp() if CT_MODEL_DIRECTORY is None else CT_MODEL_DIRECTORY
+mri_root_dir = tempfile.mkdtemp() if MRI_MODEL_DIRECTORY is None else MRI_MODEL_DIRECTORY
+
+def load_model(root_dir, model_name, model_file_name):
+    if CUSTOM_MODEL_FLAG:
+        model = Build_Custom_Model(model_name, NUM_CLASSES, pretrained=False).to(device)
+    else:
+        model = SEResNet50(spatial_dims=2, in_channels=1, num_classes=NUM_CLASSES).to(device)
+    model.load_state_dict(torch.load(os.path.join(root_dir, model_file_name)))
+    model.eval()
+    return model
+
+ct_model = load_model(ct_root_dir, CT_MODEL_NAME, CT_MODEL_FILE_NAME)
+mri_model = load_model(mri_root_dir, MRI_MODEL_NAME, MRI_MODEL_FILE_NAME)
+if LIST_MODEL_MODULES:
+    for ct_name, _ in ct_model.named_modules():
+        print(ct_name)
+
+    for mri_name, _ in mri_model.named_modules():
+        print(mri_name)
+
+# Initialize Streamlit
+st.title("Analyze")
+
+# Use Streamlit's number_input to adjust WINDOW_CENTER and WINDOW_WIDTH
+st.sidebar.header("Windowing Parameters for DICOM")
+CT_WINDOW_CENTER = st.sidebar.number_input("CT Window Center", min_value=WINDOW_CENTER_MIN, max_value=WINDOW_CENTER_MAX, value=DEFAULT_CT_WINDOW_CENTER, step=1)
+CT_WINDOW_WIDTH = st.sidebar.number_input("CT Window Width", min_value=WINDOW_WIDTH_MIN, max_value=WINDOW_WIDTH_MAX, value=DEFAULT_CT_WINDOW_WIDTH, step=1)
+MRI_WINDOW_CENTER = st.sidebar.number_input("MRI Window Center", min_value=WINDOW_CENTER_MIN, max_value=WINDOW_CENTER_MAX, value=DEFAULT_MRI_WINDOW_CENTER, step=1)
+MRI_WINDOW_WIDTH = st.sidebar.number_input("MRI Window Width", min_value=WINDOW_WIDTH_MIN, max_value=WINDOW_WIDTH_MAX, value=DEFAULT_MRI_WINDOW_WIDTH, step=1)
+
+uploaded_ct_file = st.file_uploader("Upload a candidate CT DICOM", type=["dcm"])
+if uploaded_ct_file is not None:
+    # Save the uploaded file to a temporary location
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".dcm") as temp_file:
+        temp_file.write(uploaded_ct_file.getvalue())
+
+        # Apply evaluation transforms to the DICOM image for model prediction
+        image_tensor = eval_transforms(temp_file.name).unsqueeze(0).to(device)
+
+        # Predict
+        with torch.no_grad():
+            outputs = ct_model(image_tensor).sigmoid().to("cpu").numpy()
+            prob = outputs[0][0]
+            CLOTS_CLASSIFICATION = False
+            if(prob >= CT_INFERENCE_THRESHOLD):
+                CLOTS_CLASSIFICATION=True
+
+        st.header("CT Classification")
+        st.subheader(f"Ischaemic Stroke : {CLOTS_CLASSIFICATION}")
+        st.subheader(f"Confidence : {prob * 100:.1f}%")
+
+        # Load the original DICOM image for download
+        download_image_tensor = original_transforms(temp_file.name).unsqueeze(0).to(device)
+        download_image = download_image_tensor.squeeze()
+
+        # Transform the download image and apply windowing
+        transformed_download_image = DICOM_Utils.transform_image_for_display(download_image)
+        windowed_download_image = DICOM_Utils.apply_windowing(transformed_download_image, CT_WINDOW_CENTER, CT_WINDOW_WIDTH)
+
+        # Streamlit button to trigger image download
+        image_data = image_to_bytes(Image.fromarray(windowed_download_image))
+        st.download_button(
+            label="Download CT Image",
+            data=image_data,
+            file_name="downloaded_ct_image.png",
+            mime="image/png"
+        )
+
+        # Load the original DICOM image for display
+        display_image_tensor = cam_original_transforms(temp_file.name).unsqueeze(0).to(device)
+        display_image = display_image_tensor.squeeze()
+
+        # Transform the image and apply windowing
+        transformed_image = DICOM_Utils.transform_image_for_display(display_image)
+        windowed_image = DICOM_Utils.apply_windowing(transformed_image, CT_WINDOW_CENTER, CT_WINDOW_WIDTH)
+        st.image(Image.fromarray(windowed_image), caption="Original CT Visualization", use_column_width=True)
+
+        # Expand to three channels
+        windowed_image = np.expand_dims(windowed_image, axis=2)
+        windowed_image = np.tile(windowed_image, [1, 1, 3])
+
+        # Ensure both are of float32 type
+        windowed_image = windowed_image.astype(np.float32)
+
+        # Normalize to [0, 1] range
+        windowed_image = np.float32(windowed_image) / 255
+
+        # Build the CAM (Class Activation Map)
+        target_layers = [ct_model.model.norm]
+        cam = GradCAM(model=ct_model, target_layers=target_layers, reshape_transform=reshape_transform, use_cuda=True)
+        grayscale_cam = cam(input_tensor=image_tensor, targets=[ClassifierOutputTarget(CAM_CLASS_ID)])
+        grayscale_cam = grayscale_cam[0, :]
+
+        # Now you can safely call the show_cam_on_image function
+        visualization = show_cam_on_image(windowed_image, grayscale_cam, use_rgb=True)
+        st.image(Image.fromarray(visualization), caption="CAM CT Visualization", use_column_width=True)
+
+uploaded_mri_file = st.file_uploader("Upload a candidate MRI DICOM", type=["dcm"])
+if uploaded_mri_file is not None:
+    # Save the uploaded file to a temporary location
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".dcm") as temp_file:
+        temp_file.write(uploaded_mri_file.getvalue())
+
+        # Apply evaluation transforms to the DICOM image for model prediction
+        image_tensor = eval_transforms(temp_file.name).unsqueeze(0).to(device)
+
+        # Predict
+        with torch.no_grad():
+            outputs = mri_model(image_tensor).sigmoid().to("cpu").numpy()
+            prob = outputs[0][0]
+            CLOTS_CLASSIFICATION = False
+            if(prob >= MRI_INFERENCE_THRESHOLD):
+                CLOTS_CLASSIFICATION=True
+
+        st.header("MRI Classification")
+        st.subheader(f"Ischaemic Stroke : {CLOTS_CLASSIFICATION}")
+        st.subheader(f"Confidence : {prob * 100:.1f}%")
+
+        # Load the original DICOM image for download
+        download_image_tensor = original_transforms(temp_file.name).unsqueeze(0).to(device)
+        download_image = download_image_tensor.squeeze()
+
+        # Transform the download image and apply windowing
+        transformed_download_image = DICOM_Utils.transform_image_for_display(download_image)
+        windowed_download_image = DICOM_Utils.apply_windowing(transformed_download_image, MRI_WINDOW_CENTER, MRI_WINDOW_WIDTH)
+
+        # Streamlit button to trigger image download
+        image_data = image_to_bytes(Image.fromarray(windowed_download_image))
+        st.download_button(
+            label="Download MRI Image",
+            data=image_data,
+            file_name="downloaded_mri_image.png",
+            mime="image/png"
+        )
+
+        # Load the original DICOM image for display
+        display_image_tensor = cam_original_transforms(temp_file.name).unsqueeze(0).to(device)
+        display_image = display_image_tensor.squeeze()
+
+        # Transform the image and apply windowing
+        transformed_image = DICOM_Utils.transform_image_for_display(display_image)
+        windowed_image = DICOM_Utils.apply_windowing(transformed_image, MRI_WINDOW_CENTER, MRI_WINDOW_WIDTH)
+        st.image(Image.fromarray(windowed_image), caption="Original MRI Visualization", use_column_width=True)
+
+        # Expand to three channels
+        windowed_image = np.expand_dims(windowed_image, axis=2)
+        windowed_image = np.tile(windowed_image, [1, 1, 3])
+
+        # Ensure both are of float32 type
+        windowed_image = windowed_image.astype(np.float32)
+
+        # Normalize to [0, 1] range
+        windowed_image = np.float32(windowed_image) / 255
+
+        # Build the CAM (Class Activation Map)
+        target_layers = [mri_model.model.norm]
+        cam = GradCAM(model=mri_model, target_layers=target_layers, reshape_transform=reshape_transform, use_cuda=True)
+        grayscale_cam = cam(input_tensor=image_tensor, targets=[ClassifierOutputTarget(CAM_CLASS_ID)])
+        grayscale_cam = grayscale_cam[0, :]
+
+        # Now you can safely call the show_cam_on_image function
+        visualization = show_cam_on_image(windowed_image, grayscale_cam, use_rgb=True)
+        st.image(Image.fromarray(visualization), caption="CAM MRI Visualization", use_column_width=True)

requirement.txt

@@ -0,0 +1,3 @@
+streamlit
+gradio
+monai