Semantic segmentation pipeline for underwater acoustic (echogram) data using PyTorch Lightning + segmentation-models-pytorch. The notebook Train_Model.ipynb covers:

1. Path setup & data directory structure
2. Mapping processed ZARR files to EVR region annotation files
3. Preprocessing & patch generation (multi-class + binary masks)
4. PyTorch dataset & augmentations
5. Train / validation / test split by day (to reduce leakage)
6. U-Net (ResNet encoder) training with Focal Loss + metrics
7. Threshold selection & prediction utilities

Original notebook was created by @eovchinn and is is located here.

The code here is accompanying a blog article here:

It is recommended to use a Python virtual environment (>=3.9). Below are instructions for macOS / Linux (bash/zsh) and Windows (PowerShell).

python3 -m venv .venv
source .venv/bin/activate

pip install -r requirements.txt

If you need a specific CUDA build of PyTorch, visit: https://pytorch.org/get-started/locally/ and install the wheel they recommend, then run the rest of the dependency install (you can omit torch/torchvision from requirements.txt by editing it first).

jupyter lab

Then open Train_Model.ipynb.

If you add the environment after Jupyter was installed globally, you may need to install an IPython kernel:

python -m ipykernel install --user --name echogram-seg --display-name "Echogram Seg (venv)"

Select that kernel inside Jupyter Lab.

Update the DATA_BASE_DIR variable near the top of the notebook. Expected structure (example):

DATA_BASE_DIR/
  EVR_region_files/            # .evr annotation files
  processed/                   # processed ZARR subdirectories
  ML_data/
    imgs/                      # generated .npy image patches
    masks/                     # generated multi-class masks
    binary_masks/              # generated binary masks

The preprocessing cells will populate ML_data/* after scanning ZARR + EVR pairs.

Run cells in order until the preprocessing section. This will:

• Load & concatenate ZARR datasets for a date
• Align region annotations (EVR)
• Slice the echogram into width-constrained patches
• Normalize channels & write .npy arrays for images, masks, and binary masks

You can restrict processing to a single date by inserting a break or filtering the loop.

Key parameters to tweak:

• batch_size
• encoder (e.g., resnet18, resnet50, timm-efficientnet-b0, etc.)
• Learning rate in the optimizer (Adam currently with lr=5e-3)
• alpha / gamma in FocalLoss
• Early stopping patience

Launch training by executing the trainer cell:

trainer.fit(pl_model, train_loader, val_loader)

Check TensorBoard for progress:

tensorboard --logdir outputs/tb_logs

Visit http://localhost:6006.

After training you can:

# Load best checkpoint if not already in memory
# pl_model = SegModel.load_from_checkpoint('path/to/checkpoints/best_model.ckpt', ...)

results_df = pl_model.select_threshold(val_loader, device=device)
results_df.sort_values('IoU', ascending=False).head()

Use the best threshold to set pl_model.threshold before running trainer.test(...) or custom prediction logic.

from torch.utils.data import DataLoader
pred_loader = DataLoader(val_dataset, batch_size=4, shuffle=False)
output_dir = 'outputs/preds'
os.makedirs(output_dir, exist_ok=True)
pl_model.predict(pred_loader, device=device, output_dir=output_dir, file_prefix='val')

Each batch’s raw logits will be saved as .npy. Post-process with sigmoid > threshold for binary masks.

• Set random.seed, np.random.seed, and torch.manual_seed at the start for deterministic splits.
• Pin library versions in requirements.txt if you move to production.
• Record the run_name (includes timestamp) to associate checkpoints and TensorBoard logs.

MIT