----------------------------------------------------------------------
This is the API documentation for the cellclass library.
----------------------------------------------------------------------


## Datasets

Loading labelled crops for training.


ROIDataset(rois: 'npt.NDArray[np.uint8]', labels: 'npt.NDArray[np.int_]', transform: 'Callable[[torch.Tensor], torch.Tensor] | None' = None) -> 'None'

Dataset of image regions (ROIs) and integer class labels.

Args:
    rois: Images array of shape (N, C, Y, X).
    labels: Integer class labels of shape (N,).
    transform: Optional transform applied to each ROI tensor.

prepare_dataset(data_dir: 'str', output: 'str', channels: 'list[str] | None' = None, ignore: 'list[str] | None' = None, single_label: 'bool' = True) -> 'None'

Prepare an .npz dataset from a directory of .npy cell image files.

Args:
    data_dir: Directory containing .npy files and optionally metadata.json.
    output: Full path for the output .npz file (e.g. '/path/to/rois.npz').
    channels: Channel names. If None, read from metadata.json in data_dir.
    ignore: Labels to exclude (default: ['unassigned']).
    single_label: Skip masks with more than one label (default: True).


## Models

Model construction and evaluation.


Model(*values)

Supported classification model architectures.

create_model(model: 'Union[Model, str]', num_classes: 'int', num_channels: 'int' = 1, weights: 'str | None' = None, freeze_weights: 'bool' = False, dropout: 'float' = 0.5) -> 'nn.Module | None'

Create a classification model for cell image data.

Args:
    model: Model architecture enum value or string name.
    num_classes: Number of output classes.
    num_channels: Number of input image channels.
    weights: Pretrained weights identifier (e.g. ``'DEFAULT'``).
    freeze_weights: Freeze pretrained weights; only new layers are trained.
    dropout: Dropout probability for the final classification layer.

Returns:
    Instantiated ``nn.Module``, or ``None`` if the model is not recognised.

test_epoch(model: 'nn.Module', validation_loader: 'DataLoader[tuple[torch.Tensor, torch.Tensor]]', device: 'torch.device') -> 'tuple[torch.Tensor, torch.Tensor]'

Run one evaluation pass and return true and predicted labels.

Args:
    model: Network model (must already be on ``device``).
    validation_loader: DataLoader yielding (images, labels) batches.
    device: Device to run inference on.

Returns:
    Tuple of (y_true, y_pred) as 1-D CPU tensors.


----------------------------------------------------------------------
This is the CLI documentation for the package.
----------------------------------------------------------------------

## CLI: cellclass

```
Usage: cellclass [OPTIONS] COMMAND [ARGS]...

  Train and evaluate cell-image classifiers.

Options:
  --version  Show the version and exit.
  --help     Show this message and exit.

Commands:
  batch    Generate training runs from BATCH.
  dataset  Convert cell images in DIR to a training dataset.
  extract  Extract a TorchScript model from training settings FILE.
  sample   Sample example images from DATASET.
  sbatch   Submit a CellClass training job to SLURM.
  test     Evaluate a trained model against a dataset.
  train    Train a model using an ROI dataset or resume from saved state.
```

### cellclass sample

```
Usage: cellclass sample [OPTIONS] DATASET

  Sample example images from DATASET.

Options:
  --output DIRECTORY  Output path (default: same directory as dataset).
  --samples INTEGER   [default: 10]
  --crop INTEGER      Crop NxN dimension (default: full image).
  --help              Show this message and exit.
```

### cellclass dataset

```
Usage: cellclass dataset [OPTIONS] DIR

  Convert cell images in DIR to a training dataset.

Options:
  --name TEXT                     Output dataset name.  [default: rois]
  --out DIRECTORY                 Output dataset directory.
  --ignore TEXT                   Labels to ignore. Accepts one or more
                                  values.
  --duplicates / --no-duplicates  Log duplicates.
  --channels TEXT                 Channel names. Accepts one or more values.
  --single-label / --no-single-label
                                  Only use single-label masks.
  --help                          Show this message and exit.
```

### cellclass extract

```
Usage: cellclass extract [OPTIONS] FILE

  Extract a TorchScript model from training settings FILE.

Options:
  --name TEXT                     Model file prefix (default uses model
                                  metadata).
  --save / --no-save              Save model files.
  --overwrite / --no-overwrite    Overwrite existing model files.
  --model [densenet121|densenet161|densenet169|densenet201|squeezenet1_0|shufflenet2x1_0|shufflenet2x1_5|shufflenet2x2_0|efficientnetb3|efficientnetb4|efficientnetb3s|efficientnetb4s|resenext50_32x4d|alexnet|alexnet_att|googlenet|googlenet_att|simplegooglenet|simplegooglenet_att|simplegooglenet_ch_att|simplegooglenet_ch_sp_att]
                                  Model override.
  --channels TEXT                 Input channel overrides.
  --labels TEXT                   Class label overrides.
  --help                          Show this message and exit.
```

### cellclass batch

```
Usage: cellclass batch [OPTIONS] BATCH

  Generate training runs from BATCH.

Options:
  --dry-run / --no-dry-run        Perform a dry run.
  --background / --no-background  Run each training process in the background.
  --script TEXT                   [default: batch.sh]
  --cmd TEXT                      Training program.
  --help                          Show this message and exit.
```

### cellclass test

```
Usage: cellclass test [OPTIONS] dataset.npz

  Evaluate a trained model against a dataset.

Options:
  --size INTEGER                  Number of images to use.
  --data-seed INTEGER             [default: 42]
  -d, --device TEXT               [default: cuda]
  --model [densenet121|densenet161|densenet169|densenet201|squeezenet1_0|shufflenet2x1_0|shufflenet2x1_5|shufflenet2x2_0|efficientnetb3|efficientnetb4|efficientnetb3s|efficientnetb4s|resenext50_32x4d|alexnet|alexnet_att|googlenet|googlenet_att|simplegooglenet|simplegooglenet_att|simplegooglenet_ch_att|simplegooglenet_ch_sp_att]
  -n, --name FILE                 Training checkpoint name.
  -w, --weights FILE              Model weights.
  -s, --script FILE               Model script.
  --testing-size FLOAT            [default: 0.2]
  --batch-size INTEGER            [default: 128]
  --num-workers INTEGER           [default: 0]
  --log-level INTEGER             [default: 20]
  --pin-memory / --no-pin-memory
  --reorder INTEGER               Reorder class labels.
  --plot-matrix / --no-plot-matrix
  --matrix-file TEXT              Save the confusion matrix image.
  --matrix-csv TEXT               Save the confusion matrix CSV.
  --help                          Show this message and exit.
```

