| import os |
| import numpy as np |
| import SimpleITK as sitk |
|
|
| |
| def crop_z_axis(data, group_size=30, keep_layers=9): |
|
|
| z_size = data.shape[0] |
| total_layers = z_size // group_size |
|
|
| if total_layers <= keep_layers: |
| return data |
|
|
| remove_layers = total_layers - keep_layers |
| remove_front = (remove_layers // 2) * group_size |
| remove_back = (remove_layers - remove_layers // 2) * group_size |
|
|
| cropped = data[remove_front:z_size-remove_back, :, :] |
|
|
| return cropped |
|
|
|
|
| def resample_spacing(Image,Label,NewSpacing): |
| Size=Image.GetSize() |
| Spacing=Image.GetSpacing() |
| Origin = Image.GetOrigin() |
| Direction = Image.GetDirection() |
|
|
| NewSize=[int(Size[0]*Spacing[0]/NewSpacing[0]),int(Size[1]*Spacing[1]/NewSpacing[1]),int(Size[2]*Spacing[2]/NewSpacing[2])] |
| Resample = sitk.ResampleImageFilter() |
| Resample.SetOutputDirection(Direction) |
| Resample.SetOutputOrigin(Origin) |
| Resample.SetSize(NewSize) |
| Resample.SetInterpolator(sitk.sitkBSpline) |
| Resample.SetOutputSpacing(NewSpacing) |
| NewImage = Resample.Execute(Image) |
|
|
| Resample = sitk.ResampleImageFilter() |
| Resample.SetOutputDirection(Direction) |
| Resample.SetOutputOrigin(Origin) |
| Resample.SetSize(NewSize) |
| Resample.SetOutputSpacing(NewSpacing) |
| Resample.SetInterpolator(sitk.sitkNearestNeighbor) |
| NewLabel = Resample.Execute(Label) |
|
|
| return NewImage,NewLabel |
|
|
|
|
| |
| def crop_xy_axis(image, mask, crop_size=(200, 200)): |
| if np.sum(mask) == 0: |
| return np.zeros((image.shape[0], crop_size[0], crop_size[1]), dtype=image.dtype), np.zeros((mask.shape[0], crop_size[0], crop_size[1]), dtype=mask.dtype) |
| |
| non_zero_coords = np.argwhere(mask > 0) |
| z_min, y_min, x_min = non_zero_coords.min(axis=0) |
| z_max, y_max, x_max = non_zero_coords.max(axis=0) |
|
|
| center_y = (y_min + y_max) // 2 |
| center_x = (x_min + x_max) // 2 |
|
|
| crop_height, crop_width = crop_size |
| y_start = max(center_y - crop_height // 2, 0) |
| y_end = min(center_y + crop_height // 2, image.shape[1]) |
| x_start = max(center_x - crop_width // 2, 0) |
| x_end = min(center_x + crop_width // 2, image.shape[2]) |
|
|
| cropped_image = image[:, y_start:y_end, x_start:x_end] |
| cropped_mask = mask[:, y_start:y_end, x_start:x_end] |
|
|
| pad_y_before = max(0, crop_height // 2 - center_y) |
| pad_y_after = max(0, crop_height // 2 - (image.shape[1] - center_y)) |
| pad_x_before = max(0, crop_width // 2 - center_x) |
| pad_x_after = max(0, crop_width // 2 - (image.shape[2] - center_x)) |
|
|
| cropped_image = np.pad( |
| cropped_image, |
| ((0, 0), (pad_y_before, pad_y_after), (pad_x_before, pad_x_after)), |
| mode='constant', |
| constant_values=0 |
| ) |
|
|
| cropped_mask = np.pad( |
| cropped_mask, |
| ((0, 0), (pad_y_before, pad_y_after), (pad_x_before, pad_x_after)), |
| mode='constant', |
| constant_values=0 |
| ) |
|
|
| return cropped_image, cropped_mask |
|
|
|
|
| def preprocess_patient(img_path, seg_path, save_img_path, save_seg_path, target_spacing=(0.9375, 0.9375, 0.2451)): |
|
|
| img = sitk.ReadImage(img_path) |
| seg = sitk.ReadImage(seg_path) |
|
|
| img_arr = sitk.GetArrayFromImage(img) |
| seg_arr = sitk.GetArrayFromImage(seg) |
|
|
| |
| cropped_img_arr = crop_z_axis(img_arr) |
| cropped_seg_arr = crop_z_axis(seg_arr) |
|
|
| cropped_img = sitk.GetImageFromArray(cropped_img_arr) |
| cropped_img.SetDirection(img.GetDirection()) |
| cropped_img.SetOrigin(img.GetOrigin()) |
| cropped_img.SetSpacing(img.GetSpacing()) |
|
|
| cropped_seg = sitk.GetImageFromArray(cropped_seg_arr) |
| cropped_seg.SetDirection(img.GetDirection()) |
| cropped_seg.SetOrigin(img.GetOrigin()) |
| cropped_seg.SetSpacing(img.GetSpacing()) |
|
|
| |
| resample_img, resample_seg = resample_spacing(cropped_img, cropped_seg, target_spacing) |
|
|
| |
| resample_img_arr = sitk.GetArrayFromImage(resample_img) |
| resample_seg_arr = sitk.GetArrayFromImage(resample_seg) |
|
|
| output_img_arr, output_seg_arr = crop_xy_axis(resample_img_arr, resample_seg_arr, crop_size=(200,200)) |
|
|
| |
| output_img = sitk.GetImageFromArray(output_img_arr) |
| output_seg = sitk.GetImageFromArray(output_seg_arr) |
|
|
| output_img.SetDirection(resample_img.GetDirection()) |
| output_img.SetOrigin(resample_img.GetOrigin()) |
| output_img.SetSpacing(resample_img.GetSpacing()) |
|
|
| output_seg.SetDirection(resample_img.GetDirection()) |
| output_seg.SetOrigin(resample_img.GetOrigin()) |
| output_seg.SetSpacing(resample_img.GetSpacing()) |
|
|
| sitk.WriteImage(output_img, save_img_path) |
| sitk.WriteImage(output_seg, save_seg_path) |
|
|
|
|
| def run_preprocessing(): |
|
|
| root="./ACDC4D" |
|
|
| img_dir=os.path.join(root,"img") |
| seg_dir=os.path.join(root,"seg") |
|
|
| output_dir=os.path.join("./processed") |
|
|
| os.makedirs(output_dir,exist_ok=True) |
|
|
|
|
| files=[f for f in os.listdir(img_dir) if f.endswith("_3d.nii.gz")] |
|
|
| for fname in files: |
| pid=fname.split("_")[0] |
| print(f"Processing {pid}") |
|
|
| patient_dir=os.path.join(output_dir, pid) |
| os.makedirs(patient_dir, exist_ok=True) |
|
|
| img_path=os.path.join(img_dir, fname) |
| seg_path=os.path.join(seg_dir, fname) |
|
|
| save_img=os.path.join(patient_dir, f"{pid}_sa_image.nii.gz") |
| save_seg=os.path.join(patient_dir, f"{pid}_sa_image-seg.nii.gz") |
|
|
| preprocess_patient(img_path, seg_path, save_img, save_seg) |
|
|
| print("All preprocessing finished.") |
|
|
|
|
| if __name__=="__main__": |
|
|
| run_preprocessing() |
