Update app.py

app.py

@@ -7,81 +7,140 @@ import soundfile as sf
 model_path = 'model.onnx'  # Replace with the actual path to your model
 model_inference = rt.InferenceSession(model_path)
 
-# Streamlit UI
-st.title('TB Cough Sound Analysis')
-
-# Audio file uploader
-audio_file = st.file_uploader("Upload a cough sound recording", type=['wav', 'mp3'])
-
-# Clinical data input fields
-age = st.number_input('Age', min_value=0, step=1, format='%d')
-height = st.number_input('Height (in cm)', min_value=0.0, step=0.1, format='%f')
-weight = st.number_input('Weight (in kg)', min_value=0.0, step=0.1, format='%f')
-reported_cough_dur = st.number_input('Reported Cough Duration (in days)', min_value=0, step=1, format='%d')
-heart_rate = st.number_input('Heart Rate (beats per minute)', min_value=0, step=1, format='%d')
-temperature = st.number_input('Temperature (in Celsius)', min_value=0.0, step=0.1, format='%f')
-
-sex = st.radio('Sex', ['Male', 'Female'])
-tb_prior = st.radio('Previous TB', ['No', 'Not sure', 'Yes'])
-tb_prior_Pul = st.radio('Previous Pulmonary TB', ['No', 'Yes'])
-tb_prior_Extrapul = st.radio('Previous Extrapulmonary TB', ['No', 'Yes'])
-tb_prior_Unknown = st.radio('Previous TB Unknown', ['No', 'Yes'])
-hemoptysis = st.radio('Hemoptysis', ['No', 'Yes'])
-weight_loss = st.radio('Weight Loss', ['No', 'Yes'])
-smoke_lweek = st.radio('Smoking Last Week', ['No', 'Yes'])
-fever = st.radio('Fever', ['No', 'Yes'])
-night_sweats = st.radio('Night Sweats', ['No', 'Yes'])
-
-# Convert categorical inputs to binary
-sex_Female = 1 if sex == 'Female' else 0
-sex_Male = 1 if sex == 'Male' else 0
-tb_prior_No = 1 if tb_prior == 'No' else 0
-tb_prior_Not_sure = 1 if tb_prior == 'Not sure' else 0
-tb_prior_Yes = 1 if tb_prior == 'Yes' else 0
-tb_prior_Pul_No = 1 if tb_prior_Pul == 'No' else 0
-tb_prior_Pul_Yes = 1 if tb_prior_Pul == 'Yes' else 0
-tb_prior_Extrapul_No = 1 if tb_prior_Extrapul == 'No' else 0
-tb_prior_Extrapul_Yes = 1 if tb_prior_Extrapul == 'Yes' else 0
-tb_prior_Unknown_No = 1 if tb_prior_Unknown == 'No' else 0
-tb_prior_Unknown_Yes = 1 if tb_prior_Unknown == 'Yes' else 0
-hemoptysis_No = 1 if hemoptysis == 'No' else 0
-hemoptysis_Yes = 1 if hemoptysis == 'Yes' else 0
-weight_loss_No = 1 if weight_loss == 'No' else 0
-weight_loss_Yes = 1 if weight_loss == 'Yes' else 0
-smoke_lweek_No = 1 if smoke_lweek == 'No' else 0
-smoke_lweek_Yes = 1 if smoke_lweek == 'Yes' else 0
-fever_No = 1 if fever == 'No' else 0
-fever_Yes = 1 if fever == 'Yes' else 0
-night_sweats_No = 1 if night_sweats == 'No' else 0
-night_sweats_Yes = 1 if night_sweats == 'Yes' else 0
-
-# Combine all inputs
-clinical_data = [age, height, weight, reported_cough_dur, heart_rate, temperature,
-                 sex_Female, sex_Male, tb_prior_No, tb_prior_Not_sure, tb_prior_Yes,
-                 tb_prior_Pul_No, tb_prior_Pul_Yes, tb_prior_Extrapul_No, tb_prior_Extrapul_Yes,
-                 tb_prior_Unknown_No, tb_prior_Unknown_Yes, hemoptysis_No, hemoptysis_Yes,
-                 weight_loss_No, weight_loss_Yes, smoke_lweek_No, smoke_lweek_Yes,
-                 fever_No, fever_Yes, night_sweats_No, night_sweats_Yes]
-clinical_data = np.array(clinical_data).reshape(1, -1).astype(np.float32)
-
-# Button to make prediction
-if st.button('Predict'):
-    if audio_file is not None:
-        # Process audio file
-        raw_values, rate = sf.read(audio_file)
-        mu = np.nanmean(raw_values)
-        std = np.nanstd(raw_values)
-        audio_data = (raw_values - mu) / std
-        audio_data = np.pad(audio_data, (0, 22050 - len(audio_data)), 'constant').reshape(1, -1, 1).astype(np.float32)
-        
-        # Make prediction
-        input_name = model_inference.get_inputs()[0].name
-        input_name2 = model_inference.get_inputs()[1].name
-        label_name = model_inference.get_outputs()[0].name
-        onnx_pred = model_inference.run([label_name], {input_name: audio_data, input_name2: clinical_data})
-        result = onnx_pred[0]
-        
-        # Display result
-        st.write('Prediction:', result)
-    else:
-        st.error('Please upload an audio file.')
+# Function to preprocess audio data
+def preprocess_audio(audio_file):
+    raw_values, rate = sf.read(audio_file)
+    mu = np.nanmean(raw_values)
+    std = np.nanstd(raw_values)
+    audio_data = (raw_values - mu) / std
+    audio_data = np.pad(audio_data, (0, 22050 - len(audio_data)), 'constant').reshape(1, -1, 1).astype(np.float32)
+    return audio_data
+
+# Function to preprocess clinical data
+def preprocess_clinical_data(age, height, weight, reported_cough_dur, heart_rate, temperature, sex, tb_prior, tb_prior_Pul, tb_prior_Extrapul, tb_prior_Unknown, hemoptysis, weight_loss, smoke_lweek, fever, night_sweats):
+    sex_Female = 1 if sex == 'Female' else 0
+    sex_Male = 1 if sex == 'Male' else 0
+    tb_prior_No = 1 if tb_prior == 'No' else 0
+    tb_prior_Not_sure = 1 if tb_prior == 'Not sure' else 0
+    tb_prior_Yes = 1 if tb_prior == 'Yes' else 0
+    tb_prior_Pul_No = 1 if tb_prior_Pul == 'No' else 0
+    tb_prior_Pul_Yes = 1 if tb_prior_Pul == 'Yes' else 0
+    tb_prior_Extrapul_No = 1 if tb_prior_Extrapul == 'No' else 0
+    tb_prior_Extrapul_Yes = 1 if tb_prior_Extrapul == 'Yes' else 0
+    tb_prior_Unknown_No = 1 if tb_prior_Unknown == 'No' else 0
+    tb_prior_Unknown_Yes = 1 if tb_prior_Unknown == 'Yes' else 0
+    hemoptysis_No = 1 if hemoptysis == 'No' else 0
+    hemoptysis_Yes = 1 if hemoptysis == 'Yes' else 0
+    weight_loss_No = 1 if weight_loss == 'No' else 0
+    weight_loss_Yes = 1 if weight_loss == 'Yes' else 0
+    smoke_lweek_No = 1 if smoke_lweek == 'No' else 0
+    smoke_lweek_Yes = 1 if smoke_lweek == 'Yes' else 0
+    fever_No = 1 if fever == 'No' else 0
+    fever_Yes = 1 if fever == 'Yes' else 0
+    night_sweats_No = 1 if night_sweats == 'No' else 0
+    night_sweats_Yes = 1 if night_sweats == 'Yes' else 0
+
+    clinical_data = [age, height, weight, reported_cough_dur, heart_rate, temperature,
+                     sex_Female, sex_Male, tb_prior_No, tb_prior_Not_sure, tb_prior_Yes,
+                     tb_prior_Pul_No, tb_prior_Pul_Yes, tb_prior_Extrapul_No, tb_prior_Extrapul_Yes,
+                     tb_prior_Unknown_No, tb_prior_Unknown_Yes, hemoptysis_No, hemoptysis_Yes,
+                     weight_loss_No, weight_loss_Yes, smoke_lweek_No, smoke_lweek_Yes,
+                     fever_No, fever_Yes, night_sweats_No, night_sweats_Yes]
+    clinical_data = np.array(clinical_data).reshape(1, -1).astype(np.float32)
+    return clinical_data
+
+# Main function to run the app
+def main():
+    st.title('TB Cough Sound Analysis')
+
+    # Initialize state variables
+    page = st.session_state.get('page', 0)
+    st.session_state['page'] = page
+
+    # Create multi-step form
+    if page == 0:
+        st.write('**Step 1: Upload Cough Sound Recording**')
+        audio_file = st.file_uploader("Upload a cough sound recording", type=['wav', 'mp3'])
+        if audio_file is not None:
+            st.audio(audio_file)
+            st.session_state['audio_file'] = audio_file
+            if st.button('Next'):
+                st.session_state['page'] = 1
+
+    elif page == 1:
+        st.write('**Step 2: Enter Clinical Data**')
+        age = st.number_input('Age', min_value=0, step=1, format='%d')
+        height = st.number_input('Height (in cm)', min_value=0.0, step=0.1, format='%f')
+        weight = st.number_input('Weight (in kg)', min_value=0.0, step=0.1, format='%f')
+        reported_cough_dur = st.number_input('Reported Cough Duration (in days)', min_value=0, step=1, format='%d')
+        heart_rate = st.number_input('Heart Rate (beats per minute)', min_value=0, step=1, format='%d')
+        temperature = st.number_input('Body Temperature (in °C)', min_value=0.0, step=0.1, format='%f')
+        sex = st.radio('Sex', ['Male', 'Female'])
+        tb_prior = st.selectbox('TB Prior', ['No', 'Not sure', 'Yes'])
+        tb_prior_Pul = st.selectbox('TB Prior - Pulmonary', ['No', 'Yes'])
+        tb_prior_Extrapul = st.selectbox('TB Prior - Extrapulmonary', ['No', 'Yes'])
+        tb_prior_Unknown = st.selectbox('TB Prior - Unknown', ['No', 'Yes'])
+        hemoptysis = st.selectbox('Hemoptysis', ['No', 'Yes'])
+        weight_loss = st.selectbox('Weight Loss', ['No', 'Yes'])
+        smoke_lweek = st.selectbox('Smoking Last Week', ['No', 'Yes'])
+        fever = st.selectbox('Fever', ['No', 'Yes'])
+        night_sweats = st.selectbox('Night Sweats', ['No', 'Yes'])
+
+        if st.button('Next'):
+            st.session_state['page'] = 2
+
+    elif page == 2:
+        st.write('**Step 3: Review Input Data**')
+        st.write('**Clinical Data:**')
+        st.write(f"- Age: {age}")
+        st.write(f"- Height: {height} cm")
+        st.write(f"- Weight: {weight} kg")
+        st.write(f"- Reported Cough Duration: {reported_cough_dur} days")
+        st.write(f"- Heart Rate: {heart_rate} bpm")
+        st.write(f"- Temperature: {temperature} °C")
+        st.write(f"- Sex: {sex}")
+        st.write(f"- TB Prior: {tb_prior}")
+        st.write(f"- TB Prior - Pulmonary: {tb_prior_Pul}")
+        st.write(f"- TB Prior - Extrapulmonary: {tb_prior_Extrapul}")
+        st.write(f"- TB Prior - Unknown: {tb_prior_Unknown}")
+        st.write(f"- Hemoptysis: {hemoptysis}")
+        st.write(f"- Weight Loss: {weight_loss}")
+        st.write(f"- Smoking Last Week: {smoke_lweek}")
+        st.write(f"- Fever: {fever}")
+        st.write(f"- Night Sweats: {night_sweats}")
+
+        if st.button('Next'):
+            st.session_state['page'] = 3
+
+    elif page == 3:
+        st.write('**Step 4: Prediction**')
+        if 'audio_file' in st.session_state:
+            audio_data = preprocess_audio(st.session_state['audio_file'])
+            clinical_data = preprocess_clinical_data(age, height, weight, reported_cough_dur, heart_rate, temperature,
+                                                     sex, tb_prior, tb_prior_Pul, tb_prior_Extrapul, tb_prior_Unknown, hemoptysis, weight_loss, smoke_lweek, fever,
+                                                     night_sweats)
+            input_name = model_inference.get_inputs()[0].name
+            input_name2 = model_inference.get_inputs()[1].name
+            label_name = model_inference.get_outputs()[0].name
+            onnx_pred = model_inference.run([label_name], {input_name: audio_data, input_name2: clinical_data})
+            result = onnx_pred[0]
+
+            # Interpret prediction
+            if result >= 0.5:
+                prediction_text = "Tuberculosis is found based on the audio data and clinical information."
+                probability_label = "Positive Probability"
+            else:
+                prediction_text = "No Tuberculosis is found based on the audio data and clinical information."
+                probability_label = "Negative Probability"
+
+            st.write('**Prediction:**')
+            st.write(f"{prediction_text}")
+            st.write(f"{probability_label}: {result}")
+        else:
+            st.error('Please upload an audio file.')
+
+        if st.button('Previous'):
+            st.session_state['page'] = 2
+
+if __name__ == '__main__':
+    main()