Mouse Tumor Segmentation with Detectron2

This repository contains the full pipeline for automated instance segmentation of subcutaneous mouse tumors using Detectron2 (Mask R-CNN). Data is sourced from Biopticon's TI-2 structured-light scanner. Three model variants are trained and evaluated: RGB, Depth, and RGD (Red-Green-Depth).

The RGD model replaces the blue channel of the texture image with a normalized depth map derived from the scanner's point cloud, producing the strongest overall segmentation results.

---

Repository Structure

---

Installation & Setup

git clone https://github.com/your-repository.git
cd tumor-segmentation
pip install -e .            # installs deps + makes `Detectron2` and `preprocessing` importable
pip install -e ".[dev]"     # plus pytest / black / ipykernel

The editable install is required: it adds src/pipeline/ and src/util/preprocessing/ to your import path so the training scripts run from anywhere without sys.path hacks.

Detectron2 (and SAM, if needed) are not on PyPI and must be installed from source separately:

pip install git+https://github.com/facebookresearch/detectron2.git
pip install git+https://github.com/facebookresearch/segment-anything.git

After pip install -e ., the training entry point is available as a console script:

tumor-segmentation-train --model-name <name> --model-type <group> --modality rgb

…which is equivalent to python3 src/pipeline/training_scripts/train.py ….

Training was run on the Princeton Della cluster using NVIDIA A100 GPUs. The SLURM scripts under src/pipeline/training_scripts/slurm/sbatch_scripts/ are configured for that environment — adjust --account and --partition as needed.

---

Running the Pipeline

The pipeline has a few prerequisites that must be satisfied before training. In particular, the processed_data/ directory — where cached, split data lives — is produced by the first training run with --split-cache true.

Initial Setup

The dataset is hosted as a HuggingFace dataset repository and is included here as a git submodule. On initial clone of this repository it will not be populated. The guide below sets up the HuggingFace submodule.

Note on data access. The dataset is currently gated behind HuggingFace authentication. If you do not have access, contact the maintainers. A public release with an open license is planned — see the Data section.

1. Install git lfs

git lfs install

2. Authenticate with HuggingFace

Generate a private access token from your HuggingFace account (instructions), install the CLI, and log in:

pip install -U "huggingface_hub[cli]"
huggingface-cli login

3. Initialize the submodule

git submodule update --init --recursive

Launch Training

The recommended entry point is the train.sh wrapper, which runs the pytest data-validation suite and then calls train.py:

cd src/pipeline/training_scripts

./train.sh \
  --name <model-name> \
  --model_type <output-subdir> \
  --iter <num-iterations> \
  --modality rgb \      # one of: rgb | depth | rgd
  --split-cache         # pass on first run to preprocess and cache data

train.sh arguments

All boolean flags are bare switches — pass --split-cache to enable, omit it to disable. (The legacy --rgb true / --depth true / --rgd true triplet has been replaced by a single --modality choice.)

train.py CLI (direct invocation)

train.sh is a thin wrapper that translates its named flags into the keyword arguments expected by train.py. If you want to call the Python entry point directly, the signature is:

python3 train.py --model-name <name> --model-type <group> --modality {rgb,depth,rgd} [options] [-- KEY VALUE ...]

Any positional arguments after -- are passed straight to cfg.merge_from_list, so e.g. -- SOLVER.MAX_ITER 10000 SOLVER.BASE_LR 0.0001 overrides those fields. CLI overrides win over the YAML hierarchy.

Configuration

Configs are layered (later wins):

1. Detectron2 model zoo defaults — Mask R-CNN R-101 FPN 3x
2. configs/base.yaml — project-wide overrides (LR schedule, ROI heads, dataloader)
3. configs/{rgb,depth,rgd}.yaml — per-modality overrides, chain via _BASE_: base.yaml
4. CLI: ./train.sh ... -- SOLVER.BASE_LR 0.001

Every run snapshots its fully merged config to models/<model_type>/<model_name>/run_*/config_used.yaml — so a reviewer reading a checkpoint never has to guess what hyperparameters produced it.

Output layout

models/<model_type>/<model_name>/
  run_YYYYMMDD-HHMM/
    config_used.yaml         # fully merged config (every key)
    model_final.pth          # weights
    metrics.json             # Detectron2 default
    events.out.tfevents.*    # TensorBoard scalars
  latest -> run_YYYYMMDD-HHMM/   # symlink to the most recent run

src/pipeline/slurm_output/<model_type>/<model_name>/run_YYYYMMDD-HHMM/
  loss_plots/  AP_Fig/  IoU_fig/  outputs/  IoU_AP_Final/

Each invocation creates a fresh run_<timestamp>/ directory, so reruns never silently overwrite earlier weights.

Experiment tracking

TensorBoard is wired in by default — total_loss, loss_mask, AP, AP50, AP75, and the IoU buckets are written by Detectron2's event storage:

tensorboard --logdir models/

Weights & Biases is opt-in. Set WANDB_PROJECT before launching and a WandbWriter is appended to the writer list automatically:

pip install wandb && wandb login
WANDB_PROJECT=tumor-seg ./train.sh --name foo --model_type rgb --iter 5000 --modality rgb

---

Reproducing the Paper

Status: results in this README are produced by the commands below. Update the seed, model checkpoints, and runtime numbers if you change the recipe.

The published numbers were produced with:

Compute. All training runs were performed on a single NVIDIA A100 (40 GB) on the Princeton Della cluster.

Seed. Seeds are currently inherited from Detectron2's defaults — exact reproduction of paper numbers requires fixing the seed in cfg before broad release (tracked as TODO).

Expected runtime. TBD — fill in once a clean run has been timed end-to-end.

---

Results

Numbers below are placeholders. Replace with the final published values and link to the paper / preprint once available.

Paper / preprint: TBD.

---

Data

The dataset is hosted on HuggingFace at data/huggingface-repo/ (added as a git submodule). See data/README.md for the per-directory breakdown.

---

Known Issues

• correct_binary_masks() uses the wrong OpenCV color-conversion flag. Location: src/util/preprocessing/PreprocessingFunctions/_masks.py:87. The call passes cv2.COLOR_BAYER_BG2GRAY where cv2.COLOR_BGR2GRAY is intended. This treats the input as a Bayer-pattern raw image rather than a BGR image, so the resulting grayscale (and the binary mask derived from it) is incorrect.
  • Impact on published results: the function is invoked from src/util/preprocessing/TumorDataset/_preprocessing.py (_preprocessing.py:46, :111, :122, :135) and runs on every preprocessing pass. All currently published numbers were produced with this bug present.
  • Fix: swap the flag to cv2.COLOR_BGR2GRAY and re-run --split-cache true to regenerate the cached splits before retraining.