docs/custom_data.md

# Training on your own data

Pretraining reads **WebDataset `.tar` shards**; finetuning is supervised and reads **image/mask
directories**.

```
data/
  pretrain/                  # stage 1: WebDataset shards
    customdata-000.tar
    customdata-001.tar
    ...
  finetune/                  # stage 2: supervised splits
    train/  images/  masks/
    val/    images/  masks/
    test/   images/  masks/
```

`data/` does not have to be local. Configs support env vars and default reads
`${DATA_ROOT:-data}/...`, so you can either symlink it

```bash
ln -s /mnt/big-disk/affmae-data data
```

or point the variable elsewhere and leave the repository alone:

```bash
export DATA_ROOT=/mnt/big-disk/affmae-data
```

Relative paths resolve against the **repository root**, not your shell's working
directory, so training scripts work from anywhere. The same applies to
`CHECKPOINT_ROOT` (default `weights`) and `AFFMAE_OUTPUT_DIR` (default `output`).

---

## Stage 1 — pretraining data (WebDataset)

### Shard layout

Pretraining is masked autoencoding, so just images. Each
shard is an uncompressed `.tar` of images, and `webdataset` groups files by
basename, so one image per basename is all that is required:

```
customdata-000.tar
├── 000000.png
├── 000001.png
├── 000002.png
└── ...
```

Three things to consider when building shards:
1. **The extension is the key.** The loader calls `.decode("pil")` and then reads
   `sample["png"]`, so files must be `.png`. Rename before packing, or change the
   `map_dict(png=...)` key in `affmae/data/pretrain_dataset.py` to match your
   extension.
2. **Shards should be roughly equal in size**, ideally 100–1000 images each.
   Under multi-GPU training each rank gets a disjoint set of shards, so one giant
   shard leaves other ranks idle.
3. **Have at least as many shards as `num_workers` × ranks.** Shards are split
   across DataLoader workers *and* across ranks, so a worker that gets none makes
   WebDataset raise *"No samples found in dataset; perhaps you have fewer shards
   than workers."* Three shards with `num_workers: 8` fails, and it fails at the
   first batch rather than at startup.

### Building shards

```python
import webdataset as wds
from pathlib import Path

images = sorted(Path("/my/raw/images").glob("*.png"))

# maxcount caps images per shard; %03d must match the brace range in the config
with wds.ShardWriter("data/pretrain/customdata-%03d.tar", maxcount=1000) as sink:
    for i, path in enumerate(images):
        sink.write({"__key__": f"{i:08d}", "png": path.read_bytes()})
```

`ShardWriter` prints the shard count when it finishes. If it wrote
`customdata-000.tar` through `customdata-031.tar`, the config pattern is
`customdata-{000..031}.tar` — the brace range is **inclusive** and the digit
count must match the filenames exactly.

### Counting your samples

A WebDataset is an `IterableDataset` with no length, so the epoch boundary comes
from `data.total_samples` in the config rather than from the data. Getting it
wrong does not crash: too high repeats images, too low silently truncates the
epoch. Count once:

```bash
python -c "
import tarfile, glob
n = sum(sum(1 for m in tarfile.open(f) if m.name.endswith('.png'))
        for f in glob.glob('data/pretrain/*.tar'))
print(n)
"
```

### Running it

```bash
cp configs/template_pretrain_512.yaml configs/my_pretrain.yaml
# edit data.path, data.total_samples, data.in_channels
python pretrain.py --config configs/my_pretrain.yaml
```

Multi-GPU is mostly untested (use at your own rick) with `torchrun`; `data.batch_size` is the **global** batch and is divided
across ranks, so the effective batch does not change with GPU count:

```bash
torchrun --nproc_per_node=4 pretrain.py --config configs/my_pretrain.yaml
```

---

## Stage 2 — finetuning data (images and masks)

### Directory layout

```
data/finetune/
  train/
    images/   img_0001.tif   img_0002.tif   ...
    masks/    img_0001.tiff  img_0002.tiff  ...
  val/        images/  masks/
  test/       images/  masks/
```

Things to consider for finetuning on your data:
- **Images and masks pair by filename stem.** `img_0001.tif` needs
  `img_0001.tiff`. The loader takes the *intersection* of the two directories, so
  an unpaired file is skipped **without warning**.
- **Image extension is configurable** via `data.input_ext` (default `.tif`).
  **The mask extension is fixed at `.tiff`** in
  `affmae/data/finetune_dataset.py`.
- `val/` is the split evaluated during training. `test/` is read only by
  `evaluate.py`. All three are required by the templates; point `val` and `test`
  at the same files if you only have two splits.

