File size: 5,000 Bytes
7f24887 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
from cvseg_utils import *
import warnings
warnings.filterwarnings("ignore", message="GLCM is symmetrical, therefore Sum Average = 2 * Joint Average")
import os
import logging
from datetime import datetime
def seg_main():
parser = argparse.ArgumentParser(description='Nodule segmentation and feature extraction from CT images.')
parser.add_argument('--raw_data_path', type=str, required=True, help='Path to raw CT images')
parser.add_argument('--csv_save_path', type=str, required=True, help='Path to save the CSV files')
parser.add_argument('--dataset_csv', type=str, required=True, help='Path to the dataset CSV')
parser.add_argument('--nifti_clm_name', type=str, required=True, help='name to the nifti column name')
parser.add_argument('--unique_Annotation_id', type=str, help='Column for unique annotation ID')
parser.add_argument('--Malignant_lbl', type=str, required=True, help='Column name for malignancy labels')
parser.add_argument('--coordX', type=str, required=True, help='Column name for X coordinate')
parser.add_argument('--coordY', type=str, required=True, help='Column name for Y coordinate')
parser.add_argument('--coordZ', type=str, required=True, help='Column name for Z coordinate')
parser.add_argument('--w', type=str, required=True, help='Column name for width')
parser.add_argument('--h', type=str, required=True, help='Column name for height')
parser.add_argument('--d', type=str, required=True, help='Column name for depth')
parser.add_argument('--seg_alg', type=str, default='gmm', choices=['gmm', 'knn', 'fcm', 'otsu'], help='Segmentation algorithm to use')
parser.add_argument('--dataset_name', type=str, default='DLCS24', help='Dataset to use')
parser.add_argument('--expansion_mm', type=float, default=1.0, help='Expansion in mm')
parser.add_argument('--use_expand', action='store_true', help='Use expansion if set')
parser.add_argument('--params_json', type=str, required=True, help="Path to JSON file with radiomics parameters")
parser.add_argument('--save_the_generated_mask', action='store_true', help='Save generated segmentation mask')
parser.add_argument('--save_nifti_path', type=str, help='Path to save the nifti files')
args = parser.parse_args()
df = pd.read_csv(args.dataset_csv)
final_dect = df[args.nifti_clm_name].unique()
for dictonary_list_i, ct_filename in enumerate(final_dect):
try:
filtered_df = df[df[args.nifti_clm_name] == ct_filename].reset_index()
ct_nifti_path = os.path.join(args.raw_data_path, ct_filename)
ct_image = sitk.ReadImage(ct_nifti_path)
ct_array = sitk.GetArrayFromImage(ct_image)
spacing = ct_image.GetSpacing()
full_mask_array = np.zeros_like(ct_array, dtype=np.uint8)
for idx, row in filtered_df.iterrows():
worldCoord = np.asarray([row[args.coordX], row[args.coordY], row[args.coordZ]])
voxelCoord = ct_image.TransformPhysicalPointToIndex(worldCoord.tolist())
spacing = ct_image.GetSpacing()
w = int(row[args.w] / spacing[0])
h = int(row[args.h] / spacing[1])
d = int(row[args.d] / spacing[2])
bbox_center = [voxelCoord[2], voxelCoord[1], voxelCoord[0]]
bbox_whd = [d, h, w]
if args.seg_alg == 'gmm':
mask_image_array = segment_nodule_gmm(ct_array, bbox_center, bbox_whd)
elif args.seg_alg == 'knn':
mask_image_array = segment_nodule_kmeans(ct_array, bbox_center, bbox_whd)
elif args.seg_alg == 'fcm':
mask_image_array = segment_nodule_fcm(ct_array, bbox_center, bbox_whd)
elif args.seg_alg == 'otsu':
mask_image_array = segment_nodule_otsu(ct_array, bbox_center, bbox_whd)
if args.use_expand:
mask_image_array = expand_mask_by_distance(mask_image_array, spacing=spacing, expansion_mm=args.expansion_mm)
#print("Segmented mask sum:", np.sum(mask_image_array))
full_mask_array[mask_image_array==1] = 1
if args.save_the_generated_mask:
print("Segmented mask sum:", np.sum(full_mask_array))
combined_mask_image = sitk.GetImageFromArray(full_mask_array)
combined_mask_image.SetSpacing(ct_image.GetSpacing())
combined_mask_image.SetDirection(ct_image.GetDirection())
combined_mask_image.SetOrigin(ct_image.GetOrigin())
mask_filename = ct_filename.split('.nii')[0]+"_mask.nii.gz"
sitk.WriteImage(combined_mask_image, os.path.join(args.save_nifti_path, mask_filename))
print(f"Saved {mask_filename}")
except Exception as e:
print(f"Error processing {ct_filename}: {e}")
if __name__ == "__main__":
seg_main()
|