Colonomind Training Architecture

Welcome to the Colonomind Training Repository. This repository has been recently refactored to streamline the training and evaluation of the Mod-SE(2) CNN Hybrid model across various colonoscopy datasets (e.g., Dataset 1, Dataset 2, Public Dataset, TMC).

---

📁 Repository Structure

.
├── Legacy_Notebooks/      # Archived Jupyter Notebooks from previous experiments.
├── paper_evaluations/     # Scripts for generating paper-ready results
├── src/                   # Core modular Python code
│   ├── config.py          # Global parameters, server paths, dataset registry
│   ├── features.py        # Handcrafted feature extraction (Wavelet, GLCM)
│   ├── data_loader.py     # Image loading, preprocessing, and feature binding
│   ├── model.py           # Mod-SE(2) CNN layers and Hybrid Model architecture
│   └── train.py           # Master script for executing training runs
├── antigravity.md         # AI assistant instruction log
├── requirements.txt       # Python package dependencies
└── README.md              # This file

---

🚀 Getting Started

1. Installation

To ensure all collaborators are using the exact same environment, please install dependencies from the requirements.txt file.

pip install -r requirements.txt

2. Server Environment (NTU DGX)

On the NTU DGX server, the following paths are pre-configured:

Item	Path
Project Root	~/Clara/colono_train/
Dataset Root	~/Clara/new_drive/Dataset_Extracted/Dataset+Code/

Available Datasets (registered in src/config.py):

Short Name	Folder
dataset_1	MES classification_20250313
dataset_2	MES classification_20250724
public	MES_Colonoscopy Public Dataset
mixed	MES Mixed Data

Unified Five-Backbone Adjudication Figure

Use the saved Unified models on the labeled mixed dataset in new_drive to create real ensemble examples with CNN Grad-CAM overlays. ViT uses input-gradient saliency because its saved TF-Hub model exposes pooled features:

python src/generate_unified_adjudication_figure.py \
  --base-dir /home/D13K48009/raid/Clara/new_drive \
  --models-dir /home/D13K48009/raid/Clara/colonomind-train/Result/Intra_Unified \
  --output /home/D13K48009/raid/Clara/colonomind-train/Result/Unified_Adjudication_Figure.png

The script also writes a JSON metadata file beside the figure. It reads images from Dataset+Code/MES Mixed Data/MES0 through MES3, uses mean-probability tie-breaking for a 2:2 vote, and uses a most-severe fallback when there is no majority. The four CNN backbones use Grad-CAM to localize image regions that contribute to the selected MES prediction; ViT uses input-gradient saliency. For display, the figure reuses the repository's legacy crop [30:430, 200:550] for raw images larger than 450×550 pixels. This crop shifts the view right to remove the left-side patient overlay. Heatmaps are cropped to the same bounds; model inference still uses the original full frame. An optional extra crop is controlled by --figure-zoom.

To create an editorial figure from a fixed, manually reviewed set of examples, pass a server-local text file with one path per line, relative to Dataset+Code/. Keep patient-specific filenames out of Git. These curated examples are for illustration and must not be presented as an independent test set or used to estimate model performance. These are illustrative model outputs, not clinical advice. All five model weights plus their base_scaler.pkl and umap_model.pkl files must exist under the Unified experiment folders.

3. How to Train a Model

You no longer need to copy/paste Jupyter Notebooks for new experiments! Use the unified train.py script. It automatically handles SMOTE balancing, UMAP projections, CNN training, and artifact saving.

Option A — Using short dataset names (recommended on DGX):

python src/train.py \
  --train_dir dataset_1 \
  --test_dir dataset_2 \
  --output_dir ./results/train1_test2

Option B — Using full custom paths:

python src/train.py \
  --train_dir ~/Clara/new_drive/Dataset_Extracted/Dataset+Code/"MES classification_20250313" \
  --test_dir ~/Clara/new_drive/Dataset_Extracted/Dataset+Code/"MES classification_20250724" \
  --output_dir ./results/train1_test2 \
  --batch_size 16 \
  --epochs 20

Command Arguments:

• --train_dir: A registered short name (dataset_1, dataset_2, public, mixed) or an absolute/relative path.
• --test_dir: Same as above.
• --output_dir: Path where the trained models and artifacts will be saved.
• --batch_size: (Optional) Default is 16.
• --epochs: (Optional) Default is 20.

4. Train on ALL Datasets Combined (Recommended for Paper)

This merges all 4 datasets into one pool, performs a stratified train/val/test split, and trains the model. It also prints the Extended Data Tables 2a & 2b automatically.

python src/train_all.py --output_dir ./results/all_datasets