### Mask format
Masks are multi-channel (ie multi-label) TIFFs, one binary channel per annotated structure — **not** a single-channel label map
with integer class IDs.

`data.indices` selects which channels to train on, and background is added
implicitly:

```yaml
data:
  indices: [0, 1]      # train on mask channels 0 and 1
  num_classes: 3       # len(indices) + 1 for background
```

So `num_classes` is typically `len(indices) + 1`, depending on loss selected, and `train.class_weighting` must
have exactly `num_classes` entries, background first. A wrong length fails at
loss construction with *"weight tensor should be defined either for all or no classes"*.

The selected channels are flattened into a single label map in `__getitem__`:
each channel is binarized at `THRESHOLD = 10`, and a pixel is assigned
`position_in_indices + 1`. Two consequences:

- **Class IDs follow the order of `indices`, not the channel numbers.**
  `indices: [2, 3]` makes channel 2 class 1 and channel 3 class 2.
- **Channels are applied in order, so on overlapping annotations the later channel wins.** Put the class you care most about last.

A single-channel integer label map is **not** supported — `__getitem__` unpacks
`C, H, W` and a 2-D mask raises a `ValueError`. One-hot it into channels when preparing the data.

Raise a class's weight if it is rare — the shipped `[0.2, 2.0, 3.0]` deliberately down-weights background.

### Running it

```bash
cp configs/template_finetune_512.yaml configs/my_finetune.yaml
# edit data.base_path, data.indices, data.num_classes,
#      train.class_weighting, model.pretrained_ckpt_path
python finetune.py --config configs/my_finetune.yaml
```

You do not have to pretrain first. To start from our released backbone, fetch it
once and point the config at the result:

```python
from affmae.data.weights import EMWeights
print(EMWeights.AFFMAE_BASE_PRETRAIN_512.fetch())
# -> weights/pretrain/ckpt_epoch_399_affmae_fpw.pth
```

```yaml
model:
  pretrained_ckpt_path: "${CHECKPOINT_ROOT:-weights}/pretrain/ckpt_epoch_399_affmae_fpw.pth"
```

The encoder architecture keys in your config must match that checkpoint, so keep
the `aff_*` block from the template unchanged unless you are pretraining yourself.
See the weights table in the README for every released checkpoint.

---

## Checklist before a long run

Cheap to check, expensive to discover at hour six.

| Check | Symptom if wrong |
|---|---|
| `data.total_samples` matches your shard count | epoch silently truncated or images repeated |
| Brace range digits match filenames | `FileNotFoundError` on the first batch |
| `num_classes == len(indices) + 1` | crash in the loss, or a dead output channel |
| `len(class_weighting) == num_classes` | *"weight tensor should be defined either for all or no classes"* |
| `in_channels` matches your images (1 grey, 3 RGB) | shape mismatch in the patch embedding |
| Encoder architecture keys identical to pretraining | checkpoint loads with missing/unexpected keys |
| Dataset size is what you expect | unpaired image/mask stems dropped silently |
| Shards >= `num_workers` x ranks | *"No samples found in dataset"* at the first batch |
| `stats.py` matches your data | washed-out or clipped figures; slower convergence |
| CLAHE removed if not wanted | real intensity differences flattened |
| Info-bar crop removed for non-EM images | bottom of bright images silently truncated |

Sanity-check the pipeline before committing to a full run:

```bash
python -c "
from affmae.config import load_config
from affmae.data.finetune_dataset import build_finetune_dataloader
cfg = load_config('configs/my_finetune.yaml')
loader = build_finetune_dataloader(cfg, is_train=True)   # False reads val/
img, mask = next(iter(loader))[:2]
print('image', img.shape, img.dtype)
print('mask ', mask.shape, 'classes present', mask.unique().tolist())
"
```

`mask.unique()` should be class indices in `[0, num_classes)`. If you see values
above that, `indices` and `num_classes` disagree.

---

## Adapting the preprocessing

Everything in this section is tuned for greyscale electron micrographs. None of it
is wrong for other data, but none of it is neutral either — these are the four
places where "trained on EM" is baked in.

### 1. Images are single-channel

We trained on greyscale. Two consequences:

- `data.in_channels: 1` in the config. Set it to 3 for RGB; the patch embedding
  and the normalization tuple both follow it.
- The **pretraining loader converts to greyscale regardless**:
  `apply_custom_processing` in `affmae/data/pretrain_dataset.py` calls
  `pil_image.convert("L")`. For RGB pretraining, remove that call — otherwise
  `in_channels: 3` silently receives a replicated grey channel.

### 2. Normalization statistics

`affmae/data/stats.py` holds four numbers measured on our data:

