Update app.py

app.py

@@ -105,7 +105,7 @@ WINDOW_WIDTH_MAX = 3000
 # Evaluation Transforms
 eval_transforms = Compose(
     [
-        LoadImage(image_only=True), 
+        # 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 +115,7 @@ eval_transforms = Compose(
 # CAM Transforms
 cam_transforms = Compose(
     [
-        LoadImage(image_only=True), 
+        # LoadImage(image_only=True), 
         AsChannelFirst(), 
         Resize(spatial_size=SPATIAL_SIZE)
     ]
@@ -124,7 +124,7 @@ cam_transforms = Compose(
 # Original Transforms
 original_transforms = Compose(
     [
-        LoadImage(image_only=True), 
+        # LoadImage(image_only=True), 
         AsChannelFirst()
     ]
 )
@@ -138,15 +138,15 @@ def image_to_bytes(image):
 # 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)
+# 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")
+# create_dir("CT_tempDir")
+# create_dir("MRI_tempDir")
 
-# Get the current working directory
-current_directory = os.getcwd()
+# # Get the current working directory
+# current_directory = os.getcwd()
 
 set_determinism(seed=SEED)
 torch.manual_seed(SEED)
@@ -184,140 +184,155 @@ 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:
-    # 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
-    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)
+    # 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
+
+    # Convert the data type to float32
+    dicom_array = dicom_array.astype(np.float32)
+
+    # 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}")
+
+    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
+    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(dicom_array).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(dicom_array).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
+#     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)