Update app.py

app.py

@@ -1,62 +1,46 @@
 import os
 import numpy as np
-import pandas as pd
 import matplotlib.pyplot as plt
-from sklearn.ensemble import RandomForestClassifier
-from sklearn.model_selection import train_test_split
-from sklearn.metrics import accuracy_score
-import gradio as gr
 import joblib
-
-# Step 1: Generate or Load Sample Data
-np.random.seed(42)
-data = pd.DataFrame({
-    "Pressure": np.random.randint(50, 200, 200),
-    "Temperature": np.random.randint(300, 700, 200),
-    "Material_Type": np.random.randint(1, 5, 200),
-    "Defect_Type": np.random.choice([0, 1, 2], 200, p=[0.6, 0.3, 0.1])
-})
-
-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)
-
-model = RandomForestClassifier(n_estimators=100, random_state=42)
-model.fit(X_train, y_train)
-joblib.dump(model, "defect_model.pkl")
-y_pred = model.predict(X_test)
-print(f"Model Accuracy: {accuracy_score(y_test, y_pred):.2f}")
+import gradio as gr
 
 # Visualization Function
 def plot_simulation(pressure, temperature, material_type):
     if material_type <= 0 or temperature <= 0:
-        raise ValueError("Material type and temperature must be greater than zero.")
+        raise ValueError("Temperature and Material Type must be positive.")
     stress = np.linspace(pressure * 0.5, pressure * 1.5, 100)
     strain = stress / (temperature * 0.01 * material_type)
-    plt.figure(figsize=(8, 5))
+    plt.figure(figsize=(6, 4))
     plt.plot(stress, strain, label="Stress-Strain Curve")
     plt.xlabel("Stress (MPa)")
     plt.ylabel("Strain (%)")
     plt.title("Stress-Strain Simulation")
     plt.legend()
     plt.grid(True)
-    plot_path = "simulation_plot.png"
-    plt.savefig(plot_path)
+    output_path = "simulation_plot.png"
+    plt.savefig(output_path)
     plt.close()
-    return plot_path
+    return output_path
 
 # Prediction Function
 def predict_and_visualize(pressure, temperature, material_type):
-    if not os.path.exists("defect_model.pkl"):
-        return "Error: Model file not found. Train the model first.", None
+    model_path = "defect_model.pkl"
+    if not os.path.exists(model_path):
+        return "Error: Model file not found.", None
+    try:
+        model = joblib.load(model_path)
+        prediction = model.predict([[pressure, temperature, material_type]])[0]
+        defect_map = {0: "No Defect", 1: "Defect"}
+        defect_prediction = defect_map.get(prediction, "Unknown")
+    except Exception as e:
+        return f"Prediction Error: {str(e)}", None
+    
+    try:
+        plot_path = plot_simulation(pressure, temperature, material_type)
+    except Exception as e:
+        return defect_prediction, f"Visualization Error: {str(e)}"
     
-    loaded_model = joblib.load("defect_model.pkl")
-    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]
-    plot_path = plot_simulation(pressure, temperature, material_type)
-    return f"Predicted Defect: {defect_prediction}", plot_path
+    return defect_prediction, plot_path
 
 # Gradio Interface
 interface = gr.Interface(
@@ -71,7 +55,7 @@ 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."
+    description="Predict defect type and visualize stress-strain simulation."
 )
 
 if __name__ == "__main__":