MiT capstone: forest bark beetle survey detection (PyTorch)

Object-detection pipeline for aerial / oblique forest survey imagery with Pascal-VOC-style XML annotations (tree species + damage class + bounding boxes). The reference dataset layout matches the Spruce Bark Beetle survey tree (data/Data_Set_Spruce_Bark_Beetle/); the same code can be pointed at any root that contains matching Images/ + Annotations/ pairs.

Canonical repository: github.com/markrthomas/MiT_capstone_beetle_kill

An older placeholder repo (PyTorchProject-BeetleKill) was archived after this project moved here.

---

What this repository contains

Piece	Purpose
train_forest_detector.py	Train Faster R-CNN (ResNet50-FPN) (torchvision) on your XML + JPG pairs
infer_forest_detector.py	Run a saved checkpoint on one image or a directory tree; JSON + optional overlays
eval_metrics.py	Per-class precision / recall / F1 vs ground truth (IoU-matched), including subset-aware evaluation
Makefile	setup, train, infer, eval, all, CPU/GPU smoke tests, clean
scripts/run_train.sh, scripts/run_infer.sh	Env-var wrappers used by Make (PYTHON_BIN, sizes, etc.)
scripts/connect_github_origin.sh	Optional helper to merge an existing GitHub repo (legacy)
scripts/strip_large_git_history_and_push.sh	Danger: rewrite history to drop multi-GB blobs; see Git & large files
scripts/tune_spruce_hd_f1.py	Sweep confidence thresholds to maximize spruce_hd F1 (see Tuning for best spruce_hd F1)
docs/	Extra notes (DATA_AND_LABEL_NOTES.md, PIPELINE_QUICKSTART.md); PDFs via docs/Makefile

Legacy scratch files (data_loader.py, try_read_data.py) remain for reference but are not part of the main pipeline.

---

Model and labels

• Detector: torchvision.models.detection.fasterrcnn_resnet50_fpn
• Training labels (per box) come from XML <tree> and <damage> combined according to --label-mode:
  • tree — species only
  • damage — damage token only (e.g. other, hd)
  • tree_damage — default for bark-beetle work: e.g. spruce_hd, pine_other

In the reference dataset, hd is rare vs other; expect class imbalance. Prefer longer training, sensible thresholds, and per-class metrics over loss alone.

---

Dataset layout (expected)

Survey folders should look like:

data/Data_Set_Spruce_Bark_Beetle/
  vertical/
    <site_date>/
      Images/
        <uuid>.jpg
      Annotations/
        <uuid>.xml

Each <object> in XML should include <tree>, <damage>, and <bndbox> with xmin, ymin, xmax, ymax.

Sample discovery uses the XML first:

• Preferred: XML <filename> matched against the sibling Images/ directory
• Fallback: annotation stem plus common image suffixes (.jpg, .jpeg, .png)

Invalid boxes are ignored during training/evaluation, and coordinates are clamped to image bounds before building model targets.

Git: the large data/Data_Set_Spruce_Bark_Beetle/ tree is gitignored (clone or copy data locally; do not commit multi-GB imagery to GitHub).

---

Environment

• Python: 3.10+ recommended (matches requirements.txt pins).
• OS: Linux / WSL2 tested.

Create a virtualenv and install

cd /path/to/MiT_capstone_beetle_kill
python3 -m venv .venv
source .venv/bin/activate   # Windows: .venv\Scripts\activate
pip install --upgrade pip
pip install -r requirements.txt

Or:

make setup

Verify:

python3 -c "import torch, torchvision; print(torch.__version__, torchvision.__version__)"

---

Quick sanity check (smoke test)

End-to-end CPU train (tiny subset) → infer → eval:

make smoke-test

make smoke-test now includes explicit artifact validation (fails if expected output files are missing or if inference produces zero predicted boxes), so regressions are caught automatically.

GPU smoke (fails fast with a message if CUDA is unavailable):

make smoke-test-gpu

Artifacts land under .smoke/ or .smoke-gpu/ (ignored by git).

