Update app.py

app.py

@@ -21,7 +21,7 @@ from keras.layers import Dense, GlobalAveragePooling2D
 from keras.models import Model
 from pydicom.pixel_data_handlers.util import apply_voi_lut
 
-# Environment Configuration
+# Environment Configuration ###############################################################
 os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = "./da-kalbe-63ee33c9cdbb.json"
 bucket_name = "da-kalbe-ml-result-png"
 storage_client = storage.Client()
@@ -29,20 +29,20 @@ bucket_result = storage_client.bucket(bucket_name)
 bucket_name_load = "da-ml-models"
 bucket_load = storage_client.bucket(bucket_name_load)
 
+# Utility Functions #######################################################################
+
+# Upload Image #
 st.sidebar.title("Configuration")
 uploaded_file = st.sidebar.file_uploader("Upload Original Image", type=["png", "jpg", "jpeg", "dcm"])
 enhancement_type = st.sidebar.selectbox(
-    "Enhancement Type", 
+    "Enhancement Type",
     ["Invert", "High Pass Filter", "Unsharp Masking", "Histogram Equalization", "CLAHE"]
 )
 
 st.sidebar.title("Detection")
 uploaded_detection = st.sidebar.file_uploader("Upload image to detect", type=["png", "jpg", "jpeg", "dcm"])
-# enhancement_type = st.sidebar.selectbox(
-#     "Enhancement Type", 
-#     ["Invert", "High Pass Filter", "Unsharp Masking", "Histogram Equalization", "CLAHE"]
-# )
 
+# object detection ########################################################################
 H_detection = 224
 W_detection = 224
 
@@ -84,21 +84,8 @@ def draw_bbox(image, bbox):
     image = cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
     return image
 
+# Upload to GCS ###########################################################################
 
-st.title("AI INTEGRATION FOR CHEST X-RAY")
-
-st.header("IMAGE ENHANCEMENT")
-
-@st.cache_resource
-def load_gradcam_model():
-    model = keras.models.load_model('./model_renamed.h5', compile=False)
-    return model
-
-
-# Utility Functions
-
-# Dictionaries to track InstanceNumbers and StudyInstanceUIDs per filename
-# Initialize session state for instance numbers and study UIDs
 if 'instance_numbers' not in st.session_state:
     st.session_state.instance_numbers = {}
 if 'study_uids' not in st.session_state:
@@ -225,7 +212,12 @@ def upload_folder_images(original_image_path, enhanced_image_path, file_name):
     upload_to_gcs(original_dicom_bytes, folder_name + '/' + 'original_image.dcm', content_type='application/dicom')
     upload_to_gcs(enhanced_dicom_bytes, folder_name + '/' + enhancement_name + '.dcm', content_type='application/dicom')
 
+# Grad cam ################################################################################
 
+@st.cache_resource
+def load_gradcam_model():
+    model = keras.models.load_model('./model_renamed.h5', compile=False)
+    return model
 
 def get_mean_std_per_batch(image_path, H=320, W=320):
     sample_data = []
@@ -277,20 +269,8 @@ def grad_cam(input_model, img_array, cls, layer_name):
 
     return cam
 
-
 # Compute Grad-CAM
 def compute_gradcam(model_gradcam, img_path, layer_name='bn'):
-    # base_model  = keras.applications.DenseNet121(weights = './densenet.hdf5', include_top = False)
-    # x = base_model.output
-    # x = keras.layers.GlobalAveragePooling2D()(x)
-    # predictions = keras.layers.Dense(14, activation = "sigmoid")(x)
-    # model_gradcam = keras.Model(inputs=base_model.input, outputs=predictions)
-    # model_gradcam.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
-    #           loss='sparse_categorical_crossentropy')
-    # model.load_weights('./pretrained_model.h5')
-    # Load the original model
-
-    # Now use this modified model in your application
     model_gradcam = load_gradcam_model()
     preprocessed_input = load_image(img_path)
     predictions = model_gradcam.predict(preprocessed_input)
@@ -298,10 +278,10 @@ def compute_gradcam(model_gradcam, img_path, layer_name='bn'):
     original_image = load_image(img_path, preprocess=False)
 
     # Assuming you have 14 classes as previously mentioned