You can customize the split ratio:

python src/train_all.py --output_dir ./results/all_datasets --split 70 15 15 --epochs 50

This saves an additional split_info.pkl file containing exact image counts per class and per source for your paper tables.

5. Batch Training (All Dataset Combinations)

To train across every dataset combination at once (intra-domain, cross-domain, and multi-domain), use the batch script:

chmod +x run_all_experiments.sh
./run_all_experiments.sh

This will automatically run all experiments and organize results into:

results/
├── intra/                          # Same dataset train & test
│   ├── dataset_1/
│   ├── dataset_2/
│   ├── public/
│   └── mixed/
├── cross/                          # Different dataset train & test
│   ├── train_dataset_1_test_dataset_2/
│   ├── train_dataset_1_test_public/
│   ├── train_dataset_2_test_dataset_1/
│   └── ... (12 combinations total)
└── multi/                          # Mixed-domain training
    ├── train_mixed_test_dataset_1/
    ├── train_mixed_test_dataset_2/
    └── train_mixed_test_public/

Each folder contains: best_hybrid_model.h5, scaler.pkl, label_encoder.pkl, umap_model.pkl.

6. Super Agent Hybrid Routing (New!)

The training scripts have been completely refactored to eliminate test-data leakage. The system now implements a clean One-Shot Training method for the Super Agent (LightGBM) using purely the training dataset.

Hybrid Routing Logic: During evaluation, the system utilizes a Confidence Threshold (default 0.70):

• If Deep Learning Confidence ≥ 0.70 ➔ Prediction is handled purely by the Deep Learning model.
• If Deep Learning Confidence < 0.70 ➔ The case is delegated to the Super Agent (LightGBM) which utilizes UMAP and Handcrafted features for the final prediction.

You can override the default threshold using the --threshold argument.

---

📊 Paper Evaluations Workflow

Follow this order to generate all paper-ready figures and tables:

Step 1: Train Models

# Train a single experiment
python src/train.py --train_dir dataset_1 --test_dir dataset_2 --output_dir ./results/cross/train_dataset_1_test_dataset_2

# Or train ALL experiments at once
./run_all_experiments.sh

Step 2: Run Evaluations (no training needed)

# Data exclusion flowchart & distribution tables
python paper_evaluations/data_stats.py

Step 3: Run Evaluations (needs trained model)

# Confusion Matrix, 95% CI, Sensitivity/Specificity, Feature Importance
python paper_evaluations/run_analysis.py --model_dir ./results/cross/train_dataset_1_test_dataset_2

# Computational performance (params, latency, model size)
python paper_evaluations/run_benchmark.py --model_path ./results/cross/train_dataset_1_test_dataset_2/best_hybrid_model.h5

Step 4: Baselines & Ablation (trains its own models)

# Compare against ResNet, DenseNet, EfficientNet, ViT
python paper_evaluations/run_baselines.py --model resnet

# Run ablation scenarios (1-6)
python paper_evaluations/run_ablation.py --scenario 1 --domain cross

Step 5: Architecture Diagram

Open paper_evaluations/architecture.md in any Mermaid-compatible viewer (GitHub, VS Code, etc.).

Available Scripts:

Script	Needs Training First?	Description
data_stats.py	❌ No	Exclusion flowchart & data split tables
run_baselines.py	❌ No (trains its own)	ResNet, DenseNet, EfficientNet, ViT comparison
run_ablation.py	❌ No (trains its own)	6 ablation scenarios
run_analysis.py	✅ Yes	CM, 95% CI, Sensitivity/Specificity, Feature Importance
run_benchmark.py	✅ Yes	Model size, params, inference latency
architecture.md	❌ No	End-to-end system diagram

Run all scripts from the project root directory (~/Clara/colono_train/).

---

🤝 Team Collaboration Guidelines

To keep this repository clean and lightweight for all team members, please adhere to the following rules:

1. Do not push Datasets: Never commit .jpg, .png, or .csv dataset files to Git. Keep them local.
2. Do not push Large Model Weights: Files like .h5 or .pkl generated in your results/ folder should remain local unless requested. If you must share them, use Google Drive, AWS S3, or Git LFS.
3. Avoid Jupyter Notebook Outputs: If you must create a .ipynb file for testing, clear the outputs (Cell > All Output > Clear) before committing. Notebook outputs inflate the repository size and break git history.
4. Modify src/, don't duplicate: If you want to change the CNN architecture or add a new handcrafted feature, edit src/model.py or src/features.py. Do not create a duplicate src_v2/ folder. Keep the code centralized.

Happy Coding! 🚀