Update app.py

app.py

@@ -1,72 +1,71 @@
+import os
+import sys
 import streamlit as st
 import pandas as pd
 import pickle
-from models.optimizer import optimize_design
-import pandas as pd
-import os
-import sys
 
 # Ensure the project root is in the Python path
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
+# Import optimization logic
 from models.optimizer import optimize_design
 
-def main():
-    print("Welcome to the Press Tool AI application!")
-
-    # Load sample data
-    data_path = "data/sample_data.csv"
-    data = pd.read_csv(data_path)
-    print("Sample Data Loaded:")
-    print(data.head())
-
-    # Call optimization logic
-    optimized_result = optimize_design(data)
-    print("Optimization Complete. Results:")
-    print(optimized_result)
-
-if __name__ == "__main__":
-    main()
-
 
 # Load pre-trained model
 def load_model():
+    """
+    Load the pre-trained defect prediction model from the models directory.
+    """
     with open("models/defect_model.pkl", "rb") as file:
         model = pickle.load(file)
     return model
 
-# Predict defects
+
+# Predict defects using the loaded model
 def predict_defects(model, data):
+    """
+    Predict defect types based on input data using the pre-trained model.
+    """
     predictions = model.predict(data)
     return predictions
 
-st.title("Press Tool AI: Defect Prediction and Optimization")
 
-# File upload
-uploaded_file = st.file_uploader("Upload Design Parameters (CSV)", type="csv")
-if uploaded_file:
-    data = pd.read_csv(uploaded_file)
-    st.write("Uploaded Data:")
-    st.dataframe(data)
+def main():
+    """
+    Main Streamlit web app for defect prediction and design optimization.
+    """
+    st.title("Press Tool AI: Defect Prediction and Optimization")
+
+    # File upload
+    uploaded_file = st.file_uploader("Upload Design Parameters (CSV)", type="csv")
+    if uploaded_file:
+        # Load uploaded CSV data
+        data = pd.read_csv(uploaded_file)
+        st.write("Uploaded Data:")
+        st.dataframe(data)
+
+        # Load pre-trained defect prediction model
+        model = load_model()
+
+        # Predict defects
+        st.subheader("Defect Predictions:")
+        predictions = predict_defects(model, data)
+        data['Predicted Defects'] = predictions
+        st.dataframe(data)
+
+        # Optimize design parameters
+        st.subheader("Optimized Parameters:")
+        optimized_data = optimize_design(data)
+        st.dataframe(optimized_data)
 
-    # Load AI model
-    model = load_model()
-    
-    # Predict defects
-    st.subheader("Defect Predictions:")
-    predictions = predict_defects(model, data)
-    data['Predicted Defects'] = predictions
-    st.dataframe(data)
+        # Provide a download button for optimized results
+        st.download_button(
+            label="Download Optimized Results",
+            data=optimized_data.to_csv(index=False),
+            file_name="optimized_design.csv",
+            mime="text/csv",
+        )
 
-    # Optimize design
-    st.subheader("Optimized Parameters:")
-    optimized_data = optimize_design(data)
-    st.dataframe(optimized_data)
 
-    # Download results
-    st.download_button(
-        label="Download Results",
-        data=optimized_data.to_csv(index=False),
-        file_name="optimized_design.csv",
-        mime="text/csv",
-    )
+if __name__ == "__main__":
+    main()