Update app.py

app.py

@@ -2,7 +2,7 @@
 import os
 import io
 import torch
-import tempfile
+# import shutil
 import numpy as np
 import streamlit as st
 
@@ -135,25 +135,36 @@ def image_to_bytes(image):
     image.save(byte_stream, format='PNG')
     return byte_stream.getvalue()
 
+# if os.path.exists("tempDir"):
+#     shutil.rmtree(os.path.join("tempDir"))
+
+def create_dir(dirname: str):
+    if not os.path.exists(dirname):
+        os.makedirs(dirname, exist_ok=True)
+
+create_dir("CT_tempDir")
+create_dir("MRI_tempDir")
+
+# Get the current working directory
+current_directory = os.getcwd()
+
 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),  map_location=device))
+    model.load_state_dict(torch.load(os.path.join(root_dir, model_file_name), map_location=device))
     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)
+ct_model = load_model(CT_MODEL_DIRECTORY, CT_MODEL_NAME, CT_MODEL_FILE_NAME)
+mri_model = load_model(MRI_MODEL_DIRECTORY, MRI_MODEL_NAME, MRI_MODEL_FILE_NAME)
 if LIST_MODEL_MODULES:
     for ct_name, _ in ct_model.named_modules():
         print(ct_name)
@@ -166,145 +177,147 @@ 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)
+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)
+
+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
+    mri_temp_path = os.path.join("MRI_tempDir", uploaded_mri_file.name)
+    with open(mri_temp_path, "wb") as f:
+        f.write(uploaded_mri_file.getbuffer())
+
+        full_mri_temp_path = current_directory +"\\"+ mri_temp_path
+
+        # Apply evaluation transforms to the DICOM image for model prediction
+        image_tensor = eval_transforms(full_mri_temp_path).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(full_mri_temp_path).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_transforms(full_mri_temp_path).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)
 
 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())
+    ct_temp_path = os.path.join("CT_tempDir", uploaded_ct_file.name)
+    with open(ct_temp_path, "wb") as f:
+        f.write(uploaded_ct_file.getbuffer())
 
-        print(tempfile.name)
+        full_ct_temp_path = current_directory +"\\"+ ct_temp_path
 
         # Apply evaluation transforms to the DICOM image for model prediction
-        image_tensor = eval_transforms(temp_file.name).unsqueeze(0).to(device)
-
-        print(image_tensor)
-
-#         # 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_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_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)
+        image_tensor = eval_transforms(full_ct_temp_path).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(full_ct_temp_path).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_transforms(full_ct_temp_path).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)