def __init__(self, **kwargs: Any) -> None: fg_classes = ["spinalcord", "lung_r", "lung_l", "heart", "esophagus"] fg_display_names = [ "SpinalCord", "Lung_R", "Lung_L", "Heart", "Esophagus" ] super().__init__( architecture="UNet3D", feature_channels=[32], kernel_size=3, azure_dataset_id=AZURE_DATASET_ID, crop_size=(64, 224, 224), test_crop_size=(128, 512, 512), image_channels=["ct"], ground_truth_ids=fg_classes, ground_truth_ids_display_names=fg_display_names, colours=[(255, 255, 255)] * len(fg_classes), fill_holes=[False] * len(fg_classes), largest_connected_component_foreground_classes=[ "lung_r", "lung_l", "heart" ], num_dataload_workers=8, norm_method=PhotometricNormalizationMethod.CtWindow, level=40, window=400, class_weights=equally_weighted_classes(fg_classes, background_weight=0.02), train_batch_size=8, inference_batch_size=1, inference_stride_size=(64, 256, 256), start_epoch=0, num_epochs=140, l_rate=1e-3, min_l_rate=1e-5, l_rate_polynomial_gamma=0.9, optimizer_type=OptimizerType.Adam, opt_eps=1e-4, adam_betas=(0.9, 0.999), momentum=0.9, weight_decay=1e-4, save_start_epoch=100, save_step_epochs=20, test_start_epoch=140, use_mixed_precision=True, use_model_parallel=True, monitoring_interval_seconds=0, test_diff_epochs=1, test_step_epochs=1, loss_type=SegmentationLoss.Mixture, mixture_loss_components=[ MixtureLossComponent(0.5, SegmentationLoss.Focal, 0.2), MixtureLossComponent(0.5, SegmentationLoss.SoftDice, 0.1) ], ) self.add_and_validate(kwargs)
def test_construct_loss_function() -> None: model_config = DummyModel() model_config.loss_type = SegmentationLoss.Mixture # Weights deliberately do not sum to 1.0. weights = [1.5, 0.5] model_config.mixture_loss_components = [ MixtureLossComponent(weights[0], SegmentationLoss.CrossEntropy, 0.2), MixtureLossComponent(weights[1], SegmentationLoss.SoftDice, 0.1)] loss_fn = ModelTrainingStepsForSegmentation.construct_loss_function(model_config) assert isinstance(loss_fn, MixtureLoss) assert len(loss_fn.components) == len(weights) assert loss_fn.components[0][0] == weights[0] / sum(weights) assert loss_fn.components[1][0] == weights[1] / sum(weights)
def __init__(self, num_structures: int = 0, **kwargs: Any) -> None: """ :param num_structures: number of structures from STRUCTURE_LIST to predict (default: all structures) :param kwargs: other args from subclass """ # Number of training epochs num_epochs = 120 # Number of structures to predict; if positive but less than the length of STRUCTURE_LIST, the relevant prefix # of STRUCTURE_LIST will be predicted. if num_structures <= 0 or num_structures > len(STRUCTURE_LIST): num_structures = len(STRUCTURE_LIST) ground_truth_ids = STRUCTURE_LIST[:num_structures] colours = COLOURS[:num_structures] fill_holes = FILL_HOLES[:num_structures] ground_truth_ids_display_names = [f"zz_{x}" for x in ground_truth_ids] # The amount of GPU memory required increases with both the number of structures and the # number of feature channels. The following is a sensible default to avoid out-of-memory, # but you can override is by passing in another (singleton list) value for feature_channels # from a subclass. num_feature_channels = 32 if num_structures <= 20 else 26 bg_weight = 0.02 if len(ground_truth_ids) > 1 else 0.25 # In case of vertical overlap between brainstem and spinal_cord, we separate them # by converting brainstem voxels to cord, as the latter is clinically more sensitive. # We do the same to separate SPC and MPC; in this case, the direction of change is unimportant, # so we choose SPC-to-MPC arbitrarily. slice_exclusion_rules = [] summed_probability_rules = [] if "brainstem" in ground_truth_ids and "spinal_cord" in ground_truth_ids: slice_exclusion_rules.append( SliceExclusionRule("brainstem", "spinal_cord", False)) if "external" in ground_truth_ids: summed_probability_rules.append( SummedProbabilityRule("spinal_cord", "brainstem", "external")) if "spc_muscle" in ground_truth_ids and "mpc_muscle" in ground_truth_ids: slice_exclusion_rules.append( SliceExclusionRule("spc_muscle", "mpc_muscle", False)) if "external" in ground_truth_ids: summed_probability_rules.append( SummedProbabilityRule("mpc_muscle", "spc_muscle", "external")) if "optic_chiasm" in ground_truth_ids and "pituitary_gland" in ground_truth_ids: slice_exclusion_rules.append( SliceExclusionRule("optic_chiasm", "pituitary_gland", True)) if "external" in ground_truth_ids: summed_probability_rules.append( SummedProbabilityRule("optic_chiasm", "pituitary_gland", "external")) super().__init__( should_validate=False, # we'll validate after kwargs are added num_gpus=4, num_epochs=num_epochs, save_start_epoch=num_epochs, save_step_epoch=num_epochs, architecture="UNet3D", kernel_size=3, train_batch_size=4, inference_batch_size=1, feature_channels=[num_feature_channels], crop_size=(96, 288, 288), test_crop_size=(144, 512, 512), inference_stride_size=(72, 256, 256), image_channels=["ct"], norm_method=PhotometricNormalizationMethod.CtWindow, level=50, window=600, start_epoch=0, l_rate=1e-3, min_l_rate=1e-5, l_rate_polynomial_gamma=0.9, optimizer_type=OptimizerType.Adam, opt_eps=1e-4, adam_betas=(0.9, 0.999), momentum=0.9, test_diff_epochs=1, test_step_epochs=1, use_mixed_precision=True, use_model_parallel=True, monitoring_interval_seconds=0, num_dataload_workers=4, loss_type=SegmentationLoss.Mixture, mixture_loss_components=[ MixtureLossComponent(0.5, SegmentationLoss.Focal, 0.2), MixtureLossComponent(0.5, SegmentationLoss.SoftDice, 0.1) ], ground_truth_ids=ground_truth_ids, ground_truth_ids_display_names=ground_truth_ids_display_names, largest_connected_component_foreground_classes=ground_truth_ids, colours=colours, fill_holes=fill_holes, class_weights=equally_weighted_classes( ground_truth_ids, background_weight=bg_weight), slice_exclusion_rules=slice_exclusion_rules, summed_probability_rules=summed_probability_rules, ) self.add_and_validate(kwargs)
def __init__(self, ground_truth_ids: List[str], ground_truth_ids_display_names: Optional[List[str]] = None, colours: Optional[List[TupleInt3]] = None, fill_holes: Optional[List[bool]] = None, class_weights: Optional[List[float]] = None, slice_exclusion_rules: Optional[List[SliceExclusionRule]] = None, summed_probability_rules: Optional[List[SummedProbabilityRule]] = None, num_feature_channels: Optional[int] = None, **kwargs: Any) -> None: """ Creates a new instance of the class. :param ground_truth_ids: List of ground truth ids. :param ground_truth_ids_display_names: Optional list of ground truth id display names. If present then must be of the same length as ground_truth_ids. :param colours: Optional list of colours. If present then must be of the same length as ground_truth_ids. :param fill_holes: Optional list of fill hole flags. If present then must be of the same length as ground_truth_ids. :param class_weights: Optional list of class weights. If present then must be of the same length as ground_truth_ids + 1. :param slice_exclusion_rules: Optional list of SliceExclusionRules. :param summed_probability_rules: Optional list of SummedProbabilityRule. :param num_feature_channels: Optional number of feature channels. :param kwargs: Additional arguments that will be passed through to the SegmentationModelBase constructor. """ # Number of training epochs num_epochs = 120 num_structures = len(ground_truth_ids) colours = colours or generate_random_colours_list(RANDOM_COLOUR_GENERATOR, num_structures) fill_holes = fill_holes or [True] * num_structures ground_truth_ids_display_names = ground_truth_ids_display_names or [f"zz_{x}" for x in ground_truth_ids] # The amount of GPU memory required increases with both the number of structures and the # number of feature channels. The following is a sensible default to avoid out-of-memory, # but you can override is by passing in another (singleton list) value for feature_channels # from a subclass. num_feature_channels = num_feature_channels or (32 if num_structures <= 20 else 26) bg_weight = 0.02 if len(ground_truth_ids) > 1 else 0.25 class_weights = class_weights or equally_weighted_classes(ground_truth_ids, background_weight=bg_weight) # In case of vertical overlap between brainstem and spinal_cord, we separate them # by converting brainstem voxels to cord, as the latter is clinically more sensitive. # We do the same to separate SPC and MPC; in this case, the direction of change is unimportant, # so we choose SPC-to-MPC arbitrarily. slice_exclusion_rules = slice_exclusion_rules or [] summed_probability_rules = summed_probability_rules or [] super().__init__( should_validate=False, # we'll validate after kwargs are added num_epochs=num_epochs, recovery_checkpoint_save_interval=10, architecture="UNet3D", kernel_size=3, train_batch_size=1, inference_batch_size=1, feature_channels=[num_feature_channels], crop_size=(96, 288, 288), test_crop_size=(144, 512, 512), inference_stride_size=(72, 256, 256), image_channels=["ct"], norm_method=PhotometricNormalizationMethod.CtWindow, level=50, window=600, start_epoch=0, l_rate=1e-3, min_l_rate=1e-5, l_rate_polynomial_gamma=0.9, optimizer_type=OptimizerType.Adam, opt_eps=1e-4, adam_betas=(0.9, 0.999), momentum=0.9, use_mixed_precision=True, use_model_parallel=True, monitoring_interval_seconds=0, num_dataload_workers=2, loss_type=SegmentationLoss.Mixture, mixture_loss_components=[MixtureLossComponent(0.5, SegmentationLoss.Focal, 0.2), MixtureLossComponent(0.5, SegmentationLoss.SoftDice, 0.1)], ground_truth_ids=ground_truth_ids, ground_truth_ids_display_names=ground_truth_ids_display_names, largest_connected_component_foreground_classes=ground_truth_ids, colours=colours, fill_holes=fill_holes, class_weights=class_weights, slice_exclusion_rules=slice_exclusion_rules, summed_probability_rules=summed_probability_rules, ) self.add_and_validate(kwargs)