Update app.py

app.py

@@ -2,7 +2,7 @@
 import os
 import io
 import torch
-# import shutil
+import pydicom
 import numpy as np
 import streamlit as st
 
@@ -20,7 +20,6 @@ from monai.transforms import (
     EnsureChannelFirst,
     AsDiscrete,
     Compose,
-    LoadImage,
     RandFlip,
     RandRotate,
     RandZoom,
@@ -105,7 +104,6 @@ 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)
@@ -115,7 +113,6 @@ eval_transforms = Compose(
 # CAM Transforms
 cam_transforms = Compose(
     [
-        # LoadImage(image_only=True), 
         AsChannelFirst(), 
         Resize(spatial_size=SPATIAL_SIZE)
     ]
@@ -124,7 +121,6 @@ cam_transforms = Compose(
 # Original Transforms
 original_transforms = Compose(
     [
-        # LoadImage(image_only=True), 
         AsChannelFirst()
     ]
 )
@@ -135,18 +131,26 @@ 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"))
+# Convert the file size from bytes to megabytes
+def bytes_to_megabytes(file_size_bytes):
+    # Convert bytes to MB (1 MB = 1024 * 1024 bytes)
+    file_size_megabytes = round(file_size_bytes / (1024 * 1024), 2)
+    return str(file_size_megabytes) + " MB"  # Rounding to 2 decimal places for readability
 
-# def create_dir(dirname: str):
-#     if not os.path.exists(dirname):
-#         os.makedirs(dirname, exist_ok=True)
+def meta_tensor_to_numpy(meta_tensor):
+    """
+    Convert a PyTorch MetaTensor to a NumPy array
+    """
+    # Ensure the MetaTensor is on the CPU
+    meta_tensor = meta_tensor.cpu()
 
-# create_dir("CT_tempDir")
-# create_dir("MRI_tempDir")
+    # Convert the MetaTensor to a PyTorch tensor
+    torch_tensor = meta_tensor.to(dtype=torch.float32)
 
-# # Get the current working directory
-# current_directory = os.getcwd()
+    # Convert the PyTorch tensor to a NumPy array
+    numpy_array = torch_tensor.detach().numpy()
+
+    return numpy_array
 
 set_determinism(seed=SEED)
 torch.manual_seed(SEED)
@@ -187,11 +191,6 @@ CT_WINDOW_WIDTH = st.sidebar.number_input("CT Window Width", min_value=WINDOW_WI
 
 uploaded_mri_file = st.file_uploader("Upload a candidate MRI DICOM", type=["dcm"])
 if uploaded_mri_file is not None:
-    # To check file details
-    file_details = {"FileName": uploaded_mri_file.name, "FileType": uploaded_mri_file.type, "FileSize": uploaded_mri_file.size}
-    st.write(file_details)
-
-    import pydicom
     # Read DICOM file into NumPy array
     dicom_data = pydicom.dcmread(uploaded_mri_file)
     dicom_array = dicom_data.pixel_array
@@ -202,14 +201,13 @@ if uploaded_mri_file is not None:
     # Then add a channel dimension
     dicom_array = dicom_array[:, :, np.newaxis]
 
-    # Check the shape and dtype of dicom_array
-    st.write(f"Shape of dicom_array: {dicom_array.shape}")
-    st.write(f"Data type of dicom_array: {dicom_array.dtype}")
+    # To check file details
+    file_details = {"File_Name": uploaded_mri_file.name, "File_Type": uploaded_mri_file.type, "File_Size": bytes_to_megabytes(uploaded_mri_file.size), "File_Dimension": str((dicom_array.shape[0],dicom_array.shape[1]))}
+    st.write(file_details)
 
     transformed_array = eval_transforms(dicom_array)
 
     # Convert to PyTorch tensor and move to device
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     image_tensor = transformed_array.clone().detach().unsqueeze(0).to(device)
 
     # Predict
@@ -226,11 +224,11 @@ if uploaded_mri_file is not None:
 
     # Load the original DICOM image for download
     download_image_tensor = original_transforms(dicom_array).unsqueeze(0).to(device)
-    download_image = download_image_tensor.squeeze()
+    download_image_tensor = 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)
+    download_image_numpy = meta_tensor_to_numpy(download_image_tensor)
+    windowed_download_image = DICOM_Utils.apply_windowing(download_image_numpy, MRI_WINDOW_CENTER, MRI_WINDOW_WIDTH)
 
     # Streamlit button to trigger image download
     image_data = image_to_bytes(Image.fromarray(windowed_download_image))
@@ -243,11 +241,11 @@ if uploaded_mri_file is not None:
 
     # Load the original DICOM image for display
     display_image_tensor = cam_transforms(dicom_array).unsqueeze(0).to(device)
-    display_image = display_image_tensor.squeeze()
+    display_image_tensor = 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)
+    display_image_numpy = meta_tensor_to_numpy(display_image_tensor)
+    windowed_image = DICOM_Utils.apply_windowing(display_image_numpy, MRI_WINDOW_CENTER, MRI_WINDOW_WIDTH)
     st.image(Image.fromarray(windowed_image), caption="Original MRI Visualization", use_column_width=True)
 
     # Expand to three channels
@@ -270,72 +268,78 @@ if uploaded_mri_file is not None:
     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
-#     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())
-
-#         full_ct_temp_path = current_directory +"\\"+ ct_temp_path
-
-#         # Apply evaluation transforms to the DICOM image for model prediction
-#         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)
+uploaded_ct_file = st.file_uploader("Upload a candidate CT DICOM", type=["dcm"])
+if uploaded_ct_file is not None:
+    # Read DICOM file into NumPy array
+    dicom_data = pydicom.dcmread(uploaded_ct_file)
+    dicom_array = dicom_data.pixel_array
+
+    # Convert the data type to float32
+    dicom_array = dicom_array.astype(np.float32)
+
+    # Then add a channel dimension
+    dicom_array = dicom_array[:, :, np.newaxis]
+
+    # To check file details
+    file_details = {"File_Name": uploaded_ct_file.name, "File_Type": uploaded_ct_file.type, "File_Size": bytes_to_megabytes(uploaded_ct_file.size), "File_Dimension": str((dicom_array.shape[0],dicom_array.shape[1]))}
+    st.write(file_details)
+
+    transformed_array = eval_transforms(dicom_array)
+
+    # 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(dicom_array).unsqueeze(0).to(device)
+    download_image_tensor = download_image_tensor.squeeze()
+
+    # Transform the download image and apply windowing
+    download_image_numpy = meta_tensor_to_numpy(download_image_tensor)
+    windowed_download_image = DICOM_Utils.apply_windowing(download_image_numpy, 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(dicom_array).unsqueeze(0).to(device)
+    display_image_tensor = display_image_tensor.squeeze()
+
+    # Transform the image and apply windowing
+    display_image_numpy = meta_tensor_to_numpy(display_image_tensor)
+    windowed_image = DICOM_Utils.apply_windowing(display_image_numpy, 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=USE_CUDA)
+    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)