safety_api / main.py
mahm123's picture
Upload 4 files
b43bf6a verified
Raw
History Blame Contribute Delete
1.25 kB
import tensorflow as tf
from fastapi import FastAPI
from pydantic import BaseModel
import numpy as np
import os
app = FastAPI()
# Load your specific model
MODEL_PATH = os.path.join(os.path.dirname(__file__), "driving_behavior_model.keras")
model = tf.keras.models.load_model(MODEL_PATH)
# Define the labels used in your notebook
CLASSES = ["Safe", "Moderate", "Dangerous"]
class ModelInput(BaseModel):
# Expecting a list of 10 lists, each containing 6 sensor values
# Example: [[ax, ay, az, gx, gy, gz], [...], ... (10 times)]
data: list
@app.post("/predict")
async def predict(input_params: ModelInput):
# 1. Convert input to numpy array
input_data = np.array(input_params.data) # Shape: (10, 6)
# 2. Add the batch dimension to make it (1, 10, 6)
input_data = np.expand_dims(input_data, axis=0)
# 3. Run prediction
prediction_probs = model.predict(input_data)
# 4. Get the index of the highest probability
predicted_class_index = np.argmax(prediction_probs[0])
# 5. Return readable result
return {
"prediction_index": int(predicted_class_index),
"label": CLASSES[predicted_class_index],
"confidence": float(np.max(prediction_probs[0]))
}