from fastapi import FastAPI, HTTPException from fastapi.responses import FileResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from contextlib import asynccontextmanager import torch import torch.nn as nn import joblib import numpy as np import pandas as pd import warnings import os warnings.filterwarnings("ignore", category=UserWarning) class LSTMModel(nn.Module): def __init__(self, input_size, hidden_size, num_layers, output_size): super(LSTMModel, self).__init__() self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=0.2) self.fc1 = nn.Linear(hidden_size, 32) self.relu = nn.ReLU() self.fc2 = nn.Linear(32, output_size) def forward(self, x): out, _ = self.lstm(x) out = self.fc1(out[:, -1, :]) out = self.relu(out) out = self.fc2(out) return out # Global variables for model and data model = None scaler = None test_df_processed = None feature_cols = None SEQ_LENGTH = 30 @asynccontextmanager async def lifespan(app: FastAPI): # Load artifacts on startup global model, scaler, test_df_processed, feature_cols print("===== Starting Application Startup =====") try: print("Loading model and scaler...") # Ensure files exist if not os.path.exists("scaler.pkl") or not os.path.exists("lstm_model.pth"): raise FileNotFoundError("Missing model or scaler artifacts!") scaler = joblib.load("scaler.pkl") model = LSTMModel(input_size=45, hidden_size=64, num_layers=2, output_size=1) # Explicitly load to CPU to avoid issues on CPU-only environments model.load_state_dict(torch.load("lstm_model.pth", map_location=torch.device('cpu'))) model.eval() print("Loading and preprocessing test_FD001.txt...") if not os.path.exists("test_FD001.txt"): raise FileNotFoundError("test_FD001.txt not found!") op_cols = ["op_1", "op_2", "op_3"] sensor_cols = [f"sensor_{i}" for i in range(1, 22)] columns = ["engine_id", "cycle"] + op_cols + sensor_cols test_df = pd.read_csv("test_FD001.txt", sep=r"\s+", header=None) test_df.columns = columns drop_sensors = ["sensor_1", "sensor_5", "sensor_6", "sensor_10", "sensor_16", "sensor_18", "sensor_19"] test_df = test_df.drop(columns=drop_sensors) test_df = test_df.sort_values(["engine_id", "cycle"]) top_sensors = ["sensor_11", "sensor_9", "sensor_4", "sensor_12", "sensor_14", "sensor_7", "sensor_15", "sensor_21", "sensor_2"] for sensor in top_sensors: test_df[f"{sensor}_rollmean"] = test_df.groupby("engine_id")[sensor].rolling(window=5, min_periods=1).mean().reset_index(level=0, drop=True) test_df[f"{sensor}_rollstd"] = test_df.groupby("engine_id")[sensor].rolling(window=5, min_periods=1).std().reset_index(level=0, drop=True) test_df[f"{sensor}_delta"] = test_df.groupby("engine_id")[sensor].diff() test_df = test_df.fillna(0) feature_cols = [c for c in test_df.columns if c not in ["RUL", "engine_id", "max_cycle"]] # Scale test set test_df[feature_cols] = scaler.transform(test_df[feature_cols].values) test_df_processed = test_df print("Successfully loaded artifacts and preprocessed test dataset!") except Exception as e: print(f"CRITICAL STARTUP ERROR: {str(e)}") # We don't raise here so the health check can still pass and we can see logs # But we should probably log the traceback import traceback traceback.print_exc() yield # Clean up on shutdown if needed print("===== Application Shutdown =====") app = FastAPI(title="AeroPulse RUL API", lifespan=lifespan) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Artifacts are now loaded via the lifespan handler @app.get("/engine/{engine_id}/history") def get_engine_history(engine_id: int): if test_df_processed is None: raise HTTPException(status_code=500, detail="Data not loaded properly on server startup.") engine_data = test_df_processed[(test_df_processed["engine_id"] == engine_id)] if engine_data.empty: raise HTTPException(status_code=404, detail="Engine not found in dataset.") features = engine_data[feature_cols].values if len(features) >= SEQ_LENGTH: X_seq = features[-SEQ_LENGTH:] else: pad = np.zeros((SEQ_LENGTH - len(features), len(feature_cols))) X_seq = np.vstack([pad, features]) X_test_t = torch.tensor(X_seq, dtype=torch.float32).unsqueeze(0) # shape (1, 30, 45) with torch.no_grad(): pred = model(X_test_t) # Get last 20 cycles for the charts (unscaled versions preferably, but we scaled them inplace) # Let's get them from unscaled inverse transform or just return the scaled for now since # the frontend draws the shape. For accuracy, let's reverse transform. history = engine_data.tail(20) unscaled_hist = scaler.inverse_transform(history[feature_cols].values) history_df = pd.DataFrame(unscaled_hist, columns=feature_cols) history_df["cycle"] = history["cycle"].values cycles = history_df["cycle"].tolist() sensor_11 = history_df["sensor_11"].tolist() sensor_14 = history_df["sensor_14"].tolist() return { "engine_id": engine_id, "current_cycle": int(cycles[-1]), "predicted_rul": float(pred.item()), "history": { "cycles": cycles, "sensor_11": sensor_11, "sensor_14": sensor_14 } } @app.get("/fleet") def get_fleet_status(): if test_df_processed is None: raise HTTPException(status_code=500, detail="Data not loaded.") X_test_seq = [] engine_ids = [] cycles = [] for engine_id, group in test_df_processed.groupby("engine_id"): engine_ids.append(engine_id) cycles.append(group["cycle"].iloc[-1]) features = group[feature_cols].values if len(features) >= SEQ_LENGTH: X_test_seq.append(features[-SEQ_LENGTH:]) else: pad = np.zeros((SEQ_LENGTH - len(features), len(feature_cols))) X_test_seq.append(np.vstack([pad, features])) X_test_t = torch.tensor(np.array(X_test_seq), dtype=torch.float32) with torch.no_grad(): preds = model(X_test_t).flatten().numpy() fleet = [] for i in range(len(engine_ids)): fleet.append({ "id": int(engine_ids[i]), "cycles": int(cycles[i]), "rul": float(preds[i]) }) # Sort fleet by RUL ascending fleet.sort(key=lambda x: x["rul"]) return fleet @app.get("/health") def health_check(): return { "status": "Healthy" if test_df_processed is not None else "Degraded", "artifacts_loaded": test_df_processed is not None } @app.get("/engine_bg.png") def serve_bg(): return FileResponse("engine_bg.png") @app.get("/") def serve_ui(): return FileResponse("index.html") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)