Predicting discrete pixel output

[KartikeyKansal1]

Hi @benjijamorris
I have paired set of training images, one which is a microscopic image of a cell and the other has cell areas with pixel values only as 0,1,2, and 3. Looks something like below.

I essentially wanna predict the pixel values out of (0,1,2 and 3) based off the microscopic image input only. Is there any way to do this in cyto-dl? I've given the labelfree and segmentation a try but they don't produce expected outputs.

[benjijamorris]

Good question! I would use the segmentation configs, but there are a few things I'd recommend changing:
Data:

• Remove the thresholdIntensity transform since the pixel values in your target are meaningful (not just [0,255])
  Model:
• Set your backbone.out_channels to 4 (the number of classes you have). This way, you're generating per-class probabilities with your model, one per channel
• Change the loss function arguments, in particular setting to_onehot_y: True (This will turn your 1-channel image with class-labeled pixels to a 4-channel image where each channel is 1 within that class and 0 outside of it), and replacing sigmoid:True with softmax: True. This will encourage your model's class predictions to be mutually exclusive and sum to 1.
  Let me know how this goes and if you have any other questions!

[KartikeyKansal1]

Hi @benjijamorris , thank you for the suggestions!

I'm getting the following error when I changed out_channels=3:

[2024-04-17 13:32:24,228][cyto_dl.utils.template_utils][INFO] - Closing loggers...
Error executing job with overrides: ['experiment=im2im/segmentation.yaml', 'trainer=cpu']
Traceback (most recent call last):
  File "/Users/kartikeykansal/Documents/tensionGAN/actin/Segmentation_240415/240417/Experiment_1/cyto-dl/cyto_dl/train.py", line 143, in main
    metric_dict, _ = train(cfg)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/cyto_dl/utils/template_utils.py", line 56, in wrap
    raise ex
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/cyto_dl/utils/template_utils.py", line 53, in wrap
    out = task_func(cfg=cfg)
  File "/Users/kartikeykansal/Documents/tensionGAN/actin/Segmentation_240415/240417/Experiment_1/cyto-dl/cyto_dl/train.py", line 98, in train
    trainer.fit(model=model, datamodule=data, ckpt_path=cfg.get("ckpt_path"))
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py", line 532, in fit
    call._call_and_handle_interrupt(
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/lightning/pytorch/trainer/call.py", line 43, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py", line 571, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py", line 980, in _run
    results = self._run_stage()
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py", line 1021, in _run_stage
    self._run_sanity_check()
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py", line 1050, in _run_sanity_check
    val_loop.run()
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/lightning/pytorch/loops/utilities.py", line 181, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 115, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 376, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_kwargs.values())
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/lightning/pytorch/trainer/call.py", line 294, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/lightning/pytorch/strategies/strategy.py", line 393, in validation_step
    return self.model.validation_step(*args, **kwargs)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/cyto_dl/models/base_model.py", line 181, in validation_step
    loss, preds, targets = self.model_step("val", batch, batch_idx)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/cyto_dl/models/im2im/multi_task.py", line 212, in model_step
    outs = self.run_forward(batch, stage, save_image, run_heads)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/cyto_dl/models/im2im/multi_task.py", line 145, in run_forward
    return self._train_forward(batch, stage, save_image, run_heads)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/cyto_dl/models/im2im/multi_task.py", line 109, in _train_forward
    return {
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/cyto_dl/models/im2im/multi_task.py", line 110, in <dictcomp>
    task: self.task_heads[task].run_head(z, batch, stage, save_image) for task in run_heads
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/cyto_dl/nn/head/base_head.py", line 99, in run_head
    loss = self._calculate_loss(y_hat, batch[self.head_name])
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/cyto_dl/nn/head/base_head.py", line 41, in _calculate_loss
    return self.loss(y_hat, y)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/monai/losses/dice.py", line 790, in forward
    dice_loss = self.dice(input, target)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/monai/losses/dice.py", line 158, in forward
    target = one_hot(target, num_classes=n_pred_ch)
  File "/Users/kartikeykansal/miniconda3/envs/cytodlenv_240417/lib/python3.9/site-packages/monai/networks/utils.py", line 192, in one_hot
    labels = o.scatter_(dim=dim, index=labels.long(), value=1)
RuntimeError: index -1 is out of bounds for dimension 1 with size 3

I also wanna make sure that are you referring class labels to the pixel intensity values 0,1,2,3 itself?

[benjijamorris]

Correct - I was assuming that you're trying to classify each pixel in the image into 0,1,2,3? I think the issue is you are doing 4 class classification and your output channels is only 3. If the loss is being dominated by your background class (0), you can set ignore_background: True in the loss function arguments.

[KartikeyKansal1]

I am getting the 4 channels in the out_channels. The error was probably due to some of my images having pixel values as -1.

But the output is still a range of pixels and not as expected. I've removed the ThresholdIntensity as well. Do you recommend giving the input 4 channels as well?

Also it says that ignore_background is an unexpected argument

[benjijamorris]

Each channel will have the probability of the pixel belonging to the class with the index of the channel you're looking at, if you want a classification for each pixel, you can find the channel-wise argmax for each pixel.

The input channels is how many channels are in the image you're passing into the network - only change this if you're changing the types of images you're passing in.

[KartikeyKansal1]

Yes I'm looking for a classification per pixel. I am hoping to generate an image something like my target class labels:

The different labels can be on different channels though

Can you help me point towards the channel-wise argmax.

[benjijamorris]

from aicsimageio import imread
import numpy as np

prediction = imread(your_path.tif).squeeze() # CZYX image
class_predictions = np.argmax(prediction, axis=0)