saad1BM's picture
Upload 4 files
2cc48b9 verified
import streamlit as st
import shutil # <-- Nayi library for file copy (Fixes cross-device link error)
import tensorflow as tf
from tensorflow.keras.models import load_model
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download
import os
# --- 1. Configuration (Constants) ---
# Yahan aapka confirmed Repo ID use ho raha hai
HF_REPO_ID = "saad1BM/brain-tumor-detection"
MODEL_FILENAME = "brain_tumor_model.keras"
MODEL_PATH = MODEL_FILENAME
IMAGE_SIZE = (224, 224)
CLASS_NAMES = ['glioma_tumor', 'meningioma_tumor', 'no_tumor', 'pituitary_tumor']
# --- 2. Model Loading (FINAL FIX: Cross-Device Link Resolved) ---
@st.cache_resource
def load_trained_model():
"""Trained model ko Hugging Face se download aur load karta hai."""
# Global variables ko function ke andar access karna
global HF_REPO_ID, MODEL_FILENAME, MODEL_PATH
# Minimum expected model size (approx 100 MB, for size check)
MIN_SIZE = 100 * 1024 * 1024
try:
# Check karein ki model pehle se download hai ya nahi ya file corrupt/incomplete hai
if not os.path.exists(MODEL_PATH) or os.path.getsize(MODEL_PATH) < MIN_SIZE:
st.info(f"Downloading model from Hugging Face: {HF_REPO_ID}/{MODEL_FILENAME}")
# hf_hub_download model ko /tmp mein download karega
downloaded_path = hf_hub_download(
repo_id=HF_REPO_ID,
filename=MODEL_FILENAME,
# Cache ko /tmp mein rakhein
cache_dir="/tmp/hf_cache",
)
st.info("Copying model file locally to avoid cross-device link error...")
# os.rename ki jagah shutil.copyfile use karein (Yeh cross-device error ko theek karta hai)
shutil.copyfile(downloaded_path, MODEL_PATH)
st.success("Model download and copy successful!")
# Load the model
model = load_model(MODEL_PATH)
return model
except Exception as e:
# Error message mein zaroori details daal di gayi hain
st.error(f"Model Load Failed. Ensure 'huggingface-hub' is in requirements.txt and file is public. Error: {e}")
return None
model = load_trained_model()
# --- 3. Prediction Function ---
def predict_image(image_file, model):
"""Uploaded image par prediction karta hai."""
if model is None:
return "Model Load Failed", 0.0
# PIL image ko array mein convert karein
img = Image.open(image_file).convert("RGB")
# Image ko model ke input size mein resize karein (224x224)
img = img.resize(IMAGE_SIZE)
img_array = np.array(img)
# Batch dimension add karein: (224, 224, 3) se (1, 224, 224, 3)
img_array = np.expand_dims(img_array, axis=0)
# Normalization (jaisa training mein kiya tha: 1./255)
img_array = img_array / 255.0
# Prediction karein
predictions = model.predict(img_array)
# Highest probability index aur score nikalien
predicted_index = np.argmax(predictions, axis=1)[0]
confidence_score = np.max(predictions) * 100
predicted_class = CLASS_NAMES[predicted_index]
return predicted_class, confidence_score
# --- 4. Streamlit UI ---
st.set_page_config(page_title="Brain Tumor Detection", layout="wide")
# a) Page Title
st.title("🧠 Brain Tumor Detection System (AI Powered)")
st.write("Upload an MRI image below to classify it as one of the tumor types or no tumor.")
st.markdown("---")
col1, col2 = st.columns(2)
with col1:
# b) Image Upload Section
uploaded_file = st.file_uploader("Upload MRI Image:", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# Image Preview
st.image(uploaded_file, caption="Uploaded MRI Image", use_column_width=True)
st.markdown("---")
# c) Prediction Button
if st.button("Detect Tumor"):
st.spinner("Analyzing image and detecting tumor...")
# Prediction
predicted_class, confidence_score = predict_image(uploaded_file, model)
# d) Output Section
# Simplified classification for the app display (Tumor / No Tumor)
if predicted_class == 'no_tumor':
result_label = f"🟢 **Prediction: No Tumor**"
elif predicted_class == 'Model Load Failed':
result_label = f"❌ **Prediction: Model Initialization Error**"
st.error("Model could not be loaded for prediction. Please check logs for Hugging Face download errors.")
else:
result_label = f"🔴 **Prediction: Tumor ({predicted_class.replace('_', ' ').title()})**"
st.success("✅ Analysis Complete")
st.subheader(result_label)
st.metric(label="Confidence Score", value=f"{confidence_score:.2f}%")
st.write("---")
with col2:
st.header("Results and Interpretation")
st.info("The system uses Transfer Learning (VGG16) to classify the image into four categories: Glioma, Meningioma, Pituitary, or No Tumor.")
if uploaded_file is None:
st.warning("Please upload an image and click 'Detect Tumor' to see the results.")