app.py fix

app.py

@@ -94,6 +94,12 @@ def make_prediction(alcohol, arrest_type, belts, contributed_to_accident, disobe
     prediction = tf["LabelEncoder"].inverse_transform(np.argmax(prediction, axis=1))
 
     # Return the prediction
+    # if prediction[0] == 0:
+    #     return 'SERO'
+    # elif prediction[0] == 1:
+    #     return 'Warning'
+    # else:
+    #     return 'Citation'
     return prediction[0]
 
 
@@ -114,12 +120,13 @@ iface = gr.Interface(fn=make_prediction,
                              gr.components.Checkbox(label='Did the violation involve any property damage?'),
                              gr.components.Dropdown(label='Choose the race of the driver', choices=list(tf['OneHotEncoder'].transformers_[0][1].categories_[2]), value='WHITE'),
                              gr.components.Checkbox(label='Did the driver fail to obey signs and markings (such as traffic control device instructions, stop lights, red signal and stop sign lines)?'),
-                             gr.components.Dropdown(label='Choose the race of the driver', choices=list(tf['OneHotEncoder'].transformers_[0][1].categories_[3]), value='NO SEARCH CONDUCTED'),
+                             gr.components.Dropdown(label='What was the outcome of the search (if conducted)?', choices=list(tf['OneHotEncoder'].transformers_[0][1].categories_[3]), value='NO SEARCH CONDUCTED'),
                              gr.components.Checkbox(label='Was the driver caught speeding?'),
                              gr.components.Slider(maximum=23, step=1, label='Time HOUR when stop occurred in 24-hour format'),
                              gr.components.Slider(minimum=2012, maximum=2024, step=1, label='Year when stop occurred'),
                              gr.components.Dropdown(label='What is the name of the subagency that conducted the traffic stop?', choices=list(tf['OneHotEncoder'].transformers_[0][1].categories_[1]), value='4th District, Wheaton'),
                              gr.components.Checkbox(label='Was the vehicle safe and up to standards (lights properly switched, registration plates attached etc.)?'),
+                             gr.components.Dropdown(label='What is the vehicle type?', choices=list(tf['OrdinalEncoder_HighCardinality'].categories_[1]), value='02 - Automobile'),
                              gr.components.Slider(minimum=1970, maximum=2023, step=1, label='Year of manufacture of the vehicle:')],
                      outputs=["text"])