Revised detection.

app-streamlit.py

@@ -27,56 +27,87 @@ bucket_result = storage_client.bucket(bucket_name)
 bucket_name_load = "da-ml-models"
 bucket_load = storage_client.bucket(bucket_name_load)
 
-model_path = os.path.join("model.h5")
-model = tf.keras.models.load_model(model_path)
+H = 224
+W = 224
 
-H, W = 512, 512
-
-test_samples_folder = 'object_detection_test_samples'
-
-def cal_iou(y_true, y_pred):
-    x1 = max(y_true[0], y_pred[0])
-    y1 = max(y_true[1], y_pred[1])
-    x2 = min(y_true[2], y_pred[2])
-    y2 = min(y_true[3], y_pred[3])
-
-    intersection_area = max(0, x2 - x1 + 1) * max(0, y2 - y1 + 1)
-
-    true_area = (y_true[2] - y_true[0] + 1) * (y_true[3] - y_true[1] + 1)
-    bbox_area = (y_pred[2] - y_pred[0] + 1) * (y_pred[3] - y_pred[1] + 1)
+@st.cache_resource
+def load_model():
+    model = tf.keras.models.load_model("model.h5", compile=False)
+    model.compile(
+        loss={
+            "bbox": "mse",
+            "class": "sparse_categorical_crossentropy"
+        },
+        optimizer=tf.keras.optimizers.Adam(),
+        metrics={
+            "bbox": ['mse'],
+            "class": ['accuracy']
+        }
+    )
+    return model
 
-    iou = intersection_area / float(true_area + bbox_area - intersection_area)
-    return iou
+def preprocess_image(image):
+    """ Preprocess the image to the required size and normalization. """
+    image = cv2.resize(image, (W, H))
+    image = (image - 127.5) / 127.5  # Normalize to [-1, +1]
+    image = np.expand_dims(image, axis=0).astype(np.float32)
+    return image
 
-df = pd.read_excel('BBox_List_2017.xlsx')
-labels_dict = dict(zip(df['Image Index'], df['Finding Label']))
+def predict(model, image):
+    """ Predict bounding box and label for the input image. """
+    pred_bbox, pred_class = model.predict(image)
+    pred_label_confidence = np.max(pred_class, axis=1)[0]
+    pred_label = np.argmax(pred_class, axis=1)[0]
+    return pred_bbox[0], pred_label, pred_label_confidence
 
-def predict(image):
-    H, W = 512, 512
+def draw_bbox(image, bbox):
+    """ Draw bounding box on the image. """
+    h, w, _ = image.shape
+    x1, y1, x2, y2 = bbox
+    x1, y1, x2, y2 = int(x1 * w), int(y1 * h), int(x2 * w), int(y2 * h)
+    image = cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
+    return image
 
-    image_resized = cv2.resize(image, (W, H))
-    image_normalized = (image_resized - 127.5) / 127.5
-    image_normalized = np.expand_dims(image_normalized, axis=0)
+st.title("Chest X-ray Disease Detection")
 
-    # Prediction
-    pred_bbox = model.predict(image_normalized, verbose=0)[0]
+st.write("Upload a chest X-ray image and click on 'Detect' to find out if there's any disease.")
 
-    # Rescale the bbox points
-    pred_x1 = int(pred_bbox[0] * image.shape[1])
-    pred_y1 = int(pred_bbox[1] * image.shape[0])
-    pred_x2 = int(pred_bbox[2] * image.shape[1])
-    pred_y2 = int(pred_bbox[3] * image.shape[0])
+model = load_model()
 
-    return (pred_x1, pred_y1, pred_x2, pred_y2)
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
-st.title("AI Integration for Chest X-Ray Imaging")
+if uploaded_file is not None:
+    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)
+
+    if st.button('Detect'):
+        st.write("Processing...")
+        input_image = preprocess_image(image)
+        pred_bbox, pred_label, pred_label_confidence = predict(model, input_image)
+
+        label_mapping = {
+            0: 'Atelectasis',
+            1: 'Cardiomegaly',
+            2: 'Effusion',
+            3: 'Infiltrate',
+            4: 'Mass',
+            5: 'Nodule',
+            6: 'Pneumonia',
+            7: 'Pneumothorax'
+        }
+
+        if pred_label_confidence < 0.01:
+            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}")
 
-# Concept 1: Select from test samples
-# st.header("Select Test Sample Images")
-# test_sample_images = [os.path.join(test_samples_folder, f) for f in os.listdir(test_samples_folder) if f.endswith('.jpg') or f.endswith('.png')]
-# test_sample_selected = st.selectbox("Select a test sample image", test_sample_images)
-# if test_sample_selected:
-#     st.image(test_sample_selected, caption='Selected Test Sample Image', use_column_width=True)
+            output_image = draw_bbox(image.copy(), pred_bbox)
+            st.image(output_image, caption='Detected Image.', use_column_width=True)
     
 
 # Utility Functions