Feeding processed image into monai zoo model for inference in python script

[pzaffino]

Hi all,
I would like to use Wholebrainseg large unest segmentation (from the monai zoo model) in my python code.
I implemented the preprocessing transforms according to what I see in the repository (including registration on atlas), so my image is ready for the inference.
I also loaded the pt model.
At this point my question is: how to feed the image into the model to get the segmentation?
Something like
out = model(img)
doesn't work.

Thanks a lot.

Paolo

[KumoLiu]

Hi @pzaffino, could you please share the error log? Thanks!

[pzaffino]

Hi @KumoLiu
thanks a lot for answering.

Playing aroud with the code I downloaded from the model zoo I get an error about the model loading.

The code is this:

import torch
import monai
from networks.unest import UNesT

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = UNesT(in_channels=1, out_channels=133)

ckpt = torch.load('UNesT_wholebrainseg_base.pt', map_location='cpu')
model.load_state_dict(ckpt['state_dict'], strict=True)
model.to(device)

but I get:

RuntimeError                              Traceback (most recent call last)
Cell In[11], line 6
      3 model = UNesT(in_channels=1, out_channels=133)
      5 ckpt = torch.load('UNesT_wholebrainseg_base.pt', map_location='cpu')
----> 6 model.load_state_dict(ckpt['state_dict'], strict=True)
      7 model.to(device)
      9 model.eval()

File ~/work/python_venv/pytorch/lib/python3.11/site-packages/torch/nn/modules/module.py:2041, in Module.load_state_dict(self, state_dict, strict)
   2036         error_msgs.insert(
   2037             0, 'Missing key(s) in state_dict: {}. '.format(
   2038                 ', '.join('"{}"'.format(k) for k in missing_keys)))
   2040 if len(error_msgs) > 0:
-> 2041     raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
   2042                        self.__class__.__name__, "\n\t".join(error_msgs)))
   2043 return _IncompatibleKeys(missing_keys, unexpected_keys)

RuntimeError: Error(s) in loading state_dict for UNesT:
	Unexpected key(s) in state_dict: "encoder3.layer.conv3.conv.weight", "encoder4.layer.conv3.conv.weight". 

Thanks a lot!

[KumoLiu]

Hi @pzaffino, I would suggest you refer to the "Utilizing Pretrained Weights from the Bundle" part in this tutorial.

Hope it can help you, thanks!