Spaces:
Sleeping
Sleeping
| import platform | |
| import torch as T | |
| import nibabel | |
| import numpy as np | |
| from monai.networks.nets.unet import UNet | |
| from monai.data import DataLoader, Dataset | |
| from monai.metrics import compute_meandice | |
| from monai.transforms import ( | |
| AddChanneld, | |
| Compose, | |
| LoadImaged, | |
| NormalizeIntensityd, | |
| ToTensord, | |
| Resized, | |
| AsDiscrete, | |
| ThresholdIntensityd, | |
| KeepLargestConnectedComponent | |
| ) | |
| from Engine.utils import ( | |
| read_yaml, | |
| load_model, | |
| store_output | |
| ) | |
| def init_eval_file(directory): | |
| with open(directory + "/evaluate.csv", "w+") as file: | |
| file.write("epoch,train_loss,val_loss,dice_score\n") | |
| def append_eval(eval_path, input_filename, segment_filename, val_loss, dice_score): | |
| with open(eval_path, 'a') as file: | |
| file.write(f"{input_filename},{segment_filename},{val_loss},{dice_score}\n") | |
| def evaluate(model, loss_function, loader, device, evaluate_config): | |
| model.eval() | |
| losses = [] | |
| dices = [] | |
| transform = KeepLargestConnectedComponent([1], connectivity = 3) | |
| with T.no_grad(): | |
| for image in loader: | |
| inp, boxes, label = image["image"].to(device), image["boxes"].to(device), image["label"].to(device) | |
| directory_split = '\\' if platform.system() == 'Windows' else '/' | |
| output = model(inp) | |
| original_filename = image['image_meta_dict']['filename_or_obj'][0] | |
| segmentation_filename = f"{evaluate_config['save_directory']}{directory_split}seg_{original_filename.split(directory_split)[-1]}" | |
| original_image = nibabel.load(original_filename) | |
| loss = loss_function(output, label).item() | |
| output = output.cpu() | |
| output[output >= evaluate_config['output_threshold']] = 1 | |
| output[output < evaluate_config['output_threshold']] = 0 | |
| output = output.to(device) | |
| output = transform(output) | |
| if(evaluate_config['save_segmentations']): | |
| store_output(output, original_image, segmentation_filename, image['image_meta_dict']['affine'].squeeze(0).numpy()) | |
| dice = compute_meandice(output, label).item() | |
| append_eval(evaluate_config['save_directory']+"evaluate.csv", original_filename, segmentation_filename, loss, dice) | |
| losses.append(loss) | |
| dices.append(dice) | |
| append_eval(evaluate_config['save_directory']+"evaluate.csv", 'Total', '', sum(losses)/len(losses), sum(dices)/len(dices)) | |
| def initiate(config_file): | |
| config = read_yaml(config_file) | |
| device = T.device(config["device"]) | |
| data = read_yaml(config["data"]["train_dataset"]) | |
| image_shape = (config["data"]["scale_dim"]["d_0"], config["data"]["scale_dim"]["d_1"], config["data"]["scale_dim"]["d_2"]) | |
| for i, d in enumerate(data['data']): | |
| data['data'][i]['image'] = data['image_prefix'] + d['label'] | |
| data['data'][i]['boxes'] = data['boxes_prefix'] + d['label'] | |
| data['data'][i]['label'] = data['label_prefix'] + d['label'] | |
| transform = Compose( | |
| [ | |
| LoadImaged(keys=["image", "boxes", "label"]), | |
| AddChanneld(keys=["image", "boxes", "label"]), | |
| ToTensord(keys=["image", "boxes", "label"]), | |
| ] | |
| ) | |
| dataset = Dataset(data['data'], transform) | |
| loader = T.utils.data.DataLoader(dataset, 1) | |
| model, loss = load_model(config['model'], eval = True) | |
| model.load_state_dict(T.load(config["model"]["weights"])) | |
| model.to(device) | |
| init_eval_file(config['evaluate']['save_directory']) | |
| evaluate(model, loss, loader, device, config['evaluate']) | |