Make the paper's inference and training settings the defaults

CHANGELOG.md

@@ -5,6 +5,17 @@ All notable changes to SubCellSeg are listed here. The format follows
 
 ## Unreleased
 
+### Changed
+- `predict` defaults now match the paper: 30% window overlap and uniform
+  (`constant`) blending, instead of 50% and Gaussian. Pass
+  `--overlap 0.5 --blend gaussian` for the previous behaviour. A voxel is
+  foreground when its probability is at least the threshold (was: above), so
+  `_seg.tif` is exactly `_prob.tif >= 128`.
+- `train` and `finetune` defaults now match the paper: batch size 32 (was 8),
+  at most 1000 epochs (was 700), 20% of the images held out for validation
+  (was 10%). Fine-tuning uses a learning rate of 5e-5 and stops after 50
+  epochs without improvement; training from scratch keeps 1e-4 and 100.
+
 ### Added
 - Benchmarks section in the README and `tools/benchmark/` scripts.
 - Contributing guide, code of conduct, security policy, citation file,

README.md

@@ -144,18 +144,21 @@ from `https://animatedcell-test-data.s3.us-west-2.amazonaws.com/`.
 
 Wall time is the whole command: starting Python and PyTorch, loading the model,
 reading the image over the network (1 to 5 s here) and writing the outputs.
-The sliding-window inference itself takes about 1.3 s for image 40 on the A100
-(63 windows at roughly 48 windows/s). The default `--batch-size 4` and
-`--patch-size 32 128 128` were used.
+The sliding-window inference itself took about 1.3 s for image 40 on the A100.
+Timings were measured with the defaults of the time, 50% window overlap and
+Gaussian blending (252 windows for image 40, in batches of 4). The current
+defaults are the paper's settings, 30% overlap and constant blending, which
+run 140 windows for the same image and are a little faster.
 
 The CPU and GPU masks for image 40 agree on 99.997 % of voxels. Against the
-classic `SEG_STRUCT` segmentation stored in the same files:
+classic `SEG_STRUCT` segmentation stored in the same files, with the current
+defaults (measured 2026-11-10):
 
 | Image | Dice | IoU | Precision | Recall |
 |---|---|---|---|---|
