PROJECT_DOCUMENTATION.md

NeuroLens AI Project Documentation

1. Project Overview

NeuroLens AI is a brain MRI analysis project focused on tumor detection, model comparison, explainability, and segmentation research. The repository combines a browser-based dashboard, a Streamlit interface, TensorFlow model training scripts, evaluation utilities, Grad-CAM explainability, and advanced segmentation experiments.

The project is intended for research, learning, and prototype development. It is not a certified clinical diagnostic system and should not be used as the only basis for medical decisions.

2. Objectives

• Detect whether an uploaded brain MRI indicates tumor presence.
• Compare multiple classification model families: CNN, transfer learning, and Vision Transformer.
• Provide interpretable outputs through model metrics and Grad-CAM visualizations.
• Support segmentation experiments using U-Net style architectures.
• Provide reproducible training, evaluation, k-fold validation, ablation, and robustness workflows.

3. Repository Structure

.
|-- app.py                    # Streamlit dashboard entry point
|-- dashboard.py              # Local HTTP server for the HTML dashboard
|-- web_dashboard/            # Frontend dashboard files
|-- src/                      # Core ML, training, evaluation, and segmentation code
|-- config.yaml               # Reference configuration for segmentation experiments
|-- requirements.txt          # Python package requirements
|-- README.md                 # Quick-start project overview
|-- LICENSE                   # MIT license
`-- Documentation/            # Additional report assets

4. Main Components

4.1 HTML Dashboard

The HTML dashboard lives in web_dashboard/ and is served by dashboard.py. It provides a polished interface for image upload, model selection, metrics display, and prediction results.

Key files:

• web_dashboard/index.html
• web_dashboard/style.css
• web_dashboard/app.js
• dashboard.py

The dashboard backend exposes local endpoints for model metadata and image prediction. It looks for model weights and metrics in these folders:

• real_eval_fixed/
• real_eval_current/
• artifacts/

Expected classification weight path:

artifacts/<model_name>/best_weights.weights.h5

where <model_name> is cnn, transfer, or vit.

4.2 Streamlit App

app.py provides an alternate Streamlit dashboard for model comparison and upload-driven prediction. It uses the same broad model families and artifact structure as the HTML dashboard.

4.3 Classification Models

Classification models are defined in src/models.py and trained through src/train.py.

Supported model choices:

• cnn: custom convolutional neural network
• transfer: transfer learning model
• vit: Vision Transformer model

Training output is saved under artifacts/<model_name>/.

4.4 Explainability

Explainability utilities live in src/explain.py and src/utils.py. The dashboard can generate Grad-CAM overlays for supported convolutional models when trained weights are available.

4.5 Segmentation Research

The active segmentation pipeline lives in:

• src/segmentation_torch.py - PyTorch Attention U-Net (GPU)
• src/train_segmentation_torch.py - GPU training script
• generate_pseudo_masks.py - synthesises masks from the classification dataset

PyTorch is used here because TF 2.21 has no native-Windows GPU support; the rest of the project stays on TF/Keras for the classifiers.

The reference TensorFlow framework is preserved for k-fold / ablation / robustness experimentation:

• src/segmentation_models.py - U-Net / Attention U-Net / Res U-Net / Multi-modal U-Net (TF)
• src/train_segmentation.py - TF training script
• src/kfold_validation.py, src/ablation_study.py, src/robustness_analysis.py

Supported segmentation model variants:

• Attention U-Net (default - active PyTorch pipeline)
• U-Net (TF reference)
• Residual U-Net (TF reference)
• Multi-modal U-Net (TF reference)

4.6 Status of "Advanced" Modules (3D MRI, Federated Learning, SSL)

src/advanced_models.py defines MRI3DTransformer, FederatedLearningServer/Client, and SelfSupervisedPretrainer. These are reference implementations:

• 3D ViT: now uses a real Conv3D-strided patch embedding (was a flat Reshape previously). Still requires volumetric data the repo does not ship.
• Federated learning: FedAvg aggregation is correct; evaluation now uses a held-out set when registered via set_eval_data() (was leaking client train data into eval).
• Self-supervised: rotation / jigsaw / SimCLR NT-Xent contrastive pre-training (the contrastive branch previously used an incoherent sigmoid + sparse_categorical_crossentropy combination - replaced with a real NT-Xent loss and two-view augmented batching).

None of these are invoked by train.py or the dashboards; they are libraries waiting for the appropriate datasets and a driver script.

5. Environment Setup

Create a virtual environment:

python -m venv .venv

Activate it on Windows:

.venv\Scripts\activate

Install dependencies:

python -m pip install --upgrade pip
python -m pip install -r requirements.txt

6. Running the Dashboards

6.1 HTML Dashboard

python dashboard.py --port 8501

Open the dashboard at:

http://localhost:8501/

6.2 Streamlit Dashboard

streamlit run app.py

7. Training Workflows

7.1 Classification Training

Example:

python src/train.py --model cnn --dataset dataset --epochs 10 --batch_size 32 --output artifacts

Common options:

• --model: choose cnn, transfer, or vit
• --dataset: dataset directory
• --epochs: number of training epochs
• --batch_size: batch size
• --learning_rate: optimizer learning rate
• --output: artifact output directory

7.2 Segmentation Training (PyTorch, GPU)

# 1. Generate pseudo-masks once (Otsu on tumor JPGs; empty masks for no_tumor):
python generate_pseudo_masks.py --source dataset_real --output dataset_real --clean

# 2. Train Attention U-Net on GPU:
python src/train_segmentation_torch.py --data_dir dataset_real \
    --epochs 25 --batch_size 8 --image_size 256 --base_filters 32 --patience 6

Outputs land in segmentation_artifacts/attention_unet/. Useful flags:

• --device {cuda,cpu}: target device (auto-detect; CUDA used when available)
• --learning_rate: Adam learning rate (default 1e-3)
• --dropout: dropout rate (default 0.2)
• --dice_weight: dice weight in the combined Dice + BCE loss (default 0.6)
• --num_workers: DataLoader workers; default 0 is safest on Windows

7.3 Reference TF Segmentation Training

python src/train_segmentation.py --data_dir dataset_real --model_type attention_unet \
    --epochs 25 --batch_size 8

Useful TF flags:

• --model_type: unet, attention_unet, res_unet, or multi_modal_unet
• --image_size: input image size
• --base_filters: base convolution filter count
• --dropout_rate: dropout rate
• --use_attention: enable attention in supported models
• --use_kfold: run k-fold validation
• --use_ablation: run ablation experiments
• --save_dir: output directory for segmentation artifacts

8. Evaluation

Classification evaluation utilities are available in src/evaluate.py. Segmentation workflows include Dice coefficient, IoU, accuracy, precision, recall, F1 score, specificity, and threshold-based analysis where supported.

Metrics can be stored as JSON files and consumed by the dashboards for comparison views.

9. Dataset Expectations

The project expects local image datasets such as dataset/, dataset_real/, or source data folders. Dataset folders may be large and are usually best managed outside normal source-control workflows unless the project intentionally includes sample data.

For classification, datasets should be organized so TensorFlow data loaders can infer labels from folder structure.

For segmentation, image-mask pairing should follow the assumptions in src/train_segmentation.py and related preprocessing utilities.

10. Artifact Management

Recommended generated-output locations:

• artifacts/ for classification weights, histories, plots, and metrics
• real_eval_current/ and real_eval_fixed/ for selected evaluation snapshots
• segmentation_models/ or a configured output folder for segmentation runs

Large generated files should be reviewed before committing to GitHub.

11. Limitations

• Model quality depends heavily on dataset quality, labeling, preprocessing, and validation strategy.
• MRI datasets can vary across scanners, protocols, institutions, and patient populations.
• Grad-CAM highlights model attention, not guaranteed clinical causality.
• Prototype dashboards do not replace clinical workflows, radiologist review, or regulatory validation.

12. Future Improvements

• Add automated tests for data loading, prediction endpoints, and model utility functions.
• Add a smaller sample dataset or mock artifacts for reproducible demos.
• Add CI checks for formatting and basic import validation.
• Improve artifact versioning and experiment tracking.
• Add clearer dataset preparation documentation for classification and segmentation tasks.

13. License

NeuroLens AI is released under the MIT License. See LICENSE for details.