File size: 3,490 Bytes
a23d562 | 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 | """
Validate an already trained cropping UNet on a test split, with the training augmentation applied.
Same purpose as val_trained_unet.py in the repository root, which validates without augmentation.
Paths not given on the command line come from the site config (see config.example.yaml):
python scripts/validation/validate_crop_unet.py --unet-path <unet_best.pth>
"""
import argparse
import torchio as tio
from auto_detect_breast_mri.config import resolve_path
from auto_detect_breast_mri.data.transforms import ImageToTensor, RandomCropOrPad, ZNormalization, masking_method
from auto_detect_breast_mri.models.unet_crop import UNet_Crop
from auto_detect_breast_mri.data.loaders import get_subjects_dataloader
from auto_detect_breast_mri.data.metadata import get_uka_metatensor
def build_parser():
parser = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--unet-path", required=True,
help="pth file of the trained cropping UNet")
parser.add_argument("--data-path", default=None,
help="image root. Config key: data_root")
parser.add_argument("--feature-path", default=None,
help="metadata export. Config key: metadata_file")
parser.add_argument("--split-files-folder", default=None,
help="folder holding the split files. Config key: split_root")
parser.add_argument("-b", "--batch-size", type=int, default=4)
parser.add_argument("-c", "--fold", type=int, default=0, help="outer fold whose test set is used")
parser.add_argument("-n", "--number-of-epochs", type=int, default=10,
help="how often the validation set is passed through. Default: 10")
parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
return parser
if __name__ == '__main__':
args = build_parser().parse_args()
path_base = resolve_path(args.data_path, "data_root", "image root")
feature_path = resolve_path(args.feature_path, "metadata_file", "metadata export")
split_root = resolve_path(args.split_files_folder, "split_root", "folder with the split files")
test_split = f"{split_root.rstrip('/')}/fold{args.fold}/stratified_test_set"
protocol = ['Dyn_0']
pre_image_shape = (32, 512, 512)
augmentation = tio.Compose([
tio.RandomFlip((0, 1, 2), flip_probability=0.5),
RandomCropOrPad(pre_image_shape),
ZNormalization(per_channel=True, percentiles=(0.5, 99.5), masking_method=masking_method),
tio.RandomNoise(std=(0.25, 0.5)),
ImageToTensor()
])
# ------------ Initialize Model ------------
in_channel = len(protocol)
CropNet = UNet_Crop(path_to_unet=args.unet_path, input_channel=in_channel, running_device=args.device,
show_validation=True, criterion='dice')
# Dataset and Loader
feature_dataframe = get_uka_metatensor(0, feature_path)
test_loader = get_subjects_dataloader(path_base, feature_dataframe, pre_image_shape,
transform=augmentation, protocol=protocol,
data_selection_file_sceleton=test_split,
batch_size=args.batch_size, fold=args.fold,
segmentation_task=True)
CropNet.validate(args.number_of_epochs, test_loader)
|