CoolFace
Apppublic

ROFARAMADAN/SehaTrack-Pro

sourceHugging Faceupdated 3mo agoView on Hugging Face
0likes
model_loader.py101 linesDownload Raw Back to root
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