---

Automated tests

Install isolated test dependencies and run pytest:

make test-install
make test
make check

make test-install attempts to populate .alltest_venv for isolated tests. If you are offline or behind a restricted network, the target now warns instead of hard-failing. make test prefers .alltest_venv when that venv has pytest installed, and otherwise falls back to python3 -m pytest from your current environment. make test runs fast, deterministic unit tests under tests/ without touching your training/inference runtime venv. make check runs both unit tests and the end-to-end smoke test in one command.

---

Training

Using Make (recommended)

make train \
  DATA_ROOT=./data/Data_Set_Spruce_Bark_Beetle \
  OUTPUT_DIR=./outputs \
  LABEL_MODE=tree_damage \
  DEVICE=cuda \
  EPOCHS=20 \
  BATCH_SIZE=2

Calling Python directly

python3 train_forest_detector.py \
  --data-root ./data/Data_Set_Spruce_Bark_Beetle \
  --output-dir ./outputs \
  --label-mode tree_damage \
  --epochs 20 \
  --batch-size 2 \
  --device cuda

Useful training flags

Flag	Default	Notes
--label-mode	(required in script)	tree, damage, or tree_damage
--epochs	10	Increase for real runs
--batch-size	2	Lower if GPU OOM
--lr	5e-4	AdamW
--weight-decay	1e-4
--val-ratio	0.2	Holdout fraction
--num-workers	4	Use 0 if DataLoader workers die on low-RAM CPU
--min-size / --max-size	800 / 1333	Lower both (e.g. 320–640) for CPU smoke or tight RAM
--max-samples	0	>0 = random subset cap for experiments
--seed	42	Reproducible split + subset shuffle
--device	cuda if available	Use cpu when needed

If a tiny subset or a one-image debug run leaves no validation split, training still writes a checkpoint and selects the best epoch by train loss instead of validation loss.

Outputs (under --output-dir):

• best_fasterrcnn.pt — checkpoint with model_state_dict, label_map, args
• label_map.json
• train_history.json

---

Inference

Using Make

make infer \
  CHECKPOINT=./outputs/best_fasterrcnn.pt \
  INPUT_PATH=./data/Data_Set_Spruce_Bark_Beetle/vertical/SomeSite/Images \
  PRED_DIR=./predictions \
  SCORE_THRESHOLD=0.5 \
  DEVICE=cuda

Using Python

python3 infer_forest_detector.py \
  --checkpoint ./outputs/best_fasterrcnn.pt \
  --input ./path/to/image_or_folder \
  --output-dir ./predictions \
  --score-threshold 0.5 \
  --save-annotated-images \
  --device cuda

Outputs:

• predictions.json — list of { image_path, image_key, predictions: [{ label_name, score, bbox_xyxy, ... }] }
• Optional overlay PNG/JPG when --save-annotated-images is set.

image_key is the dataset-relative path when you infer on a directory tree. This prevents collisions when different survey folders contain the same basename.

Use --write-json-min-score (e.g. 0.01) to keep low-score boxes in the JSON while leaving --score-threshold higher for cleaner overlays. That JSON is what you need for threshold tuning (scripts/tune_spruce_hd_f1.py).

Checkpoints are loaded with torch.load(..., weights_only=False) so full training dicts load on PyTorch 2.6+.

---

Evaluation (eval_metrics.py)

Consumes predictions.json and XML ground truth under --data-root.

python3 eval_metrics.py \
  --data-root ./data/Data_Set_Spruce_Bark_Beetle \
  --predictions-json ./predictions/predictions.json \
  --label-mode tree_damage \
  --score-threshold 0.5 \
  --iou-threshold 0.5 \
  --output-json ./predictions/metrics.json

Default behavior: metrics are computed only for images that appear in predictions.json, so subset inference does not count millions of false negatives on unseen images.

Evaluation now keys images by dataset-relative path when available, with basename fallback for older prediction files or single-directory runs. That makes per-image matching stable even when filenames repeat across sites.

