yhygao / CBIM-Medical-Image-Segmentation

A PyTorch framework for medical image segmentation
Apache License 2.0
260 stars 46 forks source link

The Inference Code #16

Closed dzzhang96 closed 1 year ago

dzzhang96 commented 1 year ago

Hi yhygao! Thank you for the codes first!

I have a dataset for inference and it was pre-processed and was not post-processed in the same way as my training dataset. Can I ask how to use the model to run inference for my dataset for inference and output the nii.gz file? It would be really helpful if you can update the codes.

Thanks again!

yhygao commented 1 year ago

I uploaded a prediction.py in the dev branch. It provides an example of how to make inference and save the predicted label map into a nii.gz file. it's not fully tested, but you can modify it to fit your case. You need to modify the preprocess steps and target_spacing in the 'preprocess' function to make sure the testing data is processed the same as the training data. My code will do the other steps like resample to the target spacing, ensembled prediction, resample back to the original spacing, and save nii.gz file.

dzzhang96 commented 1 year ago

Awesome! Big thank you for the quick response. I'll test it and keep updated

dzzhang96 commented 1 year ago

Hi @yhygao , I have tested the prediction.py. I have modified the preprocessing steps and target_spacing in the 'preprocess' function. I got the label files saved as nii.gz file, but the result is blank without any outputs, only showing the black background.

I tried to comment lines 92-96, then, the result is exactly what I need, except that the label image size is not correct. Do you have any clues? Thank you!

https://github.com/yhygao/CBIM-Medical-Image-Segmentation/blob/4c21adee784224ebe0c8561781650492c6ec4095/prediction.py#L94

yhygao commented 1 year ago

This function is used to resample the predicted label map to the original (the testing image) spacing. If you don't use the dataset conversion code in my framework, it might be because the testing data and the predicted label map have different origin corrdinates. You can try to add the following code between line 93 and 94: itk_pred.SetOrigin(itk_img.GetOrigin()) itk_pred.SetDirection(itk_img.GetDirection())

These commands will set the correct origin and direction for the predicted label. And then the ResampleLabelToRef will resample the predicted label to have the same spacing as the original testing image, i.e. resample to the correct size.

Also, make sure you set the target_spacing to be consistent with the training spacing setting.

dzzhang96 commented 1 year ago

This works perfectly!