-    labels = ['Cardiomegaly', 'Emphysema', 'Effusion', 'Hernia', 'Infiltration', 'Mass', 
-              'Nodule', 'Atelectasis', 'Pneumothorax', 'Pleural_Thickening', 
-              'Pneumonia', 'Fibrosis', 'Edema', 'Consolidation']
-    
+    labels = ['Cardiomegaly', 'Emphysema', 'Effusion', 'Hernia', 'Infiltration', 'Mass',
+                'Nodule', 'Atelectasis', 'Pneumothorax', 'Pleural_Thickening',
+                'Pneumonia', 'Fibrosis', 'Edema', 'Consolidation']
+
     for i in range(len(labels)):
         st.write(f"Generating gradcam for class {labels[i]}")
         gradcam = grad_cam(model_gradcam, preprocessed_input, i, layer_name)
@@ -310,6 +290,8 @@ def compute_gradcam(model_gradcam, img_path, layer_name='bn'):
         gradcam = cv2.addWeighted(gradcam, 0.5, original_image.squeeze().astype(np.uint8), 0.5, 0)
         st.image(gradcam, caption=f"{labels[i]}: p={predictions[0][i]:.3f}", use_column_width=True)
 
+# Image enhancement #######################################################################
+
 def calculate_mse(original_image, enhanced_image):
     mse = np.mean((original_image - enhanced_image) ** 2)
     return mse
@@ -383,7 +365,7 @@ def enhance_image(image, enhancement_type):
     else:
         raise ValueError(f"Unknown enhancement type: {enhancement_type}")
 
-# Function to add a button to redirect to the URL
+# Other Utils #############################################################################
 def redirect_button(url):
     button = st.button('Go to OHIF Viewer')
     if button:
@@ -415,6 +397,9 @@ def redirect_button(url):
 ########################### Streamlit Interface ###########################################
 ###########################################################################################
 
+st.title("AI INTEGRATION FOR CHEST X-RAY")
+
+st.header("IMAGE ENHANCEMENT")
 
 # File uploader for DICOM files
 if uploaded_file is not None:
@@ -445,7 +430,7 @@ if uploaded_file is not None:
 
             file_bytes = np.asarray(bytearray(uploaded_detection.read()), dtype=np.uint8)
             image = cv2.imdecode(file_bytes, 1)
-        
+
             # st.image(image, caption='Uploaded Image.', use_column_width=True)
 
             col1, col2 = st.columns(2)
@@ -465,7 +450,7 @@ if uploaded_file is not None:
                     st.image(img_array, caption="Original Image", use_column_width=True)
             else:
                 st.error("Unsupported image dimensions")
-            
+
             original_image = img_array
 
             # Example: convert to grayscale if it's a color image
@@ -496,7 +481,6 @@ if uploaded_file is not None:
         enhanced_image_path = "enhanced_image.png"
         cv2.imwrite(enhanced_image_path, enhanced_image)
 
-        
         # Save enhanced image to a file
         enhanced_image_path = "enhanced_image.png"
         cv2.imwrite(enhanced_image_path, enhanced_image)
@@ -504,113 +488,15 @@ if uploaded_file is not None:
         # Save original image to a file
         original_image_path = "original_image.png"
         cv2.imwrite(original_image_path, original_image)
-        upload_folder_images(original_image_path, enhanced_image_path, file_name)
-        
+
+        if st.button("Send to OHIF"):
+            upload_folder_images(original_image_path, enhanced_image_path, file_name)
+
     # Add the redirect button
     col1, col2, col3 = st.columns(3)
     with col1:
         redirect_button("https://new-ohif-viewer-k7c3gdlxua-et.a.run.app/")
 
