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)