Update app.py

app.py

@@ -100,24 +100,31 @@ def load_model_and_encodings():
         original_data = load_datasets()
         
         label_encoders = {}
-        categorical_features = ['Make', 'Model', 'Condition', 'Fuel', 'Title_status', 
-                                'Transmission', 'Drive', 'Size', 'Type', 'Paint_color']
+        categorical_features = original_data.select_dtypes(include=['object']).columns.tolist()
         
         for feature in categorical_features:
-            if feature in original_data.columns:
-                le = LabelEncoder()
-                unique_values = original_data[feature].fillna('unknown').str.strip().unique()
-                le.fit(unique_values)
-                label_encoders[feature.lower()] = le
+            le = LabelEncoder()
+            unique_values = original_data[feature].fillna('unknown').str.strip().unique()
+            le.fit(unique_values)
+            label_encoders[feature.lower()] = le
         
-        return model, label_encoders
+        return model, label_encoders, categorical_features
     except Exception as e:
         st.error(f"Error loading model: {str(e)}")
         raise e
 
-def predict_price(model, encoders, user_input):
-    encoded_features = {feature: encoders[feature].transform([value])[0] if value in encoders[feature] else 0 
-                        for feature, value in user_input.items()}
+def predict_price(model, encoders, categorical_features, user_input):
+    encoded_features = {}
+    for feature, value in user_input.items():
+        if feature.lower() in encoders:
+            encoded_features[feature] = encoders[feature.lower()].transform([value])[0]
+        elif feature in categorical_features:
+            # If it's a categorical feature but not in encoders, set to 0 (unknown)
+            encoded_features[feature] = 0
+        else:
+            # For numerical features, use the value as is
+            encoded_features[feature] = value
+    
     input_data = pd.DataFrame([encoded_features])
     predicted_price = model.predict(input_data)
     return predicted_price[0]
@@ -127,7 +134,7 @@ st.title("Auto Appraise")
 st.write("Upload a car image or take a picture to get its brand, model, overview, and expected price!")
 
 # Load model and encoders
-model, label_encoders = load_model_and_encodings()
+model, label_encoders, categorical_features = load_model_and_encodings()
 
 # Initialize OpenAI API key
 openai.api_key = st.secrets["GPT_TOKEN"]
@@ -192,21 +199,21 @@ if image is not None:
 
             for year in years:
                 user_input = {
-                    'Make': match['Make'],
-                    'Model': match['Model'],
-                    'Condition': match['Condition'],
-                    'Fuel': match['Fuel'],
-                    'Title_status': match['Title_status'],
-                    'Transmission': match['Transmission'],
-                    'Drive': match['Drive'],
-                    'Size': match['Size'],
-                    'Type': match['Type'],
-                    'Paint_color': match['Paint_color'],
-                    'Year': year
+                    'make': match['Make'],
+                    'model': match['Model'],
+                    'condition': match.get('Condition', 'unknown'),
+                    'fuel': match.get('Fuel', 'unknown'),
+                    'title_status': match.get('Title_status', 'unknown'),
+                    'transmission': match.get('Transmission', 'unknown'),
+                    'drive': match.get('Drive', 'unknown'),
+                    'size': match.get('Size', 'unknown'),
+                    'type': match.get('Type', 'unknown'),
+                    'paint_color': match.get('Paint_color', 'unknown'),
+                    'year': year
                 }
                 
-                price = predict_price(model, label_encoders, user_input)
-                predicted_prices.append(price)
+            price = predict_price(model, label_encoders, categorical_features, user_input)
+            predicted_prices.append(price)
 
             # Plotting the results
             plt.figure(figsize=(10, 5))