### cellclass sbatch

```
Usage: cellclass sbatch [OPTIONS] [TRAINING_ARGS]...

  Submit a CellClass training job to SLURM.

  Place arguments forwarded to training after ``--args``.

Options:
  --class TEXT          [default: gpu]
  --results-dir TEXT    [default: /mnt/lustre/users/gdsc/]
  -u, --username TEXT   [default: runner]
  --hours INTEGER       [default: 12]
  -m, --memory INTEGER  [default: 32]
  --no-gpu
  --no-exec
  --no-submit
  --help                Show this message and exit.
```

### cellclass train

```
Usage: cellclass train [OPTIONS] dataset.npz | state.json

  Train a model using an ROI dataset or resume from saved state.

Options:
  --size INTEGER                  Number of images to use.
  --data-seed INTEGER             [default: 42]
  -d, --device TEXT               [default: cuda]
  -s, --state TEXT                [default: training.json]
  --model [densenet121|densenet161|densenet169|densenet201|squeezenet1_0|shufflenet2x1_0|shufflenet2x1_5|shufflenet2x2_0|efficientnetb3|efficientnetb4|efficientnetb3s|efficientnetb4s|resenext50_32x4d|alexnet|alexnet_att|googlenet|googlenet_att|simplegooglenet|simplegooglenet_att|simplegooglenet_ch_att|simplegooglenet_ch_sp_att]
  -n, --name TEXT                 [default: model.pt]
  --existing [overwrite|load|error]
  --weights TEXT                  [default: DEFAULT]
  --freeze-weights / --no-freeze-weights
  --epochs INTEGER                [default: 2000]
  --batch-size INTEGER            [default: 32]
  --num-workers INTEGER           [default: 0]
  --seed INTEGER                  [default: 3735928559]
  --validation-size FLOAT         [default: 0.2]
  --lr FLOAT                      [default: 0.0001]
  --lr-scheduler [step|plateau]
  --lr-gamma FLOAT                [default: 0.1]
  --lr-step INTEGER               [default: 5]
  --lr-factor FLOAT               [default: 0.1]
  --lr-patience INTEGER           [default: 2]
  --patience INTEGER              [default: 10]
  --delta FLOAT                   [default: 0.0]
  --rel-delta FLOAT               [default: 0.0001]
  --weight-decay FLOAT            [default: 0.0001]
  --flip INTEGER                  [default: 1]
  --rotate INTEGER                [default: 180]
  --translate FLOAT               [default: 0.1]
  --loss-function [focal_loss|cross_entropy]
  --focal-gamma FLOAT             [default: 2.0]
  --loss-weights / --no-loss-weights
  --dropout FLOAT                 [default: 0.4]
  --testing-size FLOAT            [default: 0.2]
  --testing-interval INTEGER      [default: 5]
  --log-level INTEGER             [default: 20]
  --cudnn-benchmark / --no-cudnn-benchmark
  --pin-memory / --no-pin-memory
  --wandb / --no-wandb
  --entity TEXT                   Weights and Biases team/entity.
  --project TEXT                  [default: cellclass]
  --run-name TEXT                 Weights and Biases run display name.
  --tags TEXT                     Weights and Biases tags.
  --wandb-id TEXT                 Weights and Biases run ID.
  --help                          Show this message and exit.
```


----------------------------------------------------------------------
This is the User Guide documentation for the package.
----------------------------------------------------------------------


## Guide

### CellClass — ML Cell Classifier

