Update app/main.py

app/main.py

@@ -20,67 +20,69 @@ def predict(df, endpoint="simple"):
     scaled_df = scaler.transform(df)  # Scale the input data using a pre-defined scaler
 
     # Prediction
-    prediction = model.predict_proba(scaled_df)  # Make predictions using a pre-trained model
-
-    highest_proba = prediction.max(axis=1)  # Get the highest probability for each prediction
-
-    # Assign predicted labels based on the highest probabilities
-    predicted_labels = ["Patient does not have sepsis" if i == 0 else "Patient has sepsis" for i in highest_proba]
-    print(f"Predicted labels: {predicted_labels}")  # Print the predicted labels to the terminal
-    print(highest_proba)  # Print the highest probabilities to the terminal
+    prediction = model.predict(scaled_df)  # Make predictions using a pre-trained XGBoost regressor model
 
     response = []
-    for label, proba in zip(predicted_labels, highest_proba):
-        # Create a response for each prediction with the predicted label and probability
+    for eta in prediction:
+        # Create a response for each prediction with the predicted ETA
         output = {
-            "prediction": label,
-            "probability of prediction": str(round(proba * 100)) + '%'  # Convert the probability to a percentage
+            "predicted_eta": eta
         }
         response.append(output)  # Add the response to the list of responses
 
     return response  # Return the list of responses
 
 
-class Patient(BaseModel):
-    Blood_Work_R1: int
-    Blood_Pressure: int
-    Blood_Work_R3: int
-    BMI: float
-    Blood_Work_R4: float
-    Patient_age: int
-
-class Patients(BaseModel):
-    all_patients: list[Patient]
+class Trip(BaseModel):
+    Origin_lat: float
+    Origin_lon: float
+    Destination_lat: float
+    Destination_lon: float
+    Trip_distance: int # Assuming this column represents an integer value
+    total_secs: int # Assuming this column represents an integer value
+    dewpoint_2m_temperature: float
+    maximum_2m_air_temperature: float
+    mean_2m_air_temperature: float
+    mean_sea_level_pressure: float
+    minimum_2m_air_temperature: float
+    surface_pressure: float
+    total_precipitation: float
+    u_component_of_wind_10m: float
+    v_component_of_wind_10m: float
+
+class Trips(BaseModel):
+    all_trips: list[Trip]
 
     @classmethod
-    def return_list_of_dict(cls, patients: "Patients"):
-        patient_list = []
-        for patient in patients.all_patients: #for each item in all_patients,
-            patient_dict = patient.dict() #convert to a dictionary
-            patient_list.append(patient_dict) #add it to the empty list called patient_list
-        return patient_list
+    def return_list_of_dict(cls, trips: "Trips"):
+        trip_list = []
+        for trip in trips.all_trips:  # for each item in all_trips
+            trip_dict = trip.dict()    # convert to a dictionary
+            trip_list.append(trip_dict)  # add it to the empty list called trip_list
+        return trip_list
+
     
 # Endpoints
 # Root Endpoint
 @app.get("/")
 def root():
-    return {"Welcome to the Sepsis Prediction API! This API provides endpoints for predicting sepsis based on patient data."}    
+    return {"Welcome to the ETA Prediction API! This API provides endpoints for predicting ETA based on trip data."}    
 
 # Prediction endpoint
 @app.post("/predict")
-def predict_sepsis(patient: Patient):
+def predict_eta(trip: Trip):
     # Make prediction
-    data = pd.DataFrame(patient.dict(), index=[0])
-    parsed = predict(df=data)
-    return {"output": parsed}
+    data = pd.DataFrame(trip.dict(), index=[0])
+    predicted_eta = predict(df=data)
+    return {"predicted_eta": predicted_eta}
 
 # Multiple Prediction Endpoint
 @app.post("/predict_multiple")
-def predict_sepsis_for_multiple_patients(patients: Patients):
+def predict_eta_for_multiple_trips(trips: Trips):
     """Make prediction with the passed data"""
-    data = pd.DataFrame(Patients.return_list_of_dict(patients))
-    parsed = predict(df=data, endpoint="multi")
-    return {"output": parsed}
+    data = pd.DataFrame(Trips.return_list_of_dict(trips))
+    predicted_eta = predict(df=data, endpoint="multi")
+    return {"predicted_eta": predicted_eta}
 
 if __name__ == "__main__":
     uvicorn.run("main:app", reload=True)