ROFARAMADAN/SehaTrack-Pro
0
1"""2model_loader.py — SehaTrack Pro3Centralised lazy-loading with memory-safe caching.4 5WHY THIS FILE EXISTS:6- Streamlit re-runs the entire script on every user interaction.7- Without @st.cache_resource, all 3 models would reload on EVERY click.8- This file wraps every model in a cached loader so they load ONCE and9 stay in memory across all user sessions.10- Each loader returns None if the weight file is missing — the app stays11 alive and shows a friendly "model not available" message instead of crashing.12"""13 14import os15import streamlit as st16 17_HERE = os.path.dirname(os.path.abspath(__file__))18 19 20# ══════════════════════════════════════════════════════════════════════════════21# WHISPER (speech-to-text)22# ══════════════════════════════════════════════════════════════════════════════23@st.cache_resource(show_spinner="Loading Whisper speech model…")24def get_whisper():25 """26 Loads the Whisper 'small' model (~244 MB download on first run).27 Downloaded automatically to ~/.cache/whisper/ — no manual step needed.28 Returns None if whisper is not installed.29 """30 try:31 import whisper32 return whisper.load_model("small")33 except Exception as e:34 st.warning(f"⚠️ Whisper not available: {e}")35 return None36 37 38# ══════════════════════════════════════════════════════════════════════════════39# NLP SYMPTOM CLASSIFIER (HuggingFace)40# ══════════════════════════════════════════════════════════════════════════════41@st.cache_resource(show_spinner="Loading NLP symptom model…")42def get_nlp():43 """44 Loads the HuggingFace NLP classifier from the local 'model_only/' folder.45 Returns (tokenizer, model) or (None, None) if the folder is missing.46 """47 from model import NLP_MODEL_PATH48 import os, json49 from transformers import AutoTokenizer, AutoModelForSequenceClassification50 51 if not os.path.isdir(NLP_MODEL_PATH):52 st.warning(f"⚠️ NLP model folder not found at: {NLP_MODEL_PATH}")53 return None, None54 try:55 tok = AutoTokenizer.from_pretrained(NLP_MODEL_PATH)56 mdl = AutoModelForSequenceClassification.from_pretrained(NLP_MODEL_PATH)57 mdl.eval()58 return tok, mdl59 except Exception as e:60 st.warning(f"⚠️ NLP model failed to load: {e}")61 return None, None62 63 64# ══════════════════════════════════════════════════════════════════════════════65# CHEXNET (chest X-ray — PyTorch DenseNet-121)66# ══════════════════════════════════════════════════════════════════════════════67@st.cache_resource(show_spinner="Loading CheXNet X-ray model…")68def get_chexnet():69 """70 Loads best_chexnet_multimodal.pth.71 Returns model or None if the .pth file is missing.72 """73 from model import load_vision_engine, VISION_WEIGHTS74 if not os.path.isfile(VISION_WEIGHTS):75 st.warning(f"⚠️ CheXNet weights not found at: {VISION_WEIGHTS}")76 return None77 try:78 return load_vision_engine()79 except Exception as e:80 st.warning(f"⚠️ CheXNet failed to load: {e}")81 return None82 83 84# ══════════════════════════════════════════════════════════════════════════════85# KVASIR GI MODEL (EfficientNetB1 — TensorFlow/Keras)86# ══════════════════════════════════════════════════════════════════════════════87@st.cache_resource(show_spinner="Loading GI endoscopy model…")88def get_kvasir():89 """90 Loads gi_model_clean.h5 via TensorFlow/Keras.91 Returns model or None if TF is not installed or .h5 is missing.92 """93 from model import load_kvasir_engine, KVASIR_MODEL_PATH94 if not os.path.isfile(KVASIR_MODEL_PATH):95 st.warning(f"⚠️ Kvasir model not found at: {KVASIR_MODEL_PATH}")96 return None97 try:98 return load_kvasir_engine()99 except Exception as e:100 st.warning(f"⚠️ Kvasir model failed to load: {e}")101 return None