**cellclass** is the machine-learning training pipeline for the OmeroScreen ecosystem.
It trains lightweight convolutional neural networks to classify single-cell crops generated by the
[Training Widget](../../napari/user-guide/training_widget.html) and deploys the resulting
[TorchScript](https://pytorch.org/docs/stable/jit.html) models back into the
[omero-screen pipeline](../../) for automated inference.

The package follows a four-step workflow:

1. **Create dataset** — normalise and package `.npy` crop files into a single `.npz` archive.
2. **Write a training batch file** — one run per line with model and hyperparameter flags.
3. **Train** — run all jobs locally or on a SLURM cluster; results log to Weights & Biases.
4. **Extract model** — export the best checkpoint to a TorchScript `.pt` + `.json` sidecar for inference.

::: {.callout-note}
Every command is documented, option by option, in the generated
[CLI reference](../reference/cli/index.html). The pages in this guide explain
*why* and *when* to use them.
:::


### Installation

cellclass is part of the omero-screen monorepo workspace. It is installed automatically
when you run `uv sync` from the repository root.

```bash
git clone https://github.com/Helfrid/omero-screen.git
cd omero-screen
uv sync --dev
source .venv/bin/activate
```

The unified `cellclass` command provides all CellClass workflows; run
`cellclass --help` to see its subcommands, or browse the generated
[CLI reference](../reference/cli/index.html) for the full command tree. Existing
`cellclass-*` entry points remain available as compatibility aliases during the
migration.

## Weights & Biases

Training runs are logged to [Weights & Biases](https://wandb.ai). Set up your credentials once:

```bash
wandb login
```

Visit <https://wandb.ai/authorize>, log in, and paste the API key when prompted.
The key is saved to `~/.netrc` and persists across sessions and on the Artemis cluster.
Verify the setup with:

```bash
wandb status
```

By default, runs log to the **cellclass** project in your own W&B workspace — no
shared credentials needed. To collaborate with a team, pass `--entity <team>` to
`cellclass train` (e.g. `--entity hocheggerlab`), then share the project via the
W&B interface.


### Quick Start

This page walks through the complete workflow from annotated crop files to a deployed
TorchScript model, using a two-class mitosis classifier as a running example.

## Prerequisites

- Annotated `.npy` crop files produced by the
  [Training Widget](../../napari/user-guide/training_widget.html), stored in a single
  directory (e.g. `~/data/mitosis-rpe/`).
- A `metadata.json` file in that directory containing at minimum:

  ```json
  { "user_data": { "channels": ["DAPI", "Tub"] } }
  ```

  The Training Widget writes this file automatically.
- [Weights & Biases](installation.qmd#weights-biases) credentials configured.

---

## Step 1 — Create the dataset

Pack all `.npy` files into a single normalised `.npz` archive:

```bash
cellclass dataset ~/data/mitosis-rpe --name rois
```

This applies (1, 99) percentile normalisation, zeros all pixels outside each cell mask,
converts to `uint8`, and writes `~/data/mitosis-rpe/rois.npz`.

Inspect a random sample of each class as a TIFF montage:

```bash
cellclass sample ~/data/mitosis-rpe/rois.npz --output ~/data/mitosis-rpe
```

Open the resulting TIFFs in ImageJ (`Image → Color → Channels Tool → Composite`) to
verify your labels look correct before training.

---

## Step 2 — Write `train.txt`

Create a plain-text batch file, one training run per line. See the
[Model Guide](models.qmd) for recommended architectures for small cell images.

A typical starting point for 50 × 50 px crops:

```text
~/data/mitosis-rpe/rois.npz --model shufflenet2x1_0 --lr 1e-3 --lr-scheduler plateau --flip 3 --batch-size 128 -d mps --wandb
~/data/mitosis-rpe/rois.npz --model shufflenet2x1_0 --lr 1e-4 --lr-scheduler plateau --flip 3 --batch-size 128 -d mps --wandb
~/data/mitosis-rpe/rois.npz --model densenet121      --lr 1e-3 --lr-scheduler plateau --flip 3 --batch-size 64  -d mps --wandb
~/data/mitosis-rpe/rois.npz --model densenet121      --lr 1e-4 --lr-scheduler plateau --flip 3 --batch-size 64  -d mps --wandb
~/data/mitosis-rpe/rois.npz --model efficientnetb3s  --freeze-weights --lr 1e-3 --lr-scheduler plateau --flip 3 --batch-size 64 -d mps --wandb
~/data/mitosis-rpe/rois.npz --model efficientnetb3s  --freeze-weights --lr 1e-4 --lr-scheduler plateau --flip 3 --batch-size 64 -d mps --wandb
```

Key flags:

- `--flip 3` — random horizontal and vertical flips (cells have no orientation bias).
- `--lr-scheduler plateau` — adapts learning rate to convergence speed.
- `--freeze-weights` — only for `efficientnetb3s` / `efficientnetb4s` (ImageNet pretrained).
- `-d mps` — Apple Silicon GPU. Use `-d cuda` on Linux, `-d cpu` as fallback.
- `--wandb` — log to Weights & Biases.

Lines starting with `#` are ignored; comment out completed runs before regenerating
`batch.sh` to avoid overwriting results.

---

## Step 3 — Generate and run `batch.sh`

```bash
cellclass batch train.txt --script batch.sh
bash batch.sh
```

For SLURM on the Artemis cluster see the [SLURM guide](slurm.qmd).

---

## Step 4 — Review results on Weights & Biases

Open your W&B dashboard at <https://wandb.ai> and navigate to the **cellclass** project
in your workspace. Compare validation accuracy across runs and note the metadata `.json`
filename of the best run.

---

## Step 5 — Extract the model

```bash
cellclass extract best_run.json --save
```

This saves two files alongside the checkpoint:

- `shufflenet2x1_0_c2_l2.pt` — TorchScript model (architecture + weights).
- `shufflenet2x1_0_c2_l2.json` — sidecar with channel names, class labels, and image size.

The filename encodes the architecture, number of input channels (`c2`), and number of
classes (`l2`).

---

## Step 6 — Verify the model

```bash
cellclass test -s shufflenet2x1_0_c2_l2.pt ~/data/mitosis-rpe/rois.npz
```

This prints a full classification report (precision, recall, F1 per class) on the held-out
test split.

---

## Step 7 — Deploy

Copy the `.pt` + `.json` pair to the location expected by `omero-screen` and set
`OMERO_SCREEN_INFERENCE_MODEL` to the `.pt` path. The pipeline will load the model
automatically during plate analysis. See the
[pipeline configuration](../../user-guide/installation.html#configure) for details.


## Data format

Training data is stored as `.npy` files — pickled Python dicts with the following structure:

```python
{
    'data': ([img1, img2, ...], [mask1, mask2, ...]),  # images (YXC), masks (YX)
    'target': ['interphase', 'mitosis', ...]           # string class labels
}
```

The [Training Widget](../../napari/user-guide/training_widget.html) produces files in exactly
this format. Each file contains all annotations from one session (one plate / well / image
combination). Multiple files can coexist in the same directory and are combined by
`cellclass dataset`.

---

## Step 1 — Create dataset

::: {.callout-tip}
If you are working in Napari, this step is automated. Open the **Training Widget**, click
**Manage Sessions**, then **Export Dataset**. This runs `cellclass dataset` on all sessions
for the selected classifier and also writes a ready-to-use `train.txt` template — you can
skip straight to [Step 4](#step-4-generate-and-run-batchsh).
:::

To run manually from the command line:

```bash
cellclass dataset <data_dir> --name rois
```

What this does:

1. Globs all `*.npy` files in `<data_dir>`.
2. Drops any crop labelled `unassigned` (cells you reviewed but did not classify).
3. For each remaining crop: applies (1, 99) percentile stretch, centres the cell,
   and converts to `uint8`.
4. Deduplicates identical crops (by pixel hash) to avoid overfitting on repeated cells.
5. Writes `<data_dir>/rois.npz` — a single compressed numpy archive.

The `.npz` contains all crops in a single pool — train/validation/test splitting and
class imbalance handling are done by `cellclass train` at training time (see
[Step 4](#step-4-generate-and-run-batchsh)).

Channel names are read from `metadata.json` in the data directory. Override with `--channels`:

```bash
cellclass dataset ~/data/ppase_screen --channels DAPI RFP --name rois
```

Useful options:

| Flag | Default | Description |
|---|---|---|
| `--name` | `rois` | Output filename stem |
| `--channels` | from metadata.json | Channel names (order must match image data) |
| `--ignore` | none | Class labels to exclude |
| `--size` | none | Resize all crops to `N×N` pixels |

---

## Step 2 — Sample and verify

Before training, visually verify your dataset:

```bash
cellclass sample ~/data/mitosis-rpe/rois.npz --output ~/data/mitosis-rpe
```

This writes one TIFF per class containing a random grid of crops in ImageJ hyperstack
format. Open in ImageJ with `Image → Color → Channels Tool → Composite` to check
channel assignments and label quality.

---

## Step 3 — Write `train.txt`

A batch file is a plain-text list of arguments for `cellclass train`, one run per line.
Comment lines (`#`) and blank lines are ignored. `cellclass batch` converts this file into
a `batch.sh` shell script, assigning unique checkpoint filenames automatically so runs
never overwrite each other.

---

### Flag reference

::: {.callout-note}
The authoritative, always-current list of options and defaults is the generated
[`cellclass train` reference](../reference/cli/train.html). The notes below add
the reasoning behind each flag — which value to pick, and why.
:::

#### Input

`dataset.npz` *(positional)*
: Path to the `.npz` file produced by `cellclass dataset`. You can also pass a `.json`
  state file to resume a crashed or paused run.

`--size N`
: Subsample `N` images randomly from the dataset. Default is all images. Useful for quick
  sanity checks (`--size 200`) without waiting for a full run.

`--data-seed INT` *(default: 42)*
: Random seed for subsampling. Change this to get a different random subset.

---

#### Model

`--model NAME` *(default: densenet121)*
: Neural network architecture. Available options:

  | Model | Size | Notes |
  |---|---|---|
  | `densenet121` | medium | Strong general-purpose performer |
  | `efficientnetb3s` | medium | ImageNet pretrained; use with `--freeze-weights` |
  | `shufflenet2x1_0` | light | Fast training, good baseline |
  | `squeezenet1_0` | lightest | Minimal memory; useful sanity check |

  See the [Model Guide](models.qmd) for full architecture details.

`--freeze-weights` / `--no-freeze-weights`
: Freeze the pretrained backbone and only train the classification head. Use this with
  `efficientnetb3s`, which ships with ImageNet weights. Freezing means the feature extractor
  is not updated — training is faster and needs less data, but the backbone cannot adapt to
  your specific channel content. For fluorescence images that look nothing like ImageNet
  photos, `--no-freeze-weights` often gives better final accuracy at the cost of slower
  convergence.

`--weights TEXT` *(default: DEFAULT)*
: PyTorch weights identifier. `DEFAULT` loads the standard pretrained weights for the
  chosen model. Pass `None` to start from a random initialisation (rarely needed — prefer
  `--no-freeze-weights` instead).

`--dropout FLOAT` *(default: 0.4)*
: Dropout rate applied to the final feature vector before the classification layer.
  Higher values (0.5–0.6) regularise more aggressively and help when the dataset is small
  or the feature vector is large. Lower values (0.2–0.3) suit larger datasets.

---

#### Learning rate

`--lr FLOAT` *(default: 1e-4)*
: Initial learning rate for the Adam optimiser. This is the single most important
  hyperparameter to sweep. Start with `1e-3` and `1e-4`; if both diverge, try `5e-4`.
  If both converge slowly, `1e-3` is usually better.

`--lr-scheduler TYPE` *(default: step)*
: How the learning rate decays during training.

  - **`plateau`** — reduces LR by `--lr-factor` when validation loss stops improving for
    `--lr-patience` epochs. Adapts to your data automatically. **Recommended for most runs.**
  - **`step`** — multiplies LR by `--lr-gamma` every `--lr-step` epochs regardless of
    performance. Predictable but requires more tuning.

`--lr-factor FLOAT` *(default: 0.1, plateau only)*
: Multiplicative factor applied when the plateau scheduler fires. `0.1` drops the LR by
  10×. Use `0.5` for a gentler reduction if training is unstable.

`--lr-patience INT` *(default: 2, plateau only)*
: Number of epochs with no improvement before the plateau scheduler fires. Increase to
  `5–10` if your loss curve is noisy.

`--lr-gamma FLOAT` *(default: 0.1, step only)*
: Multiplicative factor applied at each step. Same semantics as `--lr-factor` but for the
  step scheduler.

`--lr-step INT` *(default: 5, step only)*
: Epoch interval between step scheduler reductions.

`--weight-decay FLOAT` *(default: 1e-4)*
: L2 regularisation on model weights (Adam's `weight_decay`). Increase to `1e-3` if you
  see overfitting (training accuracy far above validation accuracy).

---

#### Training loop

`--epochs INT` *(default: 2000)*
: Maximum number of training epochs. Early stopping (controlled by `--patience`) usually
  terminates training well before this limit. Set lower (`--epochs 100`) for quick test
  runs on a laptop.

`--batch-size INT` *(default: 32)*
: Number of crops per gradient update. Larger batches train faster but require more GPU
  memory. Recommended values: `64` for medium models (DenseNet, EfficientNet), `128` for
  lightweight models (ShuffleNet, SqueezeNet).

`--patience INT` *(default: 10)*
: Early-stopping patience: training stops when validation accuracy has not improved for
  this many epochs. Increase to `20–30` for noisy datasets where improvement is slow.

`--delta FLOAT` *(default: 0)*
: Minimum *absolute* improvement in validation accuracy to reset the patience counter.
  Leave at `0` unless you want to ignore very small improvements.

`--rel-delta FLOAT` *(default: 1e-4)*
: Minimum *relative* improvement (as a fraction of current accuracy) to count as
  progress. Prevents tiny numerical fluctuations from resetting the patience counter.

`--validation-size FLOAT` *(default: 0.2)*
: Fraction of data held out for validation during training. The remaining `1 - validation-size - testing-size` fraction is used for training.

`--seed INT` *(default: 0xDEADBEEF)*
: Random seed for model initialisation and data splits. Fix this across runs when you want
  to isolate the effect of a single hyperparameter change.

---

#### Loss function

`--loss-function TYPE` *(default: focal_loss)*
: Loss function used during training.

  - **`focal_loss`** — down-weights easy examples so the model focuses on hard or rare
    cases. **Recommended for imbalanced datasets** (e.g. many interphase cells, few mitotic).
  - **`cross_entropy`** — standard classification loss. Use when classes are balanced.

`--focal-gamma FLOAT` *(default: 2)*
: Focusing parameter for focal loss. Higher values (3–5) concentrate training even more on
  hard examples. Effective range is 0.5–5; `2` is the standard from the original paper.

`--loss-weights` / `--no-loss-weights`
: Weight each class in the loss by the inverse of its frequency. Helps with imbalanced
  datasets independently of focal loss — you can combine both. Disabled by default.

---

#### Data augmentation

`--flip INT` *(default: 1)*
: Random flip augmentation applied each epoch.

  | Value | Effect |
  |---|---|
  | `0` | No flipping |
  | `1` | Horizontal flips only |
  | `2` | Vertical flips only |
  | `3` | Both horizontal and vertical |

  Use `3` for cell images — cells have no preferred orientation.

`--rotate INT` *(default: 180)*
: Maximum rotation angle in degrees. Each crop is randomly rotated by ±`rotate` degrees.
  `180` means full random rotation, which is appropriate for round objects like cells.
  Set to `0` to disable.

`--translate FLOAT` *(default: 0.1)*
: Maximum translation as a fraction of image size. `0.1` shifts crops by up to 10% in
  each direction. Helps the model generalise to slightly off-centre cells.

---

#### Hardware

`-d DEVICE` *(default: cuda)*
: Compute device.

  | Value | When to use |
  |---|---|
  | `cuda` | NVIDIA GPU (Linux/Windows server) |
  | `cuda:0`, `cuda:1` | Specific GPU on a multi-GPU machine |
  | `mps` | Apple Silicon (M1/M2/M3 Mac) |
  | `cpu` | Fallback; very slow for anything but tiny tests |

`--num-workers INT` *(default: 0)*
: Number of parallel workers for data loading. On Linux with CUDA, `4` or `8` significantly
  reduces data loading bottlenecks. **Set to `0` on macOS** (MPS does not support forked
  workers).

`--cudnn-benchmark`
: Enables cuDNN autotuning — benchmarks convolution algorithms at startup and picks the
  fastest one for your specific input size. Adds ~30 s startup overhead but speeds up
  training. **CUDA only; omit on MPS/CPU.**

`--pin-memory`
: Uses CUDA pinned (page-locked) memory in the DataLoader for faster host→GPU transfers.
  **CUDA only; omit on MPS/CPU.**

---

#### Testing

`--testing-size FLOAT` *(default: 0.2)*
: Fraction of data held out as a final test set, evaluated once at the end of training.
  This split is never seen during training or validation.

`--testing-interval INT` *(default: 5)*
: Evaluate the held-out test set every N epochs (in addition to the final evaluation).
  Useful for monitoring generalisation during training; set higher to reduce overhead.

---

#### Weights & Biases

`--wandb`
: Enable logging to Weights & Biases. Requires a prior `wandb login`. Runs are logged to
  the `cellclass` project in your own workspace by default.

`--entity TEXT`
: W&B team or username. Defaults to your logged-in account. Pass `--entity myteam` to log
  to a shared team project.

`--project TEXT` *(default: cellclass)*
: W&B project name.

`--run-name TEXT`
: Human-readable display name for this run in the W&B UI. If not set, W&B generates one.

`--tags TAG [TAG ...]`
: One or more tags attached to the run, useful for filtering in W&B. Example:
  `--tags mitosis rpe 50px`.

---

### Standard sweep (GPU server)

The default sweep covers four architectures and two learning rates — six runs total.
This is a good starting point for any new dataset.

```text
rois.npz --model efficientnetb3s --freeze-weights --lr 1e-3 --lr-scheduler plateau --flip 3 --batch-size 64 -d cuda --cudnn-benchmark --pin-memory --num-workers 4 --wandb
rois.npz --model efficientnetb3s --freeze-weights --lr 1e-4 --lr-scheduler plateau --flip 3 --batch-size 64 -d cuda --cudnn-benchmark --pin-memory --num-workers 4 --wandb
rois.npz --model densenet121 --lr 1e-3 --lr-scheduler plateau --flip 3 --batch-size 64 -d cuda --cudnn-benchmark --pin-memory --num-workers 4 --wandb
rois.npz --model densenet121 --lr 1e-4 --lr-scheduler plateau --flip 3 --batch-size 64 -d cuda --cudnn-benchmark --pin-memory --num-workers 4 --wandb
rois.npz --model shufflenet2x1_0 --lr 1e-3 --lr-scheduler plateau --flip 3 --batch-size 128 -d cuda --cudnn-benchmark --pin-memory --num-workers 4 --wandb
rois.npz --model squeezenet1_0 --lr 1e-3 --lr-scheduler plateau --flip 3 --batch-size 128 -d cuda --cudnn-benchmark --pin-memory --num-workers 4 --wandb
```

### Quick test on Apple Silicon (M1/M2/M3)

Use a single lightweight model with reduced epochs to verify everything works before
sending to the GPU server. MPS does not support `--cudnn-benchmark`, `--pin-memory`,
or `--num-workers > 0`.

```text
rois.npz --model shufflenet2x1_0 --lr 1e-3 --lr-scheduler plateau --flip 3 --batch-size 32 -d mps --num-workers 0 --epochs 50
```

### Small or imbalanced dataset

When you have fewer than ~500 annotated cells per class, or a large class imbalance
(e.g. 10× more interphase than mitotic cells), combine focal loss with loss weights and
increase regularisation:

```text
rois.npz --model shufflenet2x1_0 --lr 1e-3 --lr-scheduler plateau --lr-patience 5 --flip 3 --rotate 180 --batch-size 64 -d cuda --loss-function focal_loss --focal-gamma 3 --loss-weights --dropout 0.5 --patience 20 --wandb
rois.npz --model densenet121 --lr 1e-4 --lr-scheduler plateau --lr-patience 5 --flip 3 --rotate 180 --batch-size 32 -d cuda --loss-function focal_loss --focal-gamma 3 --loss-weights --dropout 0.5 --patience 20 --wandb
```

### Unfreezing EfficientNet

By default `--freeze-weights` only trains the classification head. If accuracy plateaus,
add a second run without freezing to let the backbone adapt to fluorescence images:

```text
# Phase 1: fast convergence with frozen backbone
rois.npz --model efficientnetb3s --freeze-weights --lr 1e-3 --lr-scheduler plateau --flip 3 --batch-size 64 -d cuda --wandb --run-name effnet-frozen
# Phase 2: fine-tune full network at lower LR
rois.npz --model efficientnetb3s --no-freeze-weights --lr 1e-4 --lr-scheduler plateau --flip 3 --batch-size 64 -d cuda --wandb --run-name effnet-finetune
```

---

## Step 4 — Generate and run `batch.sh`

```bash
cellclass batch train.txt --script batch.sh
bash batch.sh
```

`cellclass batch` assigns unique checkpoint and metadata filenames automatically so runs
never overwrite each other. Comment out completed lines in `train.txt` before regenerating.

To run jobs in parallel (one per GPU on a multi-GPU machine):

```bash
cellclass batch train.txt --script batch.sh --background
bash batch.sh
```

Each run launches as a background process; add `--device cuda:0`, `cuda:1`, etc. per line
to assign GPUs explicitly.

### Data splits

The `.npz` is split into three non-overlapping sets at the start of each training run:

- **64 %** training
- **16 %** validation — monitored each epoch for early stopping
- **20 %** held-out test set — evaluated once at the end, never seen during training

These come from the defaults `--testing-size 0.2` and `--validation-size 0.2` (validation
is 20 % of the remaining 80 %). For small datasets try `--testing-size 0.1
--validation-size 0.1` to get an 80/10/10 split.

::: {.callout-warning}
The test split is **not saved to disk** — it is reconstructed at evaluation time by
replaying the same random split using `--data-seed`. The split is therefore only
genuinely held-out if you use **the same `.npz` and the same seed** as training. The
`.json` state file records the seed used, so always pass it to `cellclass test`:

```bash
cellclass test rois.npz.1.json --data-seed 42
```

If you add new annotations and regenerate `rois.npz`, the crop indices shift and the old
split no longer applies — the model may be tested on cells it trained on. Only regenerate
`rois.npz` when you intend to retrain from scratch.
:::

### Class imbalance

Imbalance is handled in the loss function, not by resampling:

- **Focal loss** (default) down-weights easy, correctly-classified examples so training
  concentrates on hard or rare cases. This handles moderate imbalance well on its own.
- **`--loss-weights`** additionally scales each class's loss contribution by its inverse
  frequency. Combine with focal loss when one class is 5× or more rarer than another.

The splits are not stratified, so with very small datasets (<50 cells per class) the test
set may by chance contain an unrepresentative ratio. Check per-class counts in the W&B
logs and fix `--data-seed` if this is a concern.

---

## Step 5 — Compare on Weights & Biases

Open your W&B dashboard at <https://wandb.ai> and navigate to the **cellclass** project in your workspace. The key metrics to compare:

- **val/acc** — validation accuracy (primary metric for model selection).
- **test/acc** — held-out test accuracy (reported at end of training).
- **val/loss** — validation loss curve (check for overfitting).

Download the best run's metadata JSON from the W&B artifacts panel, or find it locally
in your working directory.

---

## Step 6 — Extract the model

```bash
cellclass extract best_run.json
```

Without `--save` this prints the performance metrics without writing any files — useful
for a quick sanity check. Add `--save` to export:

```bash
cellclass extract best_run.json --save
```

Output files:

- `<model>_c<channels>_l<labels>.pt` — TorchScript model.
- `<model>_c<channels>_l<labels>.json` — sidecar metadata (channels, labels, image size).

---

## Step 7 — Test the exported model

```bash
cellclass test -s shufflenet2x1_0_c2_l2.pt ~/data/mitosis-rpe/rois.npz
```

Prints a per-class precision/recall/F1 report and overall accuracy on the test split.

---

## Resuming a crashed run

Training checkpoints are saved after each epoch. To restart from the last checkpoint:

```bash
cellclass train best_run.json
```

Passing the metadata JSON resumes from the saved checkpoint. Note: some training state
(e.g. convergence checker) is not persisted and is recreated from scratch.


## Overview

cellclass includes a curated set of convolutional architectures suited to small
fluorescence microscopy crops (typically 50 × 50 px, 1–3 channels). The models fall into
two groups: **lightweight baselines** that train quickly and are easy to deploy, and
**pretrained transfer-learning models** that leverage ImageNet features when labelled data
is scarce.

---

## Recommended models for small cell images

### Lightweight (no pretrained weights)

These models train from scratch and are well suited to grayscale or 2-channel images
where ImageNet features are less relevant.

| Model | Key strengths | Recommended batch size | Notes |
|---|---|---|---|
| `shufflenet2x1_0` | Fast, memory-efficient, excellent accuracy/speed trade-off | 128 | **First choice** for most cell classification tasks |
| `shufflenet2x1_5` | Slightly wider than `2x1_0`, marginal accuracy gain | 64 | Use when `2x1_0` plateaus |
| `shufflenet2x2_0` | Widest ShuffleNet, highest capacity | 32 | Use for complex multiclass problems |
| `squeezenet1_0` | Smallest model, fastest inference | 128 | Good sanity-check / baseline |
| `densenet121` | Strong regularisation via dense connections | 64 | Reliable performer across datasets |
| `densenet161` | More capacity than 121 | 32 | Diminishing returns on 50 px crops |

### Pretrained (ImageNet transfer learning)

Use `--freeze-weights` to freeze the backbone and only train the classification head.
This is beneficial when labelled data is limited (< ~500 crops per class).

| Model | Input size | Key strengths | Notes |
|---|---|---|---|
| `efficientnetb3s` | 300 × 300 (resized automatically) | Best accuracy with pretrained features | **Recommended pretrained model** |
| `efficientnetb4s` | 380 × 380 | Highest accuracy | Slower; use when `b3s` is insufficient |

> **Note on image resizing**: EfficientNet models internally resize inputs. For 50 × 50 px
> crops this upscaling can hurt performance compared to architectures that accept any input
> size. If you have ≥ 1000 crops per class, `densenet121` often matches or exceeds
> EfficientNet.

### Attention models

The following models add spatial or channel attention modules to GoogLeNet- or
AlexNet-based backbones. They are experimental and have shown mixed results on small crops.

| Model | Description |
|---|---|
| `googlenet` / `googlenet_att` | GoogLeNet with optional spatial attention |
| `alexnet` / `alexnet_att` | AlexNet with optional spatial attention |
| `simplegooglenet` | Lightweight GoogLeNet using a subset of inception blocks |
| `simplegooglenet_att` | + spatial attention |
| `simplegooglenet_ch_att` | + channel attention |
| `simplegooglenet_ch_sp_att` | + both channel and spatial attention |

---

## Choosing a model: decision tree

```
Do you have ≥ 500 labelled crops per class?
├── Yes → Start with shufflenet2x1_0 (fast sweep) + densenet121
│         Compare on W&B; pick best.
└── No  → Use efficientnetb3s --freeze-weights
          Transfer learning compensates for small datasets.

Is inference speed critical (production deployment)?
├── Yes → shufflenet2x1_0 or squeezenet1_0
└── No  → Any of the above

Do you have > 3 classes or very similar morphologies?
└── Yes → Try densenet161 or shufflenet2x2_0 for higher capacity
```

---

## Hyperparameter guidance

### Learning rate

Always sweep at least two values: `1e-3` and `1e-4`. The plateau scheduler adapts
automatically, so a higher starting LR rarely hurts.

### Batch size

Larger batches train faster but may reduce generalisation on small datasets. Use the
recommended defaults from the table above; halve the batch size if you run out of memory.

### Data augmentation (`--flip`)

| Value | Augmentation applied |
|---|---|
| `0` | None |
| `1` | Horizontal flip only |
| `2` | Vertical flip only |
| `3` | Both horizontal and vertical flips (**recommended**) |

Cells in fluorescence images have no canonical orientation, so flipping both axes is
almost always beneficial.

### Class imbalance

By default, cellclass weights the loss function inversely proportional to class frequency.
Disable with `--no-loss-weights` if your classes are well balanced.

---

## Performance expectations (50 × 50 px, 2-class)

Based on the mitosis-RPE and ppase-screen datasets:

| Model | Typical test accuracy | Training time (per epoch, GPU) |
|---|---|---|
| `squeezenet1_0` | 88–92 % | ~5 s |
| `shufflenet2x1_0` | 92–96 % | ~8 s |
| `densenet121` | 93–96 % | ~15 s |
| `efficientnetb3s` (frozen) | 90–95 % | ~20 s |

Accuracy is highly dataset-dependent. Always evaluate on your own held-out test split.


### SLURM — Training on the Artemis Cluster

The University of Sussex [Artemis HPC cluster](https://artemis-docs.hpc.sussex.ac.uk/artemis/)
provides GPU nodes for large-scale training. Access requires VPN.

---

## Single job

Submit one training run to SLURM:

```bash
cellclass sbatch --args ~/data/mitosis-rpe/rois.npz \
    --model shufflenet2x1_0 --lr 1e-3 --lr-scheduler plateau \
    --flip 3 --batch-size 128 -d cuda --wandb
```

`cellclass sbatch` wraps `sbatch` and passes everything after `--args` to `cellclass train`.
Use `cellclass sbatch --help` to see SLURM-specific options (partition, memory, walltime),
or read them in the generated [`cellclass sbatch` reference](../reference/cli/sbatch.html).

---

## Batch of SLURM jobs

Use the same `train.txt` batch file as for local training, but pass the SLURM submitter
as the command:

```bash
cellclass batch train.txt --script batch.sh --cmd "cellclass sbatch --args"
bash batch.sh
```

This generates one `sbatch` submission per line in `train.txt`. Each job runs independently
on the cluster; results are logged to W&B as they complete.

**Important**: Comment out completed lines in `train.txt` before regenerating `batch.sh` to
avoid resubmitting finished runs.

---

## Weights & Biases on Artemis

Your W&B credentials stored in `~/.netrc` are available on Artemis if your home directory
is mounted (it is by default). Verify before submitting jobs:

```bash
wandb status
```

If credentials are missing, run `wandb login` from an Artemis login node.

---

## Monitoring jobs

```bash
squeue -u $USER          # List your queued/running jobs
scancel <job_id>         # Cancel a job
```

Training progress is also visible in real time on your W&B dashboard at <https://wandb.ai>.


### Deploying Models for Inference

Once you have extracted a TorchScript model you can use it in two ways:

1. **omero-screen pipeline** — automatic inference during plate analysis.
2. **omero-screen-napari** — interactive classification in the Gallery / Classifier Selector widget.

---

## Exported model files

`cellclass extract best_run.json --save` produces two files:

```
shufflenet2x1_0_c2_l2.pt    # TorchScript model (architecture + weights)
shufflenet2x1_0_c2_l2.json  # Sidecar: channel names, class labels, input size
```

The filename encodes the architecture, number of input channels (`c2`), and number of
class labels (`l2`). Both files must be kept together.

---

## omero-screen pipeline

The pipeline loads classifiers from OMERO, not from local files. Publish the
`.pt` together with its `.json` sidecar once:

```bash
omero-screen models publish shufflenet2x1_0_c2_l2.pt
```

This attaches both files to your `Classifiers` project in OMERO. To publish a
retrained model under the same name, add `--replace`. A new name (for example,
with a version) keeps the old model available.

Then run the pipeline with the model **name**, without `.pt`:

```bash
omero-screen <plate_id> --inference shufflenet2x1_0_c2_l2
```

Multiple models can be chained (several names after `--inference`, or
colon-separated in `OMERO_SCREEN_INFERENCE_MODEL`):

```bash
omero-screen <plate_id> --inference mitosis_model ppase_model
```

Downloaded models are cached in `~/.cache/omero_screen/` and fetched again if
the published file changes.

Gallery images showing example cells per predicted class are attached as PNG figures to
the OMERO plate. Configure gallery output with:

```bash
export OMERO_SCREEN_INFERENCE_GALLERY_WIDTH=10   # grid size (default: 10)
export OMERO_SCREEN_INFERENCE_BATCH_SIZE=100     # inference batch size (default: 100)
```

See the [pipeline configuration](../../user-guide/installation.html#configure) for the full list of variables.

---

## omero-screen-napari

Copy the `.pt` + `.json` pair to the location specified when creating the classifier in
the [Setup Training Widget](../../napari/user-guide/training_widget.html). The
[Classifier Selector](../../napari/user-guide/session_manager.html) widget will load the
model and display per-class statistics alongside the database annotation counts.

---

## Verifying a deployed model

Before committing to a production run, evaluate the model on your full dataset:

```bash
cellclass test -s shufflenet2x1_0_c2_l2.pt ~/data/mitosis-rpe/rois.npz
```

This prints precision, recall, and F1 per class, plus overall accuracy, on the held-out
test split. A model is ready for deployment when test accuracy is ≥ 90 % and no class
has recall below 0.85.


### Workflow (legacy notes)

::: {.callout-note}
This page is the original developer walkthrough. It refers to the standalone
scripts (`run_training.py`, `batch_training.py`, `sbatch_training.py`,
`extract_model.py`, `test_model.py`), which are now subcommands of the unified
`cellclass` command — see the [CLI reference](../reference/cli/index.html).
:::


Classify cell data using neural networks.

## Set-up

Prepare the environment using [uv](https://docs.astral.sh/uv/):

    uv sync

Programs located in `src/bin` can now be executed using `uv run`. To omit the `uv run`
prefix activate the environment using:

    source .venv/bin/activate

## Data preparation

Cell classification examples are pickled numpy data files saved as a dictionary.
Each example has an image in (YXC) format, a mask for the region of interest (YX)
and a classification label (str).
The `data` key is a tuple of lists of the images and masks. The `target`
key is the label for each training example.

    {
      'data': ([img1, img2, ..., imgN], [mask1, mask2, ..., maskN]),
      'target': [label1, label2, ..., labelN]
    }

Multiple data files may be present in a dataset. The neural network should be trained on preprocessed images.
These will be preprocessed to maintain only the data of relevance to the classification problem.
A typical workflow is to extract the (1, 99) percentile of the image data, zero all pixels outside the mask
and convert to unsigned 8-bit format (for memory efficiency). The data can be saved into one numpy file
to allow fast loading into memory. For example:

    cellclass dataset ~/images/training_data --channels DAPI

Channel names must be specified explicitly using the `--channels` flag, with one name per channel in the
order they appear in the image data. For example, a two-channel DAPI/RFP dataset:

    cellclass dataset ~/images/training_data --channels DAPI RFP

Alternatively, channel names can be defined in a `metadata.json` file in the data directory (see below),
in which case the `--channels` flag can be omitted.

This will create a single `.npz` file in the data directory named `rois.npz` containing the preprocessed images.
This program has different options; see the help for details.

Examples of each class can be extracted from the dataset using:

    cellclass sample ~/images/training_data/rois.npz --output ~/images/training_data

This will create a tif file for each class containing a random selection of images. The tif file has been written
using `ImageJ` hyperstack format. Multiple channel images can be viewed in `ImageJ` using the channels tool
(`Image > Color > Channels Tool...`) and selecting `Composite` .

## Training

### Weights & Biases

It is recommended to log training runs to [Weights & Biases](https://wandb.ai/). The environment can be configured
for a user with the following command:

    wandb login

Note that the `wandb` command will be available in the `uv` environment.
Visit the url https://wandb.ai/authorize and login. This will generate a key to provide to the login command.
The environment will now be configured to automatically login to Weights & Biases.

### Training

The training can be performed using multiple options and models.
Ideally training should be performed on an accelerator device such as a GPU.
To train a model on an Apple MPS device and log to Weights & Biases:

    ./run_training.py ~/images/training_data/rois.npz -d mps --model efficientnetb3s --wandb -n test -s run1

This program has different options; see the help for details. By default the input data is split into
non-testing and testing. The non-testing data is then split into training and validation. The training
is performed on the training data and validated on the validation data. The testing data is separate and
never used to influence training (e.g. determine model convergence). It is used to provide the expected
model performance on new data.

Training will create output files using the provided prefixes. In this example it is `test` for the checkpoint
files and `run1` for the metadata file.

The checkpoint name is used to save the training state.
This allows the run to be restarted from the current state, for example in the event that training crashes
due to a fixable issue it can be restarted from the previous point. Note that some training objects
may not be supported and are created again, for example the convergence checker. Consider restarts
as a work-in-progress beta feature.

The best model is saved into a checkpoint file with the `.best` suffix. This is the best performing model
on the validation data. It should be used to extract the final model weights and used for evaluation.

`.json`: Metadata file containing details of all the arguments used to run training.
Training can be restarted by passing this file to the program:

    ./run_training.py test.json

Note: Using the `--wandb` flag will log to Weights & Biases. At the end of training the metadata and checkpoint
files are uploaded as artifacts of the training run.

### Batch Training

Multiple different options for training can be prepared into a training batch file.
These are arguments to pass to the `run_training.py` program, for example:

    /path/to/rois.npz --model densenet161 --lr-scheduler plateau --no-loss-weights --batch-size 32
    /path/to/rois.npz --model densenet161 --lr-scheduler plateau --no-loss-weights --batch-size 64
    /path/to/rois.npz --model densenet161 --lr-scheduler plateau --no-loss-weights --batch-size 128
    /path/to/rois.npz --model densenet201 --lr-scheduler plateau --no-loss-weights --batch-size 32
    /path/to/rois.npz --model densenet201 --lr-scheduler plateau --no-loss-weights --batch-size 64
    /path/to/rois.npz --model densenet201 --lr-scheduler plateau --no-loss-weights --batch-size 128

These can be processed into a script:

    ./batch_training.py train.txt --script batch.sh

This will create a batch script to execute the training program for each line in the batch file. This processing script
ignores empty lines and commented lines beginning with `#`. This allows the same batch file to be used repeatedly
to track all the training runs that have been performed on a dataset. Note that old runs should be commented out
before rebuilding the training script.

The `batch_training.py` program will create appropriate names and output filenames to avoid overwriting existing files.
It will log results to Weights & Biases and generate a W&B run name using the concatenated arguments.

By default a batch script is created allowing the command to be verified before executing. The script can also
run each training run in a background process. This is not recommended unless the number of runs is small
as all will run concurrently. This option is to be used on a high specification machine that may have multiple
graphics cards. Each run can have the device identified as one of the arguments, e.g. `--device cuda:0`.

### Training on a SLURM cluster

#### Artemis

The University of Sussex Artemis cluster can be used to perform multiple GPU based training jobs.
Documentation on Artemis is here: https://artemis-docs.hpc.sussex.ac.uk/artemis/.
This requires VPN access.

#### Batching SLURM Jobs

The program `sbatch_training.py` will create and execute a script to submit a training job to a SLURM cluster
using `sbatch`. The script accepts some parameters to configure options; see the help for details. However
any parameters after the `--args` argument are passed to the `run_training.py` program.

To run a single job on a SLURM cluster the program is run as per `run_training.py`:

    ./sbatch_training.py --args ~/images/training_data/rois.npz -d mps --model efficientnetb3s --wandb -n test -s run1

To create a batch of jobs to submit to a SLURM cluster the program name can be passed to the batching script:

    ./batch_training.py train.txt --script batch.sh --cmd "./sbatch_training.py --args"

This effectively creates the same `batch.sh` script as previously but will run `sbatch_training.py` for
each run rather than `run_training.py`.

## Extracting Models

The training data metadata can be used to extract details of the model from the checkpoint files
such as the performance metrics, or save the model weights and metadata:

    ./extract_model.py training.json
    ./extract_model.py training.json --save

The `save` option will create a file name using the model name, the number of input channels
in the images and the number of class labels. The scripted model architecture and weights
are saved to a `.pt` file and a separate `.json` metadata file contains the channel names and class labels.
The `.pt` file uses the [TorchScript](https://pytorch.org/docs/stable/jit.html) format.

## Testing Models

Test a model from the checkpoint file on a dataset; requires knowing the model class:

    ./test_model.py --model densenet201 -n run123.pt.best ~/images/training_data/rois.npz

Test a TorchScript model on a dataset; requires the `.json` sidecar metadata file describing
the architecture (channels and classes, and input image size):

    ./test_model.py -s densenet201_c1_l3.pt ~/images/training_data/rois.npz

This script will create a performance metrics report on the precision, recall and accuracy of the model.