jkim03's picture
Update app.py
6fb851a verified
from fastapi import FastAPI
from pydantic import BaseModel
import joblib
import pandas as pd
import json
# Load models
scaler = joblib.load("scaler.joblib")
gmm = joblib.load("gmm_model.joblib")
pca = joblib.load("pca.joblib")
with open("cluster_fatigue_map.json") as f:
cluster_fatigue_map = json.load(f)
# Load trained regressor
model = joblib.load("fatigue_model.joblib")
feature_cols = [
"AVRR", "SDNN", "RMSSD", "PNN50", "Coefficient_of_Variation",
"Age", "Weight", "Height"
]
# Define input schema
class FatigueInput(BaseModel):
AVRR: float
SDNN: float
RMSSD: float
PNN50: float
Coefficient_of_Variation: float
Age: int
Weight: float
Height: float
# Create FastAPI app
app = FastAPI()
@app.post("/debug")
def debug_predict(input_data: FatigueInput):
input_dict = input_data.dict()
input_df = pd.DataFrame([input_dict])[feature_cols]
scaled_input = scaler.transform(input_df)
pca_input = pca.transform(scaled_input)
return {
"raw_input": input_dict,
"scaled_input": scaled_input.tolist(),
"pca_input": pca_input.tolist(),
"scaler_mean": scaler.mean_.tolist(),
"scaler_scale": scaler.scale_.tolist()
}
@app.post("/predictGMM")
def predict(input_data: FatigueInput):
try:
input_dict = input_data.dict()
print("INPUT:", input_dict)
input_df = pd.DataFrame([input_dict])[feature_cols]
scaled_input = scaler.transform(input_df)
print("SCALED:", scaled_input)
pca_input = pca.transform(scaled_input)
print("PCA:", pca_input)
cluster = gmm.predict(pca_input)[0]
probs = gmm.predict_proba(pca_input)[0].astype(float)
print("CLUSTER:", cluster)
print("PROBS:", probs)
fatigue_level = cluster_fatigue_map[str(cluster)]
return {
"cluster": int(cluster),
"fatigue_level": fatigue_level,
"cluster_probabilities": {str(i): float(prob) for i, prob in enumerate(probs)},
"weighted_fatigue_score": sum(float(probs[i]) * cluster_fatigue_map[str(i)] for i in range(len(probs)))
}
except Exception as e:
return {"error": str(e)}
@app.post("/predict")
def predict(input_data: FatigueInput):
try:
input_df = pd.DataFrame([input_data.dict()])[feature_cols]
# Predict fatigue level
predicted_fatigue = model.predict(input_df)[0]
return {
"predicted_fatigue_level": float(predicted_fatigue),
"fatigueLevel": int(round(predicted_fatigue))
}
except Exception as e:
return {"error": str(e)}