flokabukie's picture
Update main.py
7bbc3ed
from fastapi import FastAPI
from pydantic import BaseModel
import pickle
import pandas as pd
import numpy as np
import uvicorn
import os
from sklearn.preprocessing import StandardScaler
import joblib
app = FastAPI(title="API")
"""We load a machine learning model and a scaler that help us make predictions based on data."""
model = joblib.load('model.pkl',mmap_mode='r')
scaler = joblib.load('scaler.pkl',mmap_mode='r')
def predict(df, endpoint='simple'):
# Scaling
scaled_df = scaler.transform(df)
# Prediction
prediction = model.predict_proba(scaled_df)
highest_proba = prediction.max(axis=1)
predicted_labels = ["Patient does not have sepsis" if i == 0 else "Patient has Sepsis" for i in highest_proba]
response = []
for label, proba in zip(predicted_labels, highest_proba):
output = {
"prediction": label,
"probability of prediction": str(round(proba * 100)) + '%'
}
response.append(output)
return response
class Patient(BaseModel):
Blood_Work_R1: float
Blood_Pressure: float
Blood_Work_R3: float
BMI: float
Blood_Work_R4: float
Patient_age: int
@app.get("/")
def root():
return {"API": "This is an API for sepsis prediction."}
# Prediction endpoint (Where we will input our features)
@app.post("/predict")
def predict_sepsis(patient: Patient):
# Make prediction
data = pd.DataFrame(patient.dict(), index=[0])
scaled_data = scaler.transform(data)
parsed = predict(df=scaled_data)
return {"output": parsed}
if __name__ == "__main__":
os.environ["DEBUG"] = "True" # Enable debug mode
uvicorn.run("main:app", reload=True)