Chethanand's picture
Upload 24 files
eb3afa1 verified
Raw
History Blame Contribute Delete
2.59 kB
import os
import sys
import torch
from flask import Flask, send_from_directory, jsonify
from flask_cors import CORS
# ── Project root on sys.path so existing modules resolve ──────────────────────
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from config import config
from models.model import ResNetLSTM
# ── Flask setup ───────────────────────────────────────────────────────────────
app = Flask(__name__, static_folder="static", static_url_path="")
CORS(app)
# ── Load model once at startup ────────────────────────────────────────────────
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ResNetLSTM()
_model_path = os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"models", "checkpoints", "last_checkpoint.pth"
)
print(f"[startup] Loading model from {_model_path} …")
_checkpoint = torch.load(_model_path, map_location=device, weights_only=False)
if "model_state_dict" in _checkpoint:
model.load_state_dict(_checkpoint["model_state_dict"])
else:
model.load_state_dict(_checkpoint)
model.to(device)
model.eval()
print(f"[startup] Model ready on {device}")
# ── Attach model / device to app context so blueprints can use them ───────────
app.model = model
app.device = device
# ── Register API blueprint ────────────────────────────────────────────────────
from api.predict import predict_bp
app.register_blueprint(predict_bp, url_prefix="/api")
# ── Serve SPA ─────────────────────────────────────────────────────────────────
@app.route("/", defaults={"path": ""})
@app.route("/<path:path>")
def serve(path):
if path and os.path.exists(os.path.join(app.static_folder, path)):
return send_from_directory(app.static_folder, path)
return send_from_directory(app.static_folder, "index.html")
@app.route("/api/status")
def status():
return jsonify({
"status": "ok",
"device": str(device),
"threshold": config.PREDICTION_THRESHOLD,
})
if __name__ == "__main__":
app.run(debug=False, host="0.0.0.0", port=5000)