main.py

from pathlib import Path

import joblib
import pandas as pd
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from schemas import SensorData


app = FastAPI(
    title="Predictive Maintenance API",
    description="API for predicting machine maintenance requirements",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)


BASE_DIR = Path(__file__).resolve().parent

app.mount("/static", StaticFiles(directory=BASE_DIR / "static"), name="static")
MODEL_PATH = BASE_DIR / "models" / "maintenance_model.pkl"
PREPROCESSOR_PATH = BASE_DIR / "models" / "preprocessor.pkl"

templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))

model = joblib.load(MODEL_PATH)
preprocessor = joblib.load(PREPROCESSOR_PATH)


@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    return templates.TemplateResponse(
        request=request,
        name="index.html",
        context={}
    )


@app.post("/predict")
async def predict(data: SensorData):
    input_df = pd.DataFrame([data.model_dump(by_alias=True)])

    processed_data = preprocessor.transform(input_df)

    prediction = int(model.predict(processed_data)[0])

    failure_probability = round(
        float(model.predict_proba(processed_data)[0][1] * 100),
        4
    )

    status = (
        "Maintenance required"
        if prediction == 1
        else "No maintenance required"
    )

    return {
        "prediction": prediction,
        "failure_probability": failure_probability,
        "status": status
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)