TUBESMLPred / main.py
Imanchukies's picture
Update main.py
e4b556c verified
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import joblib
import numpy as np
import os
# 1. Inisialisasi FastAPI
app = FastAPI(title="Machine Failure API - Hugging Face")
# 2. Aktifkan CORS untuk Lovable
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 3. Load Model & Scaler
MODEL_PATH = "RandomForest_model.joblib"
SCALER_PATH = "scaler.joblib"
# Load di tingkat modul agar efisien
try:
model = joblib.load(MODEL_PATH)
scaler = joblib.load(SCALER_PATH)
print("Successfully loaded model and scaler")
except Exception as e:
model = None
scaler = None
print(f"CRITICAL ERROR: {str(e)}")
# 4. Definisi Schema Input (8 Fitur)
class SensorInput(BaseModel):
air_temperature: float
process_temperature: float
rotational_speed: float
torque: float
tool_wear: float
type_l: int
type_m: int
type_h: int
# Endpoint Root (Agar tidak 404 saat dibuka pertama kali)
@app.get("/")
def health_check():
return {
"status": "online",
"message": "API 8-Fitur siap digunakan",
"model_loaded": model is not None
}
@app.post("/predict")
def predict(data: SensorInput):
if model is None or scaler is None:
return {"error": "Model or Scaler not loaded on server"}
try:
# 2. Update Array (URUTAN HARUS SAMA DENGAN SAAT TRAINING)
# Pastikan urutan ini: Air, Process, Speed, Torque, Wear, L, M, H
features = np.array([[
data.air_temperature,
data.process_temperature,
data.rotational_speed,
data.torque,
data.tool_wear,
data.type_l,
data.type_m,
data.type_h
]])
# Transform & Predict
features_scaled = scaler.transform(features)
prediction = int(model.predict(features_scaled)[0])
# Ambil Probabilitas
try:
prob = model.predict_proba(features_scaled).tolist()[0]
confidence = prob[prediction]
except:
confidence = 1.0
return {
"prediction": prediction,
"status": "Failure Detected" if prediction == 1 else "Normal",
"confidence": round(confidence, 4),
"timestamp": os.popen('date').read().strip()
}
except Exception as e:
return {"error": f"Prediction failed: {str(e)}"}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)