wholeBrainSeg_Large_UNEST_segmentation

[qwang0225]

HI, I'm trying to use wholeBrainSeg_Large_UNEST_segmentation from the model zoo, but encountered an error.

T1_image_name = os.path.join(data_path, "registered_T1.nii.gz")
data = {"image": T1_image_name}
# model_name = "wholeBrainSeg_Large_UNEST_segmentation"
# download(name=model_name, bundle_dir=".")
model_path = os.path.join(".", "wholeBrainSeg_Large_UNEST_segmentation", "models", "model.pt")
config_path = os.path.join(".", "wholeBrainSeg_Large_UNEST_segmentation", "configs", "inference.json")
# create a config instance to read the json file 
config = ConfigParser()
config.read_config(config_path)
preprocessing = config.get_parsed_content("preprocessing")
data = preprocessing(data)
model = config.get_parsed_content("network")

Here is the error message.

ModuleNotFoundError: Cannot locate class or function path: 'scripts.networks.unest_base_patch_4.UNesT'.

It seems that the issue arises from the network path in inference.json file does not exist.

   "network_def": {
        "_target_": "scripts.networks.unest_base_patch_4.UNesT",  
        "in_channels": 1,
        "out_channels": 133,
        "patch_size": 4,
        "depths": [
            2,
            2,
            8
        ],

However， it does exist, UNesT is the model class in the unest_base_patch_4.py

Could anyone help out? Thanks!

[tangy5]

Hi this issue might be because of the module path import errors.
you could double-check the file's path and location when executing the program...
Let me know if it still not working, I can help have a basic run test..

[qwang0225]

Hi, I did confirm my file’s path and location is correct, and it’s still not working. Could you help run a basic test? Thanks.

[KumoLiu]

Hi @qwang0225, did you try add the path of scripts?

sys.path.append(".../scripts")

[qwang0225]

Thanks! That worked!

[tangy5]

Thanks for help, @KumoLiu

[qwang0225]

Sorry, I got another issue. I'm trying to load the model.

model.load_state_dict(torch.load(model_path))

It seems there are missing keys.

RuntimeError                              Traceback (most recent call last)
Cell In[22], [line 1](vscode-notebook-cell:?execution_count=22&line=1)
----> [1](vscode-notebook-cell:?execution_count=22&line=1) model.load_state_dict(torch.load(model_path))

File [c:\Users\frank\anaconda3\envs\deep_learning\lib\site-packages\torch\nn\modules\module.py:2152](file:///C:/Users/frank/anaconda3/envs/deep_learning/lib/site-packages/torch/nn/modules/module.py:2152), in Module.load_state_dict(self, state_dict, strict, assign)
   [2147](file:///C:/Users/frank/anaconda3/envs/deep_learning/lib/site-packages/torch/nn/modules/module.py:2147)         error_msgs.insert(
   [2148](file:///C:/Users/frank/anaconda3/envs/deep_learning/lib/site-packages/torch/nn/modules/module.py:2148)             0, 'Missing key(s) in state_dict: {}. '.format(
   [2149](file:///C:/Users/frank/anaconda3/envs/deep_learning/lib/site-packages/torch/nn/modules/module.py:2149)                 ', '.join(f'"{k}"' for k in missing_keys)))
   [2151](file:///C:/Users/frank/anaconda3/envs/deep_learning/lib/site-packages/torch/nn/modules/module.py:2151) if len(error_msgs) > 0:
-> [2152](file:///C:/Users/frank/anaconda3/envs/deep_learning/lib/site-packages/torch/nn/modules/module.py:2152)     raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
   [2153](file:///C:/Users/frank/anaconda3/envs/deep_learning/lib/site-packages/torch/nn/modules/module.py:2153)                        self.__class__.__name__, "\n\t".join(error_msgs)))
   [2154](file:///C:/Users/frank/anaconda3/envs/deep_learning/lib/site-packages/torch/nn/modules/module.py:2154) return _IncompatibleKeys(missing_keys, unexpected_keys)

RuntimeError: Error(s) in loading state_dict for UNesT:
	Missing key(s) in state_dict: "nestViT.patch_embed.proj.weight", "nestViT.patch_embed.proj.bias", "nestViT.levels.0.pos_embed", "nestViT.levels.0.transformer_encoder.0.norm1.weight", "nestViT.levels.0.transformer_encoder.0.norm1.bias", "nestViT.levels.0.transformer_encoder.0.attn.qkv.weight", "nestViT.levels.0.transformer_encoder.0.attn.qkv.bias", "nestViT.levels.0.transformer_encoder.0.attn.proj.weight", "nestViT.levels.0.transformer_encoder.0.attn.proj.bias", "nestViT.levels.0.transformer_encoder.0.norm2.weight", "nestViT.levels.0.transformer_encoder.0.norm2.bias", "nestViT.levels.0.transformer_encoder.0.mlp.fc1.weight", "nestViT.levels.0.transformer_encoder.0.mlp.fc1.bias", "nestViT.levels.0.transformer_encoder.0.mlp.fc2.weight", "nestViT.levels.0.transformer_encoder.0.mlp.fc2.bias", "nestViT.levels.0.transformer_encoder.1.norm1.weight", "nestViT.levels.0.transformer_encoder.1.norm1.bias", "nestViT.levels.0.transformer_encoder.1.attn.qkv.weight", "nestViT.levels.0.transformer_encoder.1.attn.qkv.bias", "nestViT.levels.0.transformer_encoder.1.attn.proj.weight", "nestViT.levels.0.transformer_encoder.1.attn.proj.bias", "nestViT.levels.0.transformer_encoder.1.norm2.weight", "nestViT.levels.0.transformer_encoder.1.norm2.bias", "nestViT.levels.0.transformer_encoder.1.mlp.fc1.weight", "nestViT.levels.0.transformer_encoder.1.mlp.fc1.bias", "nestViT.levels.0.transformer_encoder.1.mlp.fc2.weight", "nestViT.levels.0.transformer_encoder.1.mlp.fc2.bias", "nestViT.levels.1.pos_embed", "nestViT.levels.1.pool.conv.weight", "nestViT.levels.1.pool.conv.bias", "nestViT.levels.1.pool.norm.weight", "nestViT.levels.1.pool.norm.bias", "nestViT.levels.1.transformer_encoder.0.norm1.weight", "nestViT.levels.1.transformer_encoder.0.norm1.bias", "nestViT.levels.1.transformer_encoder.0.attn.qkv.weight", "nestViT.levels.1.transformer_encoder.0.attn.qkv.bias", "nestViT.levels.1.transformer_encoder.0.attn.proj.weight", "nestViT.levels.1.transformer_encoder.0.attn.proj.bias", "nestViT.levels.1.transformer_encoder.0.norm2.weight", "nestViT.levels.1.transformer_encoder.0.norm2.bias", "nestViT.levels.1.transformer_encoder.0.mlp.fc1.weight", "nestViT.levels.1.transformer_encoder.0.mlp.fc1.bias", "nestViT.levels.1.transformer_encoder.0.mlp.fc2.weight", "nestViT.levels.1.transformer_encoder.0.mlp.fc2.bias", "nestViT.levels.1.transformer_encoder.1.norm1.weight", "nestViT.levels.1.transformer_encoder.1.norm1.bias", "nestViT.levels.1.transformer_encoder.1.attn.qkv.weight", "nestViT.levels.1.transformer_encoder.1.attn.qkv.bias", "nestViT.levels.1.transformer_encoder.1.attn.proj.weight", "nestViT.levels.1.transformer_encoder.1.attn.proj.bias", "nestViT.levels.1.transformer_encoder.1.norm2.weight", "nestViT.levels.1.transformer_encoder.1.norm2.bias", "nestViT.levels.1.transformer_encoder.1.mlp.fc1.weight", "nestViT.levels.1.transformer_encoder.1.mlp.fc1.bias", "nestViT.levels.1.transformer_encoder.1.mlp.fc2.weight", "nestViT.levels.1.transformer_encoder.1.mlp.fc2.bias", "nestViT.levels.2.pos_embed", "nestViT.levels.2.pool.conv.weight", "nestViT.levels.2.pool.conv.bias", "nestViT.levels.2.pool.norm.weight", "nestViT.levels.2.pool.norm.bias", "nestViT.levels.2.transformer_encoder.0.norm1.weight", "nestViT.levels.2.transformer_encoder.0.norm1.bias", "nestViT.levels.2.transformer_encoder.0.attn.qkv.weight", "nestViT.levels.2.transformer_encoder.0.attn.qkv.bias", "nestViT.levels.2.transformer_encoder.0.attn.proj.weight", "nestViT.levels.2.transformer_encoder.0.attn.proj.bias", "nestViT.levels.2.transformer_encoder.0.norm2.weight", "nestViT.levels.2.transformer_encoder.0.norm2.bias", "nestViT.levels.2.transformer_encoder.0.mlp.fc1.weight", "nestViT.levels.2.transformer_encoder.0.mlp.fc1.bias", "nestViT.levels.2.transformer_encoder.0.mlp.fc2.weight", "nestViT.levels.2.transformer_encoder.0.mlp.fc2.bias", "nestViT.levels.2.transformer_encoder.1.norm1.weight", "nestViT.levels.2.transformer_encoder.1.norm1.bias", "nestViT.levels.2.transformer_encoder.1.attn.qkv.weight", "nestViT.levels.2.transformer_encoder.1.attn.qkv.bias", "nestViT.levels.2.transformer_encoder.1.attn.proj.weight", "nestViT.levels.2.transformer_encoder.1.attn.proj.bias", "nestViT.levels.2.transformer_encoder.1.norm2.weight", "nestViT.levels.2.transformer_encoder.1.norm2.bias", "nestViT.levels.2.transformer_encoder.1.mlp.fc1.weight", "nestViT.levels.2.transformer_encoder.1.mlp.fc1.bias", "nestViT.levels.2.transformer_encoder.1.mlp.fc2.weight", "nestViT.levels.2.transformer_encoder.1.mlp.fc2.bias", "nestViT.levels.2.transformer_encoder.2.norm1.weight", "nestViT.levels.2.transformer_encoder.2.norm1.bias", "nestViT.levels.2.transformer_encoder.2.attn.qkv.weight", "nestViT.levels.2.transformer_encoder.2.attn.qkv.bias", "nestViT.levels.2.transformer_encoder.2.attn.proj.weight", "nestViT.levels.2.transformer_encoder.2.attn.proj.bias", "nestViT.levels.2.transformer_encoder.2.norm2.weight", "nestViT.levels.2.transformer_encoder.2.norm2.bias", "nestViT.levels.2.transformer_encoder.2.mlp.fc1.weight", "nestViT.levels.2.transformer_encoder.2.mlp.fc1.bias", "nestViT.levels.2.transformer_encoder.2.mlp.fc2.weight", "nestViT.levels.2.transformer_encoder.2.mlp.fc2.bias", "nestViT.levels.2.transformer_encoder.3.norm1.weight", "nestViT.levels.2.transformer_encoder.3.norm1.bias", "nestViT.levels.2.transformer_encoder.3.attn.qkv.weight", "nestViT.levels.2.transformer_encoder.3.attn.qkv.bias", "nestViT.levels.2.transformer_encoder.3.attn.proj.weight", "nestViT.levels.2.transformer_encoder.3.attn.proj.bias", "nestViT.levels.2.transformer_encoder.3.norm2.weight", "nestViT.levels.2.transformer_encoder.3.norm2.bias", "nestViT.levels.2.transformer_encoder.3.mlp.fc1.weight", "nestViT.levels.2.transformer_encoder.3.mlp.fc1.bias", "nestViT.levels.2.transformer_encoder.3.mlp.fc2.weight", "nestViT.levels.2.transformer_encoder.3.mlp.fc2.bias", "nestViT.levels.2.transformer_encoder.4.norm1.weight", "nestViT.levels.2.transformer_encoder.4.norm1.bias", "nestViT.levels.2.transformer_encoder.4.attn.qkv.weight", "nestViT.levels.2.transformer_encoder.4.attn.qkv.bias", "nestViT.levels.2.transformer_encoder.4.attn.proj.weight", "nestViT.levels.2.transformer_encoder.4.attn.proj.bias", "nestViT.levels.2.transformer_encoder.4.norm2.weight", "nestViT.levels.2.transformer_encoder.4.norm2.bias", "nestViT.levels.2.transformer_encoder.4.mlp.fc1.weight", "nestViT.levels.2.transformer_encoder.4.mlp.fc1.bias", "nestViT.levels.2.transformer_encoder.4.mlp.fc2.weight", "nestViT.levels.2.transformer_encoder.4.mlp.fc2.bias", "nestViT.levels.2.transformer_encoder.5.norm1.weight", "nestViT.levels.2.transformer_encoder.5.norm1.bias", "nestViT.levels.2.transformer_encoder.5.attn.qkv.weight", "nestViT.levels.2.transformer_encoder.5.attn.qkv.bias", "nestViT.levels.2.transformer_encoder.5.attn.proj.weight", "nestViT.levels.2.transformer_encoder.5.attn.proj.bias", "nestViT.levels.2.transformer_encoder.5.norm2.weight", "nestViT.levels.2.transformer_encoder.5.norm2.bias", "nestViT.levels.2.transformer_encoder.5.mlp.fc1.weight", "nestViT.levels.2.transformer_encoder.5.mlp.fc1.bias", "nestViT.levels.2.transformer_encoder.5.mlp.fc2.weight", "nestViT.levels.2.transformer_encoder.5.mlp.fc2.bias", "nestViT.levels.2.transformer_encoder.6.norm1.weight", "nestViT.levels.2.transformer_encoder.6.norm1.bias", "nestViT.levels.2.transformer_encoder.6.attn.qkv.weight", "nestViT.levels.2.transformer_encoder.6.attn.qkv.bias", "nestViT.levels.2.transformer_encoder.6.attn.proj.weight", "nestViT.levels.2.transformer_encoder.6.attn.proj.bias", "nestViT.levels.2.transformer_encoder.6.norm2.weight", "nestViT.levels.2.transformer_encoder.6.norm2.bias", "nestViT.levels.2.transformer_encoder.6.mlp.fc1.weight", "nestViT.levels.2.transformer_encoder.6.mlp.fc1.bias", "nestViT.levels.2.transformer_encoder.6.mlp.fc2.weight", "nestViT.levels.2.transformer_encoder.6.mlp.fc2.bias", "nestViT.levels.2.transformer_encoder.7.norm1.weight", "nestViT.levels.2.transformer_encoder.7.norm1.bias", "nestViT.levels.2.transformer_encoder.7.attn.qkv.weight", "nestViT.levels.2.transformer_encoder.7.attn.qkv.bias", "nestViT.levels.2.transformer_encoder.7.attn.proj.weight", "nestViT.levels.2.transformer_encoder.7.attn.proj.bias", "nestViT.levels.2.transformer_encoder.7.norm2.weight", "nestViT.levels.2.transformer_encoder.7.norm2.bias", "nestViT.levels.2.transformer_encoder.7.mlp.fc1.weight", "nestViT.levels.2.transformer_encoder.7.mlp.fc1.bias", "nestViT.levels.2.transformer_encoder.7.mlp.fc2.weight", "nestViT.levels.2.transformer_encoder.7.mlp.fc2.bias", "nestViT.norm.weight", "nestViT.norm.bias", "encoder1.layer.conv1.conv.weight", "encoder1.layer.conv2.conv.weight", "encoder1.layer.conv3.conv.weight", "encoder2.transp_conv_init.conv.weight", "encoder2.blocks.0.conv.weight", "encoder3.layer.conv1.conv.weight", "encoder3.layer.conv2.conv.weight", "encoder4.layer.conv1.conv.weight", "encoder4.layer.conv2.conv.weight", "decoder5.transp_conv.conv.weight", "decoder5.conv_block.conv1.conv.weight", "decoder5.conv_block.conv2.conv.weight", "decoder5.conv_block.conv3.conv.weight", "decoder4.transp_conv.conv.weight", "decoder4.conv_block.conv1.conv.weight", "decoder4.conv_block.conv2.conv.weight", "decoder4.conv_block.conv3.conv.weight", "decoder3.transp_conv.conv.weight", "decoder3.conv_block.conv1.conv.weight", "decoder3.conv_block.conv2.conv.weight", "decoder3.conv_block.conv3.conv.weight", "decoder2.transp_conv.conv.weight", "decoder2.conv_block.conv1.conv.weight", "decoder2.conv_block.conv2.conv.weight", "decoder2.conv_block.conv3.conv.weight", "decoder1.transp_conv.conv.weight", "decoder1.conv_block.conv1.conv.weight", "decoder1.conv_block.conv2.conv.weight", "decoder1.conv_block.conv3.conv.weight", "encoder10.conv.weight", "encoder10.conv.bias", "encoder10.adn.A.weight", "out.conv.conv.weight", "out.conv.conv.bias". 
	Unexpected key(s) in state_dict: "model".

[tangy5]

There is something wrong with either the version of python, monai, or other packages or pytorch.
The error shows the keyerror. You could debug a bit, e.g., change the key names.

[qwang0225]

Thanks! I was able to fix it by changing the key names.

But somehow the model perform very poorly on a sample.

Is there anything I missed?

import os
import sys
import numpy as np
import torch
import pydicom
import nibabel as nib
import matplotlib.pyplot as plt
from monai.bundle import ConfigParser, download
from monai.transforms import (
    LoadImage, 
    LoadImaged, 
    Orientation, 
    Orientationd, 
    EnsureChannelFirst, 
    EnsureChannelFirstd, 
    Compose,
    NormalizeIntensityd,
    EnsureTyped,
)
import json
import ants

# register images to MNI space
mri_image = ants.image_read(os.path.join(data_path, "BRAINIX_NIFTI_T1.nii.gz"))
mni_template = ants.image_read(ants.get_ants_data('mni')).resample_image_to_target(mri_image)
affine_transform = ants.registration(fixed=mni_template, moving=mri_image, type_of_transform='Affine')
mni_registered_image = ants.apply_transforms(fixed=mni_template, moving=mri_image, transformlist=affine_transform['fwdtransforms'][0])
ants.image_write(mni_registered_image, os.path.join(data_path, "registered_T1.nii.gz"))

# get original image 
T1_image_name = os.path.join(data_path, "registered_T1.nii.gz")
original_data = {"image": T1_image_name}
model_path = os.path.join(".", "wholeBrainSeg_Large_UNEST_segmentation", "models", "model.pt")
config_path = os.path.join(".", "wholeBrainSeg_Large_UNEST_segmentation", "configs", "inference.json")

# create a config instance 
config = ConfigParser()
config.read_config(config_path)

#preprocessing 
preprocessing = config.get_parsed_content("preprocessing")
data = preprocessing(original_data)

# Load the state dictionary into the model
sys.path.append(os.path.join(".", "wholeBrainSeg_Large_UNEST_segmentation"))
model = config.get_parsed_content("network")
checkpoint = torch.load(model_path)
mapping = {
    "old_key": "new_key",
}
state_dict = {mapping.get(k, k): v for k, v in checkpoint['model'].items()}
model.load_state_dict(state_dict)

# Infer from model 
inferer = config.get_parsed_content("inferer")
with torch.no_grad():
    data["pred"] = inferer(data["image"].unsqueeze(0).to(device), network=model)

# remove batch dimension
data["pred"] = data["pred"][0]
data["image"] = data["image"][0]

# postprocessing
postprocessing = config.get_parsed_content("postprocessing")
data = postprocessing(data)
segmentation = data["pred"][0]
segmentation = segmentation.cpu().numpy()

# print 
for i in range(22):
    plt.subplots(1,2,figsize=(16, 6))
    plt.subplot(121)
    # title = "T1 Image"
    plt.title(f"T1 image {i}")
    plt.pcolormesh(T1_images[:, :, i].T, cmap="gray")
    plt.axis("off")
    plt.subplot(122)
    plt.title(f"Segmentation {i}")
    plt.pcolormesh(segmentation[:,:, i].T, cmap="nipy_spectral")
    plt.axis("off")
    plt.show()