-    # with col2:
-        # model_detection = load_model_detection()
-        # file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
-        # image = cv2.imdecode(file_bytes, 1)
-    
-        # st.image(image, caption='Uploaded Image.', use_column_width=True)
-    
-        # st.button('Detect')
-        #     st.write("Processing...")
-        #     input_image = preprocess_image(image)
-        #     pred_bbox, pred_label, pred_label_confidence = predict(model_detection, input_image)
-    
-        #     # Updated label mapping based on the dataset
-        #     label_mapping = {
-        #         0: 'Atelectasis',
-        #         1: 'Cardiomegaly',
-        #         2: 'Effusion',
-        #         3: 'Infiltrate',
-        #         4: 'Mass',
-        #         5: 'Nodule',
-        #         6: 'Pneumonia',
-        #         7: 'Pneumothorax'
-        #     }
-    
-        #     if pred_label_confidence < 0.2:
-        #         st.write("May not detect a disease.")
-        #     else:
-        #         pred_label_name = label_mapping[pred_label]
-        #         st.write(f"Prediction Label: {pred_label_name}")
-        #         st.write(f"Prediction Bounding Box: {pred_bbox}")
-        #         st.write(f"Prediction Confidence: {pred_label_confidence:.2f}")
-    
-        #         output_image = draw_bbox(image.copy(), pred_bbox)
-        #         st.image(output_image, caption='Detected Image.', use_column_width=True)
-
-
-        # file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
-        # image = cv2.imdecode(file_bytes, 1)
-        # st.button('Auto Detect')
-            # st.write("Processing...")
-            # input_image = preprocess_image(image)
-            # pred_bbox, pred_label, pred_label_confidence = predict(model_detection, input_image)
-    
-            # # Updated label mapping based on the dataset
-            # label_mapping = {
-            #     0: 'Atelectasis',
-            #     1: 'Cardiomegaly',
-            #     2: 'Effusion',
-            #     3: 'Infiltrate',
-            #     4: 'Mass',
-            #     5: 'Nodule',
-            #     6: 'Pneumonia',
-            #     7: 'Pneumothorax'
-            # }
-    
-            # if pred_label_confidence < 0.2:
-            #     st.write("May not detect a disease.")
-            # else:
-            #     pred_label_name = label_mapping[pred_label]
-            #     st.write(f"Prediction Label: {pred_label_name}")
-            #     st.write(f"Prediction Bounding Box: {pred_bbox}")
-            #     st.write(f"Prediction Confidence: {pred_label_confidence:.2f}")
-    
-            #     output_image = draw_bbox(image.copy(), pred_bbox)
-            #     st.image(output_image, caption='Detected Image.', use_column_width=True)
-
-
-        # if st.button('Auto Detect'):
-        #     st.write("Processing...")
-        #     input_image = image
-        #     # input_image = enhancement_type
-        #     # input_image = cv2.resize(enhanced_image, (W, H))  # Resize the enhanced image to the required input size
-        #     # input_image = (input_image - 127.5) / 127.5  # Normalize to [-1, +1]
-        #     # input_image = np.expand_dims(input_image, axis=0).astype(np.float32)  # Expand dimensions and convert to float32
-            
-        #     pred_bbox, pred_label, pred_label_confidence = predict(model_detection, input_image)
-        
-        #     # Updated label mapping based on the dataset
-        #     label_mapping = {
-        #         0: 'Atelectasis',
-        #         1: 'Cardiomegaly',
-        #         2: 'Effusion',
-        #         3: 'Infiltrate',
-        #         4: 'Mass',
-        #         5: 'Nodule',
-        #         6: 'Pneumonia',
-        #         7: 'Pneumothorax'
-        #     }
-        
-        #     if pred_label_confidence < 0.2:
-        #         st.write("May not detect a disease.")
-        #     else:
-        #         pred_label_name = label_mapping[pred_label]
-        #         st.write(f"Prediction Label: {pred_label_name}")
-        #         st.write(f"Prediction Bounding Box: {pred_bbox}")
-        #         st.write(f"Prediction Confidence: {pred_label_confidence:.2f}")
-        
-        #         output_image = draw_bbox(image.copy(), pred_bbox)
-        #         st.image(output_image, caption='Detected Image.', use_column_width=True)
-    
     with col2:
         if st.button('Generate Grad-CAM'):
             model=load_gradcam_model()
@@ -627,14 +513,10 @@ st.write("Upload a chest X-ray image and click on 'Detect' to find out if there'
 
 model_detection = load_model_detection()
 
-# uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["jpg", "jpeg", "png", "dcm"])
-
 if uploaded_detection is not None:
     file_bytes = np.asarray(bytearray(uploaded_detection.read()), dtype=np.uint8)
     image = cv2.imdecode(file_bytes, 1)
 
-    # st.image(image, caption='Uploaded Image.', use_column_width=True)
-
     if st.button('Detect'):
         st.write("Processing...")
         input_image = preprocess_image(image)