How to improve nuclear segmentation and tracking

[XiangruiJi]

Hi,

I'm trying to segment and track nuclei in embryos using the light microscopy (LM) model, but the performance is not ideal.

I draw a lasso prompt in the initial frame (using a point prompt segments the entire embryo), and select the "mask" projection (the "single_point" projection loses the nuclei even by the next frame).

However, after a few frames the mask no longer aligns with the nucleus.

Do you have any suggestions on how to improve the performance? Thanks!

[anwai98]

Hi @XiangruiJi,

Thanks for your interest in micro-sam. And thanks for sharing all the things you've tried on your task at hand!

However, after a few frames the mask no longer aligns with the nucleus.

Could you try to add additional rectification points in the slices where it goes wrong and segment all frames again?

Although, I have a feeling adding subsequent rectification points might improve the segmentation quality by a bit, but it would probably still be quite some effort to do that for all nuclei!

In this case, I would suggest to train your own custom micro-sam model (annotate all nuclei in a few slices, train micro-sam on those few annotated slices, and then see if it improves your segmentation quality). There's a section in our doc for training your own models: https://computational-cell-analytics.github.io/micro-sam/micro_sam.html#training-your-own-model.

Let us know if you come across any issues in training your own model / applying the trained model on your own data.

[XiangruiJi]

Hi @anwai98 ,

Thank you for your suggestions! I am now trying the example notebook https://github.com/computational-cell-analytics/micro-sam/blob/master/notebooks/sam_finetuning.ipynb on HPC, but I find that the jupyter lab server cannot access the "checkpoints" directory generated.