To score against all annotated images (strict full-dataset FN accounting):

python3 eval_metrics.py ... --full-dataset-eval

---

Tuning for best spruce_hd F1

Rare-class F1 depends strongly on the confidence cutoff. After training, run inference on a fixed validation image set with a low JSON floor so scores are retained, then sweep thresholds:

# 1) Export low-score boxes into predictions.json (overlay can stay stricter)
python3 infer_forest_detector.py \
  --checkpoint ./outputs/best_fasterrcnn.pt \
  --input ./path/to/val_images \
  --output-dir ./tune_preds \
  --score-threshold 0.5 \
  --write-json-min-score 0.01 \
  --device cuda

# 2) Sweep score thresholds; maximize F1 for spruce_hd (same IoU matching as eval_metrics)
python3 scripts/tune_spruce_hd_f1.py \
  --data-root ./data/Data_Set_Spruce_Bark_Beetle \
  --predictions-json ./tune_preds/predictions.json \
  --label-mode tree_damage \
  --iou-threshold 0.5 \
  --target-class spruce_hd \
  --t-min 0.0 --t-max 0.6 --t-step 0.025 \
  --output-json ./tune_preds/spruce_hd_sweep.json

Use the reported best score_threshold when generating final predictions or in deployment. Re-tune if you change the model, label mode, or IoU policy.

---

Make targets and variables

Target	Meaning
make help	Print targets
make setup	venv + pip install -r requirements.txt + chmod helper scripts
make train	Train via scripts/run_train.sh
make infer	Infer via scripts/run_infer.sh
make eval	Run eval_metrics.py on PRED_DIR/predictions.json
make all	train + infer + eval
make smoke-test / make smoke-test-gpu	Short integration runs
make clean	Remove OUTPUT_DIR, PRED_DIR, .smoke*, etc. (see Makefile)

Common overrides: DEVICE, EPOCHS, BATCH_SIZE, NUM_WORKERS, MIN_SIZE, MAX_SIZE, MAX_SAMPLES, MAX_IMAGES, DATA_ROOT, OUTPUT_DIR, PRED_DIR, LABEL_MODE, SCORE_THRESHOLD, IOU_THRESHOLD.

Example lightweight CPU run:

make all DEVICE=cpu EPOCHS=2 BATCH_SIZE=1 NUM_WORKERS=0 \
  MIN_SIZE=384 MAX_SIZE=640 MAX_SAMPLES=32 MAX_IMAGES=64 \
  SCORE_THRESHOLD=0.1 IOU_THRESHOLD=0.3

---

Documentation PDFs

From docs/:

make -C docs pdf

Requires pandoc (and a PDF engine unless using pandoc’s default). See docs/Makefile.

---

Git & large files

• Never commit data/Data_Set_Spruce_Bark_Beetle/ or other huge archives; .gitignore lists common mistakes (data/archive.zip, MNIST, CIFAR-10 paths, etc.).
• If history ever contains multi-GB blobs, GitHub rejects pushes (pack exceeds maximum allowed size). Use scripts/strip_large_git_history_and_push.sh only when you understand history rewrite and coordinated force-push / new-repo workflows.

---

Troubleshooting

Symptom	What to try
CUDA OOM	Lower --batch-size, lower --min-size / --max-size, fewer workers
DataLoader worker “Killed”	NUM_WORKERS=0
No predictions	Lower --score-threshold; confirm checkpoint label_map matches --label-mode
python3-venv missing (Debian)	sudo apt install python3-venv then recreate .venv
numpy pin vs Python	requirements.txt uses versions compatible with Python 3.10

More context: docs/DATA_AND_LABEL_NOTES.md, docs/PIPELINE_QUICKSTART.md.

---

License

This repository includes an Apache License 2.0 (LICENSE). By contributing or using the code, follow that license and keep third-party dataset terms separate (survey data is not bundled in git).

---

Course / attribution

Started in the context of an MIT deep learning capstone; this README describes the technical pipeline as implemented in this repo, independent of course grading mechanics.