Update app.py

app.py

@@ -1,6 +1,4 @@
-# app.py
-
-# Import libraries
+import os
 import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
@@ -19,29 +17,22 @@ data = pd.DataFrame({
     "Defect_Type": np.random.choice([0, 1, 2], 200, p=[0.6, 0.3, 0.1])
 })
 
-# Splitting data
 X = data[["Pressure", "Temperature", "Material_Type"]]
 y = data["Defect_Type"]
 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
 
-# Step 2: Train the Model
 model = RandomForestClassifier(n_estimators=100, random_state=42)
 model.fit(X_train, y_train)
-
-# Save the model
 joblib.dump(model, "defect_model.pkl")
-
-# Test the Model
 y_pred = model.predict(X_test)
 print(f"Model Accuracy: {accuracy_score(y_test, y_pred):.2f}")
 
-# Step 3: Define Helper Functions for Visualizations
+# Visualization Function
 def plot_simulation(pressure, temperature, material_type):
-    # Simulate some stress/strain values for visualization
+    if material_type <= 0 or temperature <= 0:
+        raise ValueError("Material type and temperature must be greater than zero.")
     stress = np.linspace(pressure * 0.5, pressure * 1.5, 100)
     strain = stress / (temperature * 0.01 * material_type)
-    
-    # Create the plot
     plt.figure(figsize=(8, 5))
     plt.plot(stress, strain, label="Stress-Strain Curve")
     plt.xlabel("Stress (MPa)")
@@ -49,29 +40,25 @@ def plot_simulation(pressure, temperature, material_type):
     plt.title("Stress-Strain Simulation")
     plt.legend()
     plt.grid(True)
-    
-    # Save the plot to an image
     plot_path = "simulation_plot.png"
     plt.savefig(plot_path)
     plt.close()
     return plot_path
 
-# Step 4: Prediction and Visualization Function
+# Prediction Function
 def predict_and_visualize(pressure, temperature, material_type):
-    # Load the model
+    if not os.path.exists("defect_model.pkl"):
+        return "Error: Model file not found. Train the model first.", None
+    
     loaded_model = joblib.load("defect_model.pkl")
-    # Predict
     input_data = np.array([[pressure, temperature, material_type]])
     prediction = loaded_model.predict(input_data)[0]
     defect_map = {0: "No Defect", 1: "Crack", 2: "Wrinkle"}
     defect_prediction = defect_map[prediction]
-    
-    # Generate visualization
     plot_path = plot_simulation(pressure, temperature, material_type)
-    
     return f"Predicted Defect: {defect_prediction}", plot_path
 
-# Step 5: Create Gradio Interface with Visualization
+# Gradio Interface
 interface = gr.Interface(
     fn=predict_and_visualize,
     inputs=[
@@ -84,9 +71,8 @@ interface = gr.Interface(
         gr.Image(label="Stress-Strain Visualization")
     ],
     title="Defect Prediction & Simulation",
-    description="Predict defect type and visualize stress-strain simulation for input features: Pressure, Temperature, and Material Type."
+    description="Predict defect type and visualize stress-strain simulation for input features."
 )
 
-# Step 6: Launch the Interface
 if __name__ == "__main__":
     interface.launch()