This seems to be a bug for jupyter lab (jupyter-server/jupyter_server#950). Are there ways to get around with this?

[anwai98]

Hi @XiangruiJi,

Thanks for getting back. Since it's on your HPC system, I suggest finding the file and figure out how to download it / transfer it to your local device. I don't have an answer off the top of my head for your system, I'm sorry! :(

PS. How is the segmentation going / training outcome looking btw? Did the suggestions provided above work for your task?

[XiangruiJi]

Hi @anwai98 ,

I realized that my earlier issue was actually a path problem and I fixed it by changing the line for best_checkpoint to
best_checkpoint = os.path.join(root_dir, "models", "checkpoints", checkpoint_name, "best.pt")

However, I ran into another issue when trying the vit_l_lm model:

...

Referring to Running behind firewall. · Issue #1093 · computational-cell-analytics/micro-sam, I manually downloaded vit_l.pt from the provided link, renamed it as vit_l_lm, and moved it to the 'models' folder.

A similar error occurred when downloading the 'vit_l_lm_decoder' file.

I fixed this the same way. However, I noticed that the vit_l_decoder.pt file from the link is different from the one from BioImage.IO - Model Zoo (ViT-L) according to the file size.

(from the link) (from Model Zoo)

Which vit_l_decoder.pt file should I use? Thanks!

PS. Adding rectification points doesn't help much, so I plan to train my own model. As a first step, I am currently trying to run the example notebook.

[XiangruiJi]

I am also having difficulty adapting the example notebook https://github.com/computational-cell-analytics/micro-sam/blob/master/notebooks/sam_finetuning.ipynb for my own data. Unlike the example, my data and masks are .tif stacks containing time series of a single channel.

I have made the following changes:
raw_key, label_key= None, None

patch_shape = (73, 512, 512) (73 is the number of time points)

raw_paths = image_paths[0] (currently I only have one .tif stack in image_dir)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[39], line 2
      1 # Run training
----> 2 sam_training.train_sam(
      3     name=checkpoint_name,
      4     save_root=os.path.join(root_dir, "models"),
      5     model_type=model_type,
      6     train_loader=train_loader,
      7     val_loader=val_loader,
      8     n_epochs=n_epochs,
      9     n_objects_per_batch=n_objects_per_batch,
     10     with_segmentation_decoder=train_instance_segmentation,
     11     device=device,
     12 )

File [~/.conda/envs/microsam/lib/python3.10/site-packages/micro_sam/training/training.py:274](https://ondemand.nemo.thecrick.org/node/ga100.nemo.thecrick.org/22121/lab/tree/home/users/jix/microSAM_finetuning/~/.conda/envs/microsam/lib/python3.10/site-packages/micro_sam/training/training.py#line=273), in train_sam(name, model_type, train_loader, val_loader, n_epochs, early_stopping, n_objects_per_batch, checkpoint_path, with_segmentation_decoder, freeze, device, lr, n_sub_iteration, save_root, mask_prob, n_iterations, scheduler_class, scheduler_kwargs, save_every_kth_epoch, pbar_signals, optimizer_class, peft_kwargs, ignore_warnings, verify_n_labels_in_loader, box_distortion_factor, overwrite_training, **model_kwargs)
    272 t_start = time.time()
    273 if verify_n_labels_in_loader is not None:
--> 274     _check_loader(train_loader, with_segmentation_decoder, "train", verify_n_labels_in_loader)
    275     _check_loader(val_loader, with_segmentation_decoder, "val", verify_n_labels_in_loader)
    277 device = get_device(device)

File [~/.conda/envs/microsam/lib/python3.10/site-packages/micro_sam/training/training.py:41](https://ondemand.nemo.thecrick.org/node/ga100.nemo.thecrick.org/22121/lab/tree/home/users/jix/microSAM_finetuning/~/.conda/envs/microsam/lib/python3.10/site-packages/micro_sam/training/training.py#line=40), in _check_loader(loader, with_segmentation_decoder, name, verify_n_labels_in_loader)
     40 def _check_loader(loader, with_segmentation_decoder, name=None, verify_n_labels_in_loader=None):
---> 41     x, _ = next(iter(loader))
     43     # Raw data: check that we have 1 or 3 channels.
     44     n_channels = x.shape[1]

File [~/.conda/envs/microsam/lib/python3.10/site-packages/torch/utils/data/dataloader.py:708](https://ondemand.nemo.thecrick.org/node/ga100.nemo.thecrick.org/22121/lab/tree/home/users/jix/microSAM_finetuning/~/.conda/envs/microsam/lib/python3.10/site-packages/torch/utils/data/dataloader.py#line=707), in _BaseDataLoaderIter.__next__(self)
    705 if self._sampler_iter is None:
    706     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    707     self._reset()  # type: ignore[call-arg]
--> 708 data = self._next_data()
    709 self._num_yielded += 1
    710 if (
    711     self._dataset_kind == _DatasetKind.Iterable
    712     and self._IterableDataset_len_called is not None
    713     and self._num_yielded > self._IterableDataset_len_called
    714 ):

File [~/.conda/envs/microsam/lib/python3.10/site-packages/torch/utils/data/dataloader.py:764](https://ondemand.nemo.thecrick.org/node/ga100.nemo.thecrick.org/22121/lab/tree/home/users/jix/microSAM_finetuning/~/.conda/envs/microsam/lib/python3.10/site-packages/torch/utils/data/dataloader.py#line=763), in _SingleProcessDataLoaderIter._next_data(self)
    762 def _next_data(self):
    763     index = self._next_index()  # may raise StopIteration
--> 764     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    765     if self._pin_memory:
    766         data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)

File [~/.conda/envs/microsam/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py:52](https://ondemand.nemo.thecrick.org/node/ga100.nemo.thecrick.org/22121/lab/tree/home/users/jix/microSAM_finetuning/~/.conda/envs/microsam/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py#line=51), in _MapDatasetFetcher.fetch(self, possibly_batched_index)
     50         data = self.dataset.__getitems__(possibly_batched_index)
     51     else:
---> 52         data = [self.dataset[idx] for idx in possibly_batched_index]
     53 else:
     54     data = self.dataset[possibly_batched_index]

File [~/.conda/envs/microsam/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py:52](https://ondemand.nemo.thecrick.org/node/ga100.nemo.thecrick.org/22121/lab/tree/home/users/jix/microSAM_finetuning/~/.conda/envs/microsam/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py#line=51), in <listcomp>(.0)
     50         data = self.dataset.__getitems__(possibly_batched_index)
     51     else:
---> 52         data = [self.dataset[idx] for idx in possibly_batched_index]
     53 else:
     54     data = self.dataset[possibly_batched_index]

File [~/.conda/envs/microsam/lib/python3.10/site-packages/torch_em/data/segmentation_dataset.py:217](https://ondemand.nemo.thecrick.org/node/ga100.nemo.thecrick.org/22121/lab/tree/home/users/jix/microSAM_finetuning/~/.conda/envs/microsam/lib/python3.10/site-packages/torch_em/data/segmentation_dataset.py#line=216), in SegmentationDataset.__getitem__(self, index)
    216 def __getitem__(self, index):
--> 217     raw, labels = self._get_sample(index)
    218     initial_label_dtype = labels.dtype
    220     if self.raw_transform is not None:

File [~/.conda/envs/microsam/lib/python3.10/site-packages/torch_em/data/segmentation_dataset.py:203](https://ondemand.nemo.thecrick.org/node/ga100.nemo.thecrick.org/22121/lab/tree/home/users/jix/microSAM_finetuning/~/.conda/envs/microsam/lib/python3.10/site-packages/torch_em/data/segmentation_dataset.py#line=202), in SegmentationDataset._get_sample(self, index)
    201 # squeeze the singleton spatial axis if we have a spatial shape that is larger by one than self._ndim
    202 if self.patch_shape is not None and len(self.patch_shape) == self._ndim + 1:
--> 203     raw = raw.squeeze(1 if self._with_channels else 0)
    204     labels = labels.squeeze(1 if self._with_label_channels else 0)
    206 return raw, labels

ValueError: cannot select an axis to squeeze out which has size not equal to one

Could you provide an example of how to work with this type of data? Thanks!

[anwai98]

Hi @XiangruiJi,

Sorry for the late response.

A few quick mentions:

1. Re: model from BioImage.IO: The website is under active development, which is why I see an older model too.
   • Solution: The easiest way would be to download the micro-sam in a system with internet, and move the model checkpoints to the system without internet (I know this is some work, but the best way-to-go). I haven't got time to assort the model checkpoints, but the automatic downloading setup does the trick as expected.
2. Re: Adding rectification points doesn't help much: Ahha gotcha. Yeah the task is quite different than what micro-sam is familiar with. Let me know the progress of how it goes?
3. Re: Could you provide an example of how to work with this type of data?: Would it be possible to share the image and corresponding labels for your task? Then I can test it out and try to reproduce the error!

[anwai98]

Well I spot another thing actually:

patch_shape=(73, 512, 512)

micro-sam is an inherently 2d model (atleast at the training stage). This means, you need to train on 2d timepoints. You can simply switch to (1, 512, 512) as the patch shape and that should do the trick. But in any case, let me know if you can share the image and corresponding labels for me to test stuff!

[XiangruiJi]

Hi @anwai98 ,

Thank you for your reply. Now I can run the notebook with patch_shape=(1, 512, 512). However, I'm wondering whether I can use multiple .tif stack files (time series) in the image_dir folder for training. When I change raw_paths=image_paths[0] to raw_paths=image_paths (with raw_key, label_key= None, None) or raw_paths=image_dir (with raw_key, label_key = "*.tif", "*.tif" ), I run into errors. I have sent you my training data by email.