NodDet_3d

A modular, config-driven 3D Faster R-CNN for object detection in volumetric data (built for CT nodule detection, but domain-agnostic). The design follows Detectron2's philosophy: every component — backbone, anchor generator, RPN head, ROI head, box head, augmentations — is registered by name and swapped from a config. Nothing is hard-wired.

volume (C,Z,Y,X) ─► backbone+FPN ─► RPN ─► proposals ─► ROI head ─► boxes
                                     │                     │
                              anchor gen + matcher    pooler + box head
                                     │                     │
                            (optional) HNM / focal   (optional) HNM / cascade

Highlights

• Registry-based, plug-and-play components (backbones, RPN/ROI/box heads, anchor generators, augmentations). Add a component, reference it by name.
• 3D everything — anchors, ROIAlign, NMS, IoU, box regression all operate on (Z, Y, X) volumes with 6-DoF boxes.
• Hard-Negative Mining for both the RPN and the ROI head (shared sampler).
• Focal / RetinaNet-style or cross-entropy RPN training, config-selectable.
• Explicit FPN level assignment (size-band thresholds) or the canonical log2 rule.
• Single- or multi-stage (cascade) ROI heads.
• Modular inference: patch-level and full-scan sliding-window paths are separate, reusable methods.
• Config-toggleable augmentations — turn any on/off from the config.
• Optional auxiliary segmentation head (dice on box-derived masks).

Conventions

Thing	Order
Tensors	(C, Z, Y, X)
Boxes	[z_min, y_min, x_min, z_max, y_max, x_max]
Spacing	(spacing_z, spacing_y, spacing_x)

Install

pip install -e .
# or: pip install -r requirements.txt

Python ≥ 3.10, PyTorch ≥ 2.1. NIfTI volume loading needs nibabel (pip install -e ".[nifti]").

Repository layout

NodDet_3d/
  registry.py          # the name -> component registries
  structures/          # Boxes3D, Instances3D, ImagesList3D, IoU
  backbone/            # segresnet / vit / mednext encoders + FPN builders
  rpn/                 # anchors, matcher (IoU + ATSS), sampling (+HNM), RPN3D
  roi/                 # ROIHeads3D, CascadeROIHeads3D, box head, box regression
  layers/              # ROIPooler3D (level assignment), ROIAlign3D, ShapeSpec
  loss/                # dynamic focal loss
  metrics/             # BoxMetrics3D (COCO-style 3D AP/AR)
  data/                # Detection3DDataset, augmentations, datamodule
  nn/                  # GeneralizedRCNN3D + build_detector
  lightning/           # FasterRCNN3DLightning (train/val/test loops)
  utils/               # SlidingWindowInferer3D_Fast
configs/               # Hydra config tree (one sample recipe)
tools/                 # train.py / infer.py entrypoints

Data

The dataset is a generic CSV reader (NodDet_3d/data/base_dataset.py). One row per ground-truth box; rows sharing a series_uid describe one scan:

column	required	meaning
series_uid	✓	unique scan id
split	✓	train / val / test
volume_path	✓	path to the volume (joined onto volume_root if relative)
z_min,y_min,x_min,z_max,y_max,x_max	✓	box in voxel coords (blank row = no boxes)
label		class id (default 0)
spacing_z,spacing_y,spacing_x		voxel spacing (mm), needed for isotropic resampling
mask_path		binary nodule mask (only if load_masks: true)

Volume formats are dispatched on extension: .npy / .npz, .pt / .pth, .nii / .nii.gz. Volumes are read as (Z, Y, X) in HU.

Training samples random patches (positive crops around a box, or negatives with probability neg_patch_prob); validation/test run full scans through the sliding-window inferer.

Train

python tools/train.py \
  data.csv_path=/path/to/annotations.csv \
  data.volume_root=/path/to/volumes \
  trainer.devices=[0,1,2,3]

Everything is overridable from the CLI (Hydra). Common overrides:

