Update main.py

main.py

@@ -19,9 +19,10 @@ def load_model():
     return num_imputer, scaler, model
 
 def preprocess_input_data(input_data, num_imputer, scaler):
-    input_data_df = pd.DataFrame([input_data])
-    input_data_imputed_num = num_imputer.transform(input_data_df)
-    input_scaled_df = pd.DataFrame(scaler.transform(input_data_imputed_num), columns=input_data_df.columns)
+    input_data_df = pd.DataFrame(input_data)
+    num_columns = [col for col in input_data_df.columns if input_data_df[col].dtype != 'object']
+    input_data_imputed_num = num_imputer.transform(input_data_df[num_columns])
+    input_scaled_df = pd.DataFrame(scaler.transform(input_data_imputed_num), columns=num_columns)
     return input_scaled_df
 
 @app.get("/")
@@ -33,22 +34,21 @@ def read_root():
     return info.strip()
 
 
-
 @app.get("/sepsis/predict")
 def predict_sepsis_endpoint(PRG: float, PL: float, PR: float, SK: float, TS: float,
                             M11: float, BD2: float, Age: float, Insurance: int):
     num_imputer, scaler, model = load_model()
 
     input_data = {
-        'PRG': [PRG],
-        'PL': [PL],
-        'PR': [PR],
-        'SK': [SK],
-        'TS': [TS],
-        'M11': [M11],
-        'BD2': [BD2],
-        'Age': [Age],
-        'Insurance': [Insurance]
+        'PRG': PRG,
+        'PL': PL,
+        'PR': PR,
+        'SK': SK,
+        'TS': TS,
+        'M11': M11,
+        'BD2': BD2,
+        'Age': Age,
+        'Insurance': Insurance
     }
 
     input_scaled_df = preprocess_input_data(input_data, num_imputer, scaler)