```python
IMAGE_MEAN = 0.6266           # labelled FPW split: finetuning and inference
IMAGE_STD  = 0.2259
PRETRAIN_IMAGE_MEAN = 0.5562  # unlabelled pretraining corpus
PRETRAIN_IMAGE_STD  = 0.2396
```

**Change these for your own data.** They are used by the loaders *and* by the
renderers, which denormalize with the same numbers — so a mismatch shows up as
washed-out or clipped figures long before you suspect the statistics.

Measure yours with:

```bash
python scripts/calc_norm.py --config configs/my_pretrain.yaml
```

It reports mean/std over your shards and prints the currently configured values
next to them. It measures the *raw* pipeline deliberately: an earlier version read
through the normalized pipeline and reported the residual (~0, ~1), which looks
like a passing result and is useless.

For multi-channel data, note these are scalars broadcast across channels
(`(mean,) * n_channels`). Per-channel statistics need `create_transforms` in
`affmae/data/pretrain_dataset.py` changed to take tuples.

### 3. CLAHE

Contrast-limited adaptive histogram equalization is applied in **four** places,
because EM micrographs have low local contrast. If your images are already
well-exposed, this will flatten real intensity differences and you should remove
it:

| Where | What | Flag? |
|---|---|---|
| `affmae/data/preprocess.py` | `apply_clahe`, `clipLimit=4.25`, `8x8` tiles | yes — `preprocess_image(..., use_clahe=False)` |
| `affmae/data/finetune_dataset.py:116, 214` | same parameters, inline | **no flag, edit the call** |
| `affmae/data/pretrain_dataset.py:80` | `clipLimit=4.0`, inside `apply_custom_processing` | **no flag, edit the call** |
| `affmae/data/pretrain_dataset.py:58` | random CLAHE, an augmentation | part of `random_transform` |

Only the inference path takes a flag. The dataset classes hardcode it, so removing
CLAHE from training means editing those call sites. If you do, **re-measure the
normalization statistics afterwards** — CLAHE changes the intensity distribution,
so the old mean/std no longer describe your inputs.

### 4. The microscope info bar

`apply_custom_processing` in `affmae/data/pretrain_dataset.py` crops the bottom of
the image when it finds a row that is more than 95% pure white below the 60%
height mark. That removes the scale-bar and metadata strip our microscope burns
into every frame.

This is the most EM-specific step in the pipeline and the most likely to damage
other data: any image with a genuinely bright bottom edge — a light background, a
white border — gets silently truncated, and because the crop is variable-height
the resize then changes the aspect ratio per image. **Remove it unless your images
have the same burned-in strip.**

### 5. Augmentation

`ElasticTransform` in `affmae/data/transforms.py` plus affine and photometric
transforms are tuned for tissue sections, which deform plausibly. For rigid
subjects (documents, industrial parts) elastic warping invents geometry that
cannot occur; see `scripts/visualize_aug.py` to look at what the current pipeline
does to your images before training on it:

```bash
python scripts/visualize_aug.py --base-path data/finetune --output output/aug.png
```

## Adapting the architecture

Defaults are AFF-Base at 512px with patch size 8. What is worth changing:

| Goal | Change |
|---|---|
| Larger images | `data.img_size` 768 or 1024, and lower `batch_size` |
| Smaller/faster model | shorten `aff_depths`, narrow `aff_embed_dims` |
| More tokens kept per stage | raise `aff_ds_rates` (0.4 keeps 40% per stage) |
| Different masking difficulty | `model.mask_ratio`, pretraining only |
| Out of memory | lower `batch_size`, raise `train.num_accum` to compensate |

Two big constraints:

- **`decoder_depth` must equal the number of encoder stages** (4 for AFF-Base),
  one decoder block per stage.
- **Encoder keys must match between pretraining and finetuning.** The finetune
  config re-declares `aff_embed_dims`, `aff_depths`, `aff_num_heads`,
  `aff_nbhd_sizes`, `aff_cluster_size`, `aff_ds_rates`, `aff_mlp_ratio` and
  `aff_merging_method`; any disagreement means the checkpoint does not load into
  the encoder.

`img_size` is the exception — it may differ between stages. Cost is driven by the
token count, which scales with `img_size / patch_size` rather than with pixel
count, so moving from 512 to 1024 at patch size 8 doubles the grid on each axis.

Leave `cluster_attention_backend` and `decoder_deform_backend` at `"auto"` unless
you have profiled the alternative. `"auto"` selects the fused Triton kernels on
CUDA and the PyTorch fallback elsewhere, so the same config runs on CPU, CUDA and
Apple silicon; see *Hardware and Kernel Coverage* in the README.