Update app.py

app.py

@@ -5,12 +5,20 @@ import streamlit as st
 model = joblib.load('model.joblib')
 unique_values = joblib.load('unique_values.joblib')
 
+unique_Room =  unique_values["Room"]
+unique_Partitions =  unique_values["Partitions"]
+unique_Rows =  unique_values["Rows"]
+unique_Columns =  unique_values["Columns"]
 unique_Color =  unique_values["Color"]
 
 def main():
     st.title("Colors Prediction")
 
     with st.form("questionaire"):
+        room = st.selectbox("Room", unique_Room)
+        partition = st.selectbox("Partitions", unique_Partitions)
+        rows = st.selectbox("Rows", unique_Rows)
+        columns = st.selectbox("Columns", unique_Columns)
         color = st.selectbox("Color", unique_Color)
         
         clicked = st.form_submit_button("Predict Color")
@@ -20,16 +28,18 @@ def main():
                                                "Rows": [Rows],
                                                "Columns": [Columns],
                                                "Color": [Color]}))
-            if result == 'Red':
-                print('Red')
-            elif result == 'Blue':
-                print('Blue')
-            elif result == 'Yellow':
-                print('Yellow')
-            else:
-                print('Green')
+            predicted_color = predict_color(result)
             #result = 'Red' = Red if result[0] == 1 else 'Blue'
             st.success('The predicted color is {}'.format(result))
-
+def predict_color(prediction):
+    if prediction == 'Red':
+        return 'Red'
+    elif prediction == 'Blue':
+        return 'Blue'
+    elif prediction == 'Yellow':
+        return 'Yellow'
+    else:
+        return 'Green'
+        
 if __name__=='__main__':
     main()