Update app.py

app.py

@@ -454,6 +454,39 @@ if uploaded_file is not None:
             img_array = np.uint8(img_array)  # Convert to uint8
             img = Image.fromarray(img_array)
 
+            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)
+                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)
+
             col1, col2 = st.columns(2)
             # Check the number of dimensions of the image
             if img_array.ndim == 3: