Hi ITK community,
I have a 3D PET/CT dataset with 3D segmentations (GT) (with labels 0, 1, and 2). The CT images are typically 512x512xSizeZ, while the PET images are of various shapes like 128x128xSizeZ' or `256x256xSizeZ’, etc. All the segmentations (GT) are made on CT so they have the same Direction, Size, Origin, and Spacing as CT images. I am trying to train a 2D UNet model which uses 2D slices of these images and but my goal is to stack all the 2D predictions and finally get the 3D predicted mask in the same coordinates as the original CT/Mask (Direction. Spacing, Size, Origin). I am doing the following preprocessing:
(1) Using sitk.Resample(ctimg, ptimg, interpolator=sitk.sitkLinear) and sitk.Resample(gtimg, ptimg, interpolator=sitk.sitkLinear) to resample 3D volumes of CT and GT to PET coordinates.
(2) Now I save all the individual 2D axial slices of 3D PET, CT, and GT separately as slices that contain lesions and slices that don’t contain lesions (this information is present in the selected_slices variable). At this point, while saving 2D slices, I don’t put in any information for spacing, direction, origin, etc (is that correct?). For this I use:
def save_selected_slices(array3d, patientID, slicetype, selected_slices, dest_folder):
for i in range(len(selected_slices)):
array2d = array3d[:,:,selected_slices[i]]
img = nib.Nifti1Image(array2d, affine=np.eye(4))
save_filename = os.path.join(dest_folder, patientID + '_' + convert_to_3_digits(selected_slices[i]) + '_ax_' + slicetype + '.nii.gz' )
nib.save(img, save_filename)
(3) Before giving them as input to the model, I am resizing all the 2D slices to (128x128). The input to the model looks like this: x = [N, 2, 128, 128], where N= batchsize and 2 channels refer to the first channel being the PET slice and the second being the CT slice. The target looks like this: y = [N, 128, 128]. I save all the 2D predictions after resizing them back to their original PET sizes as 2D nifti images using:
def save_as_nifti(pred, savepath):
image = nib.Nifti1Image(pred, affine=np.eye(4))
nib.save(image, savepath)
(4) Once all the predictions on all the test slices have been made and saved, I want to stack them per patient and save them as 3D nifti images in the same coordinates as the original CT/GT images. I create a dictionary gt3dresampled_coords with patientID as the key containing coordinates information and use it to assign direction, origin, spacing, size to the axially-stacked 2d prediction. After this I find that, for the new 3D stacked images, only the size and direction as the same. The spacing is still [1,1,1] and the origin still remains [0, 0, 0]. Could you tell me what is the mistake in this whole procedure given below?
PatientIDs = []
Pred3D_dict = {}
for path in pred2d_paths:
patientID = os.path.basename(path).split('_')[0]
print(patientID)
PatientIDs.append(patientID)
Unique_PatientIDs = np.unique(PatientIDs)
for i in range(len(Unique_PatientIDs)):
Pred3D_dict[Unique_PatientIDs[i]] = []
for path in pred2d_paths:
patientID = os.path.basename(path).split('_')[0]
array2d, vox = utils.nib2numpy(path)
Pred3D_dict[patientID].append(array2d)
resampler1 = sitk.ResampleImageFilter()
resampler1.SetInterpolator(sitk.sitkNearestNeighbor)
for ptid in Unique_PatientIDs:
pred3d = np.stack(Pred3D_dict[ptid], axis=-1)
pred3d = np.transpose(pred3d, (2, 1, 0))
patientdict = gt3dresampled_coords[ptid]
resampled_origin = patientdict['Origin']
resampled_size = patientdict['Size']
resampled_spacing = patientdict['Spacing']
resampled_direction = patientdict['Direction']
pred3d_img = sitk.GetImageFromArray(pred3d)
resampler1.SetOutputOrigin(resampled_origin)
resampler1.SetSize(resampled_size)
resampler1.SetOutputSpacing(resampled_spacing)
resampler1.SetOutputDirection(resampled_direction)
resampled_pred3d = resampler1.Execute(pred3d_img)
pred3dsavepath = os.path.join(savedir_pred3d, ptid+'.nii.gz')
sitk.WriteImage(pred3d_img, pred3dsavepath)
print('saving 3d: ' + ptid)