Training an image-to-image (virtual staining) model with BiaPy - NL-BioImaging Advanced Image Analysis Course 2026

Same input as 01_training_instance_segmentation.ipynb, a different approach.

Instead of nuclei labels, the model predicts the DAPI image itself from the phalloidin channel. This is called virtual staining: predicting what a stain would have shown, without staining for it.

The output is an image, not objects. To count or measure nuclei you still need to segment it.

Again we need to organize the data into training, validation and test data.

Kernel: biapy on Research Cloud, nlbi26-day3-biapy on your own laptop.

from pathlib import Path
import shutil

import numpy as np
import yaml
import matplotlib.pyplot as plt
from IPython.display import Image, display

from skimage.io import imread

1. The data¶

The target is now the DAPI image in training_data/nuclei/. Copy it into each split, next to images/ and labels/, so train, validation and test stay the same as in 01.

DATA = Path('training_data')
DATASET = Path('dataset')

for split in ('train', 'val', 'test'):
    folder = DATASET / split / 'nuclei'
    folder.mkdir(exist_ok=True)
    for f in (DATASET / split / 'images').glob('*.tif'):
        shutil.copy(DATA / 'nuclei' / f.name, folder / f.name)
    print(split, len(list(folder.glob('*.tif'))), 'images')

2. The settings¶

As in 01: take BiaPy’s template, change a few settings, save it under a new name.

patch_size = 256     # the network trains on crops of this size, not whole images
epochs = 50          # one epoch is one pass over the training images

job_name = 'dapi_from_phalloidin'
run_id = 1

output_path = 'training_output'

from urllib.request import urlretrieve

template = Path('2d_image-to-image.yaml')
if not template.exists():
    urlretrieve(
        'https://raw.githubusercontent.com/BiaPyX/BiaPy/v3.7.1/templates/'
        'image-to-image/2d_image-to-image.yaml',
        template,
    )

print(template.read_text())

config = yaml.safe_load(template.read_text())

# where the data is: the phalloidin images as input, the DAPI images as target
config['DATA']['TRAIN']['PATH'] = 'dataset/train/images'
config['DATA']['TRAIN']['GT_PATH'] = 'dataset/train/nuclei'
config['DATA']['TEST']['PATH'] = 'dataset/test/images'
config['DATA']['TEST']['GT_PATH'] = 'dataset/test/nuclei'

# validation from its own folder, not taken out of the training images
config['DATA']['VAL']['FROM_TRAIN'] = False
config['DATA']['VAL']['PATH'] = 'dataset/val/images'
config['DATA']['VAL']['GT_PATH'] = 'dataset/val/nuclei'

# PATCH_SIZE has to be a string, that is how BiaPy reads it
config['DATA']['PATCH_SIZE'] = f'({patch_size}, {patch_size}, 1)'
config['TRAIN']['EPOCHS'] = epochs

run_config = Path('2d_image-to-image_config.yaml')
run_config.write_text(yaml.safe_dump(config, sort_keys=False))

print(run_config.read_text())

3. Train¶

A few minutes on a GPU. Run biapy_out.show() in a new cell if you want to read the output.

%%capture biapy_out
from biapy import BiaPy

biapy = BiaPy(config=str(run_config), result_dir=output_path, name=job_name,
              run_id=run_id, gpu='0')

biapy.run_job()

4. Check the training graphs¶

The loss should go down, PSNR and SSIM should go up. Read the curves as in 01: still improving, flat, or overfitting?

results = Path(output_path) / job_name / 'results' / f'{job_name}_{run_id}'

for chart in sorted((results / 'charts').glob('*.png')):
    display(Image(str(chart)))

5. What does it predict?¶

BiaPy already ran the model on the test images. Input, real DAPI and prediction side by side.

test_ids = sorted(f.stem for f in (DATASET / 'test' / 'images').glob('*.tif'))

for image_id in test_ids[:3]:
    panels = [imread(DATASET / 'test' / 'images' / f'{image_id}.tif'),
              imread(DATASET / 'test' / 'nuclei' / f'{image_id}.tif'),
              imread(results / 'per_image' / f'{image_id}.tif').squeeze()]

    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
    for ax, img, title in zip(axes, panels, ['phalloidin', 'DAPI', 'prediction']):
        ax.imshow(img, cmap='gray', vmin=np.percentile(img, 1), vmax=np.percentile(img, 99.8))
        ax.set_title(title)
        ax.axis('off')
    plt.show()

Where does the prediction differ from the real DAPI? Faint nuclei, bright ones, edges?

If you have time¶

Each of these is one line, then re-run the config cell, the training cell and the result cells. Give every run its own job_name.

config['MODEL']['ARCHITECTURE'] = 'unet'                    # plain U-Net instead of attention_unet
config['LOSS'] = {'TYPE': ['SSIM'], 'WEIGHTS': [1.0]}       # structure instead of pixel difference

Is the virtual DAPI good enough to segment? Run StarDist on the predictions in per_image/ as in 00_data_collection_stardist.ipynb (that notebook’s kernel), and compare the nuclei with the labels in dataset/test/labels/.