# warm-start the RPN from a checkpoint, train a fresh box head
python tools/train.py pretrained_ckpt=/path/rpn.ckpt pretrained_modules=[rpn] \
  optim.learning_rate_backbone=5e-6 optim.learning_rate_rpn=5e-6 optim.learning_rate_roi=1e-4

# RPN-only detector (no box head)
python tools/train.py roi_training=false

# turn hard-negative mining off on the ROI head
python tools/train.py model.roi.hard_negative_mining=false

# freeze the RPN entirely while training the head
python tools/train.py modules_to_freeze=[rpn]

Infer

python tools/infer.py \
  ckpt_path=/path/to/model.ckpt \
  data.csv_path=/path/to/annotations.csv data.volume_root=/path/to/volumes \
  test_pred_save_dir=./runs/predictions

Runs full-scan sliding-window inference over the CSV test split and writes one .pth per scan (pred_boxes, pred_scores, gt_boxes, plus metadata).

The core concepts (and where they live)

Hard-Negative Mining (HNM)

Instead of sampling background anchors/proposals uniformly, fill part of the negative budget with the hardest examples (highest predicted objectness). hard_negative_fraction=0.5 = half hardest / half random.

• Shared sampler: rpn/sampling.py::subsample_labels_hnm, roi/roi_heads.py::subsample_labels.
• Enabled per-head via model.rpn.hard_negative_mining / model.roi.hard_negative_mining.

Focal loss / RetinaNet-style RPN

model.rpn.box_class_loss_type ∈ {cross_entropy, focal_loss, dynamic_focal_loss}. With retinanet_training: true the RPN trains on all anchors (no subsampling). Implemented in rpn/rpn.py::RPN3D.losses and loss/losses.py.

FPN level assignment

Which pyramid level pools a given proposal. Either explicit size bands (model.roi.pooler.level_thresholds, e.g. [8, 16] voxels) or the canonical log2 rule (level_thresholds: null). See layers/poolers.py::assign_boxes_to_levels_3d.

Extending it

Everything below is "write a class/function, register it, name it in a config."

Add a new backbone

1. Write a build_my_backbone_with_fpn(...) -> nn.Module whose forward(x) returns {"p2": ..., "p3": ..., ...} feature maps (see backbone/fpn.py for the pattern).
2. Register it in backbone/__init__.py:

   BACKBONES.register("my_backbone_fpn")(build_my_backbone_with_fpn)

3. Point the config at it: configs/model/backbone/my_backbone.yaml with name: my_backbone_fpn + kwargs, then set the architecture's defaults: - /model/backbone@backbone: my_backbone.

Add a new box decoder / regression loss

Box encode/decode lives in roi/box_regression.py::Box3DTransform and the loss in _dense_box_regression_loss_3d (smooth_l1 | giou | diou | ciou). Add a new *_loss_3d there and a branch in the loss dispatch, then select it with model.rpn.box_reg_loss_type. For a different parameterization, subclass Box3DTransform and pass it through the builder.

Add a new ROI head

Subclass ROIHeads3D (or CascadeROIHeads3D) in roi/roi_heads.py, then:

ROI_HEADS.register("my_head")(MyROIHeads3D)

Set model.roi.name: my_head. The builder (nn/builder.py) wires the pooler, box head, predictor, and matcher for you.

Add a new augmentation

Write a class with __call__(sample) -> sample (sample has image, boxes, labels) and register it in data/augmentations.py:

@TRANSFORMS.register("my_aug")
class MyAug:
    def __init__(self, prob=0.5): ...
    def __call__(self, sample): ...

Then add - {name: my_aug, prob: 0.5} to data.augmentations in the config.

How it fits together

nn/builder.py::build_detector reads the model config, resolves each component through its registry, and assembles a GeneralizedRCNN3D (nn/FasterRCNN3d.py). The detector exposes four clean entry points: forward_train, forward_inference, forward_sliding_window (full-scan), forward_seg_only. FasterRCNN3DLightning (lightning/lightning_module.py) drives training, patch validation, and full-scan validation, keeping the two inference modes as separate methods (_patch_inference, _fullscan_inference).

License

Apache-2.0.