-| 40 | 0.533 | 0.363 | 0.732 | 0.419 |
-| 6919 | 0.877 | 0.781 | 0.902 | 0.853 |
-| 11370 | 0.846 | 0.733 | 0.883 | 0.812 |
+| 40 | 0.532 | 0.362 | 0.733 | 0.417 |
+| 6919 | 0.874 | 0.777 | 0.904 | 0.847 |
+| 11370 | 0.845 | 0.732 | 0.888 | 0.806 |
 
 To reproduce, run `tools/benchmark/bench_predict.sh cuda results 40 6919 11370`
 (it times `subcellseg predict`, records GPU memory, and runs

docs/SEGMENT.md

@@ -115,12 +115,14 @@ Run `subcellseg predict --help` to see them. The ones you are most likely to use
 | `--input`, `-i` | Image, folder of images, or web address | required |
 | `--output`, `-o` | Folder for the results | `subcellseg_output` |
 | `--channel` | Channel with the structure (number from 0, or name) | `0` |
-| `--threshold` | Score above which a voxel counts as structure | `0.5` |
+| `--threshold` | Score at or above which a voxel counts as structure | `0.5` |
 | `--save-prob` | Also save the score image | off |
 | `--checkpoint`, `-c` | Use your own model | pretrained model |
 | `--overwrite` | Redo images that are already done | off |
 | `--device` | `auto`, `cpu` or `cuda` | `auto` |
 | `--batch-size` | Lower it if the GPU runs out of memory | `4` |
+| `--overlap` | How much neighbouring windows overlap | `0.3` |
+| `--blend` | How overlapping windows are combined: `constant` or `gaussian` | `constant` |
 
 ## Problems
 
@@ -145,12 +147,16 @@ images.
 
 ## How this relates to the paper
 
-`subcellseg predict` blends overlapping windows with 50% overlap and Gaussian
-weighting, which gives smoother results at window borders. The benchmarks in
-the paper used 30% overlap and uniform weighting. To reproduce those settings:
+The defaults of `subcellseg predict` are the settings behind the numbers in
+the paper: windows of 32 x 128 x 128 voxels, 30% overlap between neighbouring
+windows, a uniform (`constant`) average where windows overlap, and a threshold
+of 0.5 on the probability map (128 in the saved `_prob.tif`).
+
+Gaussian blending gives smoother transitions at window borders. It is not what
+the paper used, and the masks differ slightly:
 
 ```
-subcellseg predict --input my_images --output results --overlap 0.3 --blend constant
+subcellseg predict --input my_images --output results --overlap 0.5 --blend gaussian
 ```
 
-<!-- RELEASE TODO: confirm these are the settings behind the paper's numbers. -->
+<!-- RELEASE TODO (done 2026-11-10): confirmed against the cyto-dl configs used for the paper: overlap 0.3, MONAI's default constant blending, threshold 0.5. -->

docs/TRAINING.md

@@ -103,7 +103,7 @@ SubCellSeg sets a few of your examples aside for checking its progress (the
 "validation" images) and learns from the rest. You will see lines like:
 
 ```
-Training on 11 image(s), validating on 1.
+Training on 10 image(s), validating on 2.
 Starting from weights in C:\Users\you\.cache\subcellseg\subcellseg_v1.pt
 Training on the GPU.
 Epoch 12: 100%|##########| 2/2 [00:01<00:00, 1.70it/s, val/loss=0.351, val/dice=0.693, train/loss=0.471]
@@ -121,7 +121,7 @@ The numbers jump around from one epoch to the next, especially with few
 images. Look at the trend, not at single values.
 
 Training stops by itself when the validation images have stopped improving
-for 100 epochs, or after 700 epochs. You can also stop it at any time with
+for 50 epochs (100 when training from scratch), or after 1000 epochs. You can also stop it at any time with
 `Ctrl+C`; the best model so far is already saved.
 
 When it finishes it tells you where the model is:
@@ -178,9 +178,10 @@ Run `subcellseg finetune --help` to see all of them.
 | `--train-csv` | The list from Step 3 | required |
 | `--output`, `-o` | Folder for the new model | `subcellseg_run` |
 | `--channel` | Channel with the structure, if images have several (number from 0, or name) | `0` |
-| `--epochs` | Stop after this many epochs at the latest | `700` |
-| `--patience` | Stop when there was no improvement for this many epochs | `100` |
-| `--batch-size` | Lower it if the GPU runs out of memory | `8` |
+| `--epochs` | Stop after this many epochs at the latest | `1000` |
+| `--patience` | Stop when there was no improvement for this many epochs | `50` (fine-tune), `100` (from scratch) |
+| `--batch-size` | Lower it if the GPU runs out of memory | `32` |
+| `--lr` | Learning rate | `5e-5` (fine-tune), `1e-4` (from scratch) |
 | `--val-csv` | Your own list of validation images, instead of an automatic choice | automatic |
 | `--cache-dir` | Folder for prepared images, for speed | off |
 | `--resume` | Continue an interrupted run from `last.ckpt` | off |
@@ -191,18 +192,20 @@ Run `subcellseg finetune --help` to see all of them.
 |---|---|
 | 2 | about 1.5 GB |
 | 4 | about 3 GB |
-| 8 (default) | about 5 GB |
+| 8 | about 5 GB |
 | 16 | about 10 GB |
+| 32 (default) | about 20 GB |
 
-A larger batch size makes each step more reliable but is not required for
-good results. The models in the paper were trained with a batch size of 32.
+The default is the batch size the models in the paper were trained with. A
+smaller batch size also gives good results and is the first thing to lower
+when the GPU runs out of memory.
 
 ## Problems
 
 **Training is extremely slow on Windows.** The batch size is probably too big
 for your graphics card. Windows does not stop with an error when GPU memory
 runs out; it quietly switches to much slower memory instead. Stop with
-`Ctrl+C` and run again with a smaller `--batch-size`, for example 4 or 2.
+`Ctrl+C` and run again with a smaller `--batch-size`, for example 8 or 4.
 
 **"CUDA out of memory".** Use a smaller `--batch-size`.
 
@@ -243,12 +246,13 @@ subcellseg train --train-csv my_data/train.csv --output my_model
 - Every epoch, one random 3D patch of 32 x 128 x 128 voxels (Z, Y, X) is cut
   from each training image, with mild random brightness and contrast changes.
 - The network is a 3D residual U-Net (MONAI `DynUNet`), trained with the
-  Generalized Dice + Focal loss, the Adam optimiser (learning rate 0.0001,
-  weight decay 0.0001) and a learning rate that decays by 0.5% per epoch.
+  Generalized Dice + Focal loss, the Adam optimiser (learning rate 0.00005
+  when fine-tuning, 0.0001 from scratch; weight decay 0.0001) and a learning
+  rate that decays by 0.5% per epoch.
 - Validation uses four fixed patches per validation image, the same ones
   every epoch, so that the numbers are comparable between epochs.
 
-These are the settings used for the models in the paper, except for the fixed
-validation patches and the default batch size.
+These are the settings used for the models in the paper, with one difference:
+the paper's runs cut a fresh random validation patch every epoch.
 
-<!-- RELEASE TODO: confirm the statement above against the paper's methods section. -->
+<!-- RELEASE TODO (done 2026-11-10): checked against the paper's Methods (Hyperparameters, Fine-tuning protocol) and the cyto-dl configs. -->

src/subcellseg/cli.py

@@ -38,7 +38,7 @@ def _add_predict(sub) -> None:
         "--threshold",
         type=float,
         default=0.5,
-        help="Probability above which a voxel is foreground (default: 0.5).",
+        help="Probability at or above which a voxel is foreground (default: 0.5).",
     )
     p.add_argument(
         "--save-prob",
@@ -47,13 +47,17 @@ def _add_predict(sub) -> None:
     )
     p.add_argument("--device", default="auto", help="auto, cpu, cuda, cuda:1, ... (default: auto).")
     p.add_argument(
-        "--overlap", type=float, default=0.5, help="Sliding-window overlap (default: 0.5)."
+        "--overlap",
+        type=float,
+        default=0.3,
+        help="Fraction by which neighbouring windows overlap (default: 0.3, as in the paper).",
     )
     p.add_argument(
         "--blend",
         choices=["gaussian", "constant"],
-        default="gaussian",
-        help="How overlapping windows are combined (default: gaussian).",
+        default="constant",
+        help="How overlapping windows are combined: constant averages them uniformly "
+        "(default, as in the paper), gaussian favours the centre of each window.",
     )
     p.add_argument(
         "--patch-size",
@@ -122,21 +126,21 @@ def _add_train(sub, name: str, finetune: bool) -> None:
     p.add_argument(
         "--val-fraction",
         type=float,
-        default=0.1,
-        help="Fraction held out for validation without --val-csv (default: 0.1).",
+        default=0.2,
+        help="Fraction held out for validation without --val-csv (default: 0.2).",
     )
-    p.add_argument("--epochs", type=int, default=700, help="Maximum epochs (default: 700).")
+    p.add_argument("--epochs", type=int, default=1000, help="Maximum epochs (default: 1000).")
     p.add_argument(
         "--patience",
         type=int,
-        default=100,
-        help="Stop after this many epochs without improvement (default: 100).",
+        help="Stop after this many epochs without improvement "
+        "(default: 50 when fine-tuning, 100 when training from scratch).",
     )
     p.add_argument(
         "--batch-size",
         type=int,
-        default=8,
-        help="Patches per step; lower it if the GPU runs out of memory (default: 8).",
+        default=32,
+        help="Patches per step; lower it if the GPU runs out of memory (default: 32).",
     )
     p.add_argument(
         "--patches-per-image",
@@ -152,7 +156,11 @@ def _add_train(sub, name: str, finetune: bool) -> None:
         metavar=("Z", "Y", "X"),
         help="Training patch size (default: 32 128 128).",
     )
-    p.add_argument("--lr", type=float, default=1e-4, help="Learning rate (default: 1e-4).")
+    p.add_argument(
+        "--lr",
+        type=float,
+        help="Learning rate (default: 5e-5 when fine-tuning, 1e-4 when training from scratch).",
+    )
     p.add_argument(
         "--num-workers", type=int, default=4, help="Data loading processes (default: 4)."
     )

src/subcellseg/data.py

@@ -56,7 +56,7 @@ def read_manifest(csv_path: PathLike, raw_col: str = "raw", seg_col: str = "seg"
 
 
 def split_items(
-    items: list[Item], val_fraction: float = 0.1, seed: int = 12345
+    items: list[Item], val_fraction: float = 0.2, seed: int = 12345
 ) -> tuple[list[Item], list[Item]]:
     """Randomly hold out a validation set. At least one image goes to each side."""
     if len(items) < 2:

src/subcellseg/inference.py

@@ -71,8 +71,8 @@ def predict_proba(
         volume: np.ndarray,
         *,
         patch_size: Sequence[int] = PATCH_SIZE,
-        overlap: float = 0.5,
-        blend: str = "gaussian",
+        overlap: float = 0.3,
+        blend: str = "constant",
         batch_size: int = 4,
         progress: bool = False,
     ) -> np.ndarray:
@@ -88,7 +88,11 @@ def predict_proba(
         overlap
             Fraction by which neighbouring windows overlap.
         blend
-            How overlapping windows are combined: "gaussian" or "constant".
+            How overlapping windows are combined: "constant" averages them
+            uniformly, "gaussian" weights the centre of each window more.
+
+        The defaults (30% overlap, constant blending) are the settings behind
+        the results in the paper.
         batch_size
             Number of windows run through the network at once.
         """
@@ -123,13 +127,14 @@ def predict_proba(
     def segment(self, volume: np.ndarray, threshold: float = 0.5, **kwargs) -> np.ndarray:
         """Binary mask of a (Z, Y, X) volume: uint8, 255 inside the structure.
 
-        Extra keyword arguments are passed to :meth:`predict_proba`.
+        A voxel is inside the structure when its probability is at least
+        ``threshold``. Extra keyword arguments are passed to :meth:`predict_proba`.
         """
         return to_mask(self.predict_proba(volume, **kwargs), threshold)
 
 
 def to_mask(probability: np.ndarray, threshold: float = 0.5) -> np.ndarray:
-    return (probability > threshold).astype(np.uint8) * 255
+    return (probability >= threshold).astype(np.uint8) * 255
 
 
 def predict(
@@ -162,7 +167,8 @@ def predict(
     channel
         Channel to segment, as an index or a channel name.
     threshold
-        Probability above which a voxel is foreground.
+        Probability at or above which a voxel is foreground. With the default
+        of 0.5 the mask equals the saved probability map thresholded at 128.
     save_probability
         Also write the probability map as ``name_prob.tif`` (uint8, 0-255).
     device

src/subcellseg/training.py

@@ -91,13 +91,13 @@ def train(
     channel: int | str = 0,
     raw_col: str = "raw",
     seg_col: str = "seg",
-    val_fraction: float = 0.1,
+    val_fraction: float = 0.2,
     patch_size: Sequence[int] = PATCH_SIZE,
-    batch_size: int = 8,
+    batch_size: int = 32,
     patches_per_image: int = 1,
-    max_epochs: int = 700,
-    patience: int = 100,
-    lr: float = 1e-4,
+    max_epochs: int = 1000,
+    patience: int | None = None,
+    lr: float | None = None,
     weight_decay: float = 1e-4,
     lr_gamma: float = 0.995,
     num_workers: int = 4,
@@ -133,7 +133,13 @@ def train(
         Random patches drawn from each image per epoch.
     max_epochs, patience
         Train for at most ``max_epochs``; stop early once the validation loss
-        has not improved for ``patience`` epochs.
+        has not improved for ``patience`` epochs. ``patience`` defaults to 50
+        when fine-tuning and 100 when training from scratch.
+    lr
+        Learning rate. Defaults to 5e-5 when fine-tuning and 1e-4 when
+        training from scratch.
+
+    The defaults are the settings used for the models in the paper.
     cache_dir
         Directory for an on-disk cache of preprocessed images. Speeds up
         training considerably when images are large.
@@ -154,6 +160,11 @@ def train(
     _quiet_lightning()
     pl.seed_everything(seed, workers=True)
     output_dir = Path(output_dir)
+    finetuning = init is not None
+    if lr is None:
+        lr = 5e-5 if finetuning else 1e-4
+    if patience is None:
+        patience = 50 if finetuning else 100
 
     train_items = read_manifest(train_csv, raw_col, seg_col)
     if val_csv is not None:
@@ -186,6 +197,7 @@ def train(
         print(f"Starting from weights in {source}")
     elif resume is None:
         print("Training from scratch.")
+    print(f"Learning rate {lr:g}; stopping after {patience} epochs without improvement.")
 
     use_gpu = device in ("gpu", "cuda") or (device == "auto" and torch.cuda.is_available())
     if use_gpu:

tests/test_subcellseg.py

@@ -71,12 +71,32 @@ def test_predict_proba(model):
 
     mask = model.segment(volume, patch_size=SMALL_PATCH)
     assert mask.dtype == np.uint8
-    assert np.array_equal(mask > 0, prob > 0.5)
+    assert np.array_equal(mask > 0, prob >= 0.5)
 
     with pytest.raises(ValueError):
         model.predict_proba(volume[0])
 
 
+def test_defaults_match_paper():
+    """The package defaults are the inference and training settings of the paper."""
+    import inspect
+
+    from subcellseg.training import train
+
+    infer = inspect.signature(SubCellSeg.predict_proba).parameters
+    assert infer["patch_size"].default == (32, 128, 128)
+    assert infer["overlap"].default == 0.3
+    assert infer["blend"].default == "constant"
+    assert inspect.signature(SubCellSeg.segment).parameters["threshold"].default == 0.5
+
+    fit = inspect.signature(train).parameters
+    assert fit["batch_size"].default == 32
+    assert fit["max_epochs"].default == 1000
+    assert fit["val_fraction"].default == 0.2
+    # Learning rate and patience depend on whether we fine-tune; resolved inside train().
+    assert fit["lr"].default is None and fit["patience"].default is None
+
+
 def test_predict_ignores_intensity_scale(model):
     volume = _volume().astype(np.float32)
     a = model.predict_proba(volume, patch_size=SMALL_PATCH)