Update app.py

app.py

@@ -117,43 +117,59 @@ if st.button("Predict"):
 
             # Show progress spinner while making predictions
             with st.spinner('Making prediction...'):
-                # Predict based on selected row
-                # prediction = model.predict([input_values])
-                # Load the pre-trained scaler and model
-                with open('scaler.pkl', 'rb') as f:
-                    scalar = pickle.load(f)
-                
-                with open('rf_clf.pkl', 'rb') as f:
-                    rf_clf = pickle.load(f)
-
-                            # Scale the new data using the pre-trained scaler
-                X_new_scaled = scalar.transform(single_sample.reshape(1, -1))
-
-                # Make predictions using the pre-trained model
-                prediction = rf_clf.predict(X_new_scaled)
-                y_pred_proba = rf_clf.predict_proba(X_new_scaled)
-                # loaded_model = pickle.load(open('Random_Forest.sav', 'rb'))
-                # prediction = loaded_model.predict(np.array(single_sample))
-
-
-            # st.write(f"Prediction : {prediction[0]}")
-            if prediction[0] == 0:
-                st.success("The website is not a phishing website.")
-                st.markdown(f'<img src="data:image/gif;base64,{data_url_ok}" alt="cat gif">', unsafe_allow_html=True,)
-            else:
-                st.error("The website is a phishing website.")
-                st.markdown(f'<img src="data:image/gif;base64,{data_url_warning}" alt="cat gif">', unsafe_allow_html=True,)
-
-            
-             # Visualize prediction confidence scores as a bar chart
-            st.write("Prediction Confidence Scores:")
-            class_names = rf_clf.classes_
-            plt.figure(figsize=(8, 4))
-            sns.barplot(x=class_names, y=y_pred_proba[0])
-            plt.title("Prediction Confidence Scores")
-            plt.xlabel("Class")
-            plt.ylabel("Probability")
-            st.pyplot(plt)
+                try:
+                    # Load the pre-trained scaler
+                    with open('scaler.pkl', 'rb') as f:
+                        scaler = pickle.load(f)
+                    # Scale the new data
+                    X_new_scaled = scaler.transform(single_sample.reshape(1, -1))
+    
+                    # Load the selected model
+                    if selected_model == 'Logistic Regression':
+                        with open('logistic_regression_model.pkl', 'rb') as f:
+                            model = pickle.load(f)
+                    elif selected_model == 'Decision Tree':
+                        with open('decision_tree_model.pkl', 'rb') as f:
+                            model = pickle.load(f)
+                    elif selected_model == 'KNN':
+                        with open('knn_model.pkl', 'rb') as f:
+                            model = pickle.load(f)
+                    elif selected_model == 'XGBoost':
+                        with open('xgboost_model.pkl', 'rb') as f:
+                            model = pickle.load(f)
+                    elif selected_model == 'Random Forest':
+                        with open('rf_clf.pkl', 'rb') as f:
+                            model = pickle.load(f)
+                    elif selected_model == 'SVM':
+                        with open('svm_model.pkl', 'rb') as f:
+                            model = pickle.load(f)
+    
+                    # Make predictions
+                    prediction = model.predict(X_new_scaled)
+                    y_pred_proba = model.predict_proba(X_new_scaled)
+    
+                    # Display the prediction result
+                    if prediction[0] == 0:
+                        st.success("The website is not a phishing website.")
+                        st.markdown(f'<img src="data:image/gif;base64,{data_url_ok}" alt="cat gif">', unsafe_allow_html=True,)
+                    else:
+                        st.error("The website is a phishing website.")
+                        st.markdown(f'<img src="data:image/gif;base64,{data_url_warning}" alt="cat gif">', unsafe_allow_html=True,)
+    
+                    # Visualize prediction confidence scores as a bar chart
+                    st.write("Prediction Confidence Scores:")
+                    class_names = model.classes_
+                    plt.figure(figsize=(8, 4))
+                    sns.barplot(x=class_names, y=y_pred_proba[0])
+                    plt.title("Prediction Confidence Scores")
+                    plt.xlabel("Class")
+                    plt.ylabel("Probability")
+                    st.pyplot(plt)
+    
+                except FileNotFoundError as e:
+                    st.error(f"Model file for {selected_model} not found: {str(e)}")
+                except Exception as e:
+                    st.error(f"An error occurred while loading the model: {str(e)}")
            
     else:
         st.error("ERROR!!! Please provide web site information for prediction !!!")