cellpose icon indicating copy to clipboard operation
cellpose copied to clipboard

3D fine-tuning of CPSAM leads to no mask output

Open shk44 opened this issue 6 months ago • 1 comments

I am trying to fine-tune the cpsam on 3D labeled images and am running into issues. I am running cellpose 4.0.4.dev0+ga75db82.d20250529 through a jupyter notebook interface on an HPC running linux. I have tried updating to the stable version of 4.0.4 to no avail. I have similarly tried the training notebook, but the same issue occurs where no masks are produced.

My training data involves one image (69 x 1024 x 1024) with thousands of puncta-like masks and initial training seems to learn ok, dropping the train_loss from ~0.1 to ~0.001 within 20 epochs, see below. However, when I use that model to try and recapitulate my masks on my training image, I receive no masks. I wonder if it is an issue of training the model in 3D vs 2D, or simply not enough training data? Thank you for your help, Cellpose is amazing!

from cellpose import io, models
io.logger_setup()

train_dir = '/home/shk441/cellpose_plugin/cellpose/train'
model_name = 'syt2_model100'

# === Load training + test data ===
output = io.load_train_test_data(
    train_dir, 
    image_filter="_img", 
    look_one_level_down=False
)
images, labels, image_names, _, _, image_names_test = output

# === Initialize model ===
model = models.CellposeModel(gpu=True, pretrained_model='cpsam')  # or 'cyto', 'nuclei'

# === Train model ===
model_path, train_losses, test_losses = train.train_seg(model.net,
    train_data=images,
    train_labels=labels,
    learning_rate=1e-5,
    weight_decay=0.1,
    n_epochs=30,
    model_name=model_name,  # saved to ~/.cellpose/models/my_new_model/
    save_every=1           # grayscale images                # if using 3D training volumes
)

Which outputs

2025-06-10 16:59:37,654 [INFO] WRITING LOG OUTPUT TO /home/shk441/.cellpose/run.log
2025-06-10 16:59:37,656 [INFO] 
cellpose version: 	4.0.4.dev0+ga75db82.d20250529 
platform:       	linux 
python version: 	3.10.13 
torch version:  	2.6.0+cu124
2025-06-10 16:59:37,666 [INFO] not all flows are present, running flow generation for all images
2025-06-10 16:59:37,705 [INFO] reading tiff with 69 planes
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 463.14it/s]
2025-06-10 16:59:37,862 [INFO] reading tiff with 69 planes

100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 427.54it/s]
2025-06-10 16:59:38,028 [INFO] 1 / 1 images in /home/shk441/cellpose_plugin/cellpose/train folder have labels
2025-06-10 16:59:38,029 [INFO] ** TORCH CUDA version installed and working. **
2025-06-10 16:59:38,030 [INFO] >>>> using GPU (CUDA)
2025-06-10 16:59:40,112 [INFO] >>>> loading model /home/shk441/.cellpose/models/cpsam
2025-06-10 16:59:40,690 [INFO] flows precomputed
2025-06-10 16:59:40,873 [INFO] >>> computing diameters
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 271.60it/s]
2025-06-10 16:59:40,881 [INFO] >>> normalizing {'lowhigh': None, 'percentile': None, 'normalize': True, 'norm3D': True, 'sharpen_radius': 0, 'smooth_radius': 0, 'tile_norm_blocksize': 0, 'tile_norm_smooth3D': 1, 'invert': False}
2025-06-10 16:59:40,937 [INFO] >>> n_epochs=30, n_train=1, n_test=None
2025-06-10 16:59:40,938 [INFO] >>> AdamW, learning_rate=0.00001, weight_decay=0.10000
2025-06-10 16:59:40,942 [INFO] >>> saving model to /home/shk441/cellpose_plugin/cellpose/models/syt2_model100

2025-06-10 16:59:41,704 [INFO] 0, train_loss=0.1140, test_loss=0.0000, LR=0.000000, time 0.76s
2025-06-10 16:59:42,375 [INFO] saving network parameters to /home/shk441/cellpose_plugin/cellpose/models/syt2_model100
(further saves snipped)
2025-06-10 16:59:54,533 [INFO] 5, train_loss=0.1215, test_loss=0.0000, LR=0.000006, time 13.59s
2025-06-10 17:00:09,755 [INFO] 10, train_loss=0.0136, test_loss=0.0000, LR=0.000010, time 28.82s
2025-06-10 17:00:40,652 [INFO] 20, train_loss=0.0016, test_loss=0.0000, LR=0.000010, time 59.71s

Given this, I can proceed to test the model on the original data:

model = models.CellposeModel(gpu=True, pretrained_model='/home/shk441/cellpose_plugin/cellpose/models/syt2_model100')

img = imread('/home/shk441/cellpose_plugin/cellpose/train/final_syt2_img.tif')
img = img[:5,:200,:200] #trimmed to save compute time
masks, flows, styles = model.eval(
    img,
    do_3D=True,
    channel_axis=None,
    z_axis=0,
    diameter=None,
     progress=True
    # auto size
)

import matplotlib.pyplot as plt
plt.imshow(np.max(masks, axis=0), cmap='nipy_spectral')
plt.title("Model prediction on training image")
plt.axis('off')
plt.tight_layout()
plt.show()

which results in,

2025-06-10 17:07:50,682 [INFO] ** TORCH CUDA version installed and working. **
2025-06-10 17:07:50,686 [INFO] >>>> using GPU (CUDA)
2025-06-10 17:07:52,934 [INFO] >>>> loading model /home/shk441/cellpose_plugin/cellpose/models/syt2_model100
2025-06-10 17:07:53,982 [INFO] running YX: 5 planes of size (1024, 1024)
2025-06-10 17:07:53,984 [INFO] 0%|          | 0/5 [00:00<?, ?it/s]
2025-06-10 17:08:06,400 [INFO] 100%|##########| 5/5 [00:12<00:00,  2.48s/it]
2025-06-10 17:08:06,458 [INFO] running ZY: 1024 planes of size (5, 1024)
2025-06-10 17:08:06,460 [INFO] 0%|          | 0/1024 [00:00<?, ?it/s]
2025-06-10 17:08:36,994 [INFO] 3%|3         | 31/1024 [00:30<16:18,  1.02it/s]
*snip*
2025-06-10 17:25:04,088 [INFO] 100%|##########| 1024/1024 [16:57<00:00,  1.01it/s]
2025-06-10 17:25:04,342 [INFO] running ZX: 1024 planes of size (5, 1024)
2025-06-10 17:25:04,344 [INFO] 0%|          | 0/1024 [00:00<?, ?it/s]
2025-06-10 17:25:35,154 [INFO] 3%|3         | 31/1024 [00:30<16:26,  1.01it/s]
*snip*
2025-06-10 17:42:03,029 [INFO] 100%|##########| 1024/1024 [16:58<00:00,  1.01it/s]
2025-06-10 17:42:03,329 [INFO] network run in 2049.41s
2025-06-10 17:42:03,713 [INFO] No cell pixels found.
2025-06-10 17:42:03,788 [INFO] masks created in 0.09s

shk44 avatar Jun 12 '25 01:06 shk44

Hi to do training for 3d segmentation you have to pre-process your masks into 2d slices. You can follow the instructions here: https://cellpose.readthedocs.io/en/latest/do3d.html#training-for-3d-segmentation

mrariden avatar Jun 12 '25 15:06 mrariden

closing due to inactivity

mrariden avatar Jul 29 '25 17:07 mrariden