Jupitern52/TextBraTS

1# Copyright 2020 - 2022 MONAI Consortium2# Licensed under the Apache License, Version 2.0 (the "License");3# you may not use this file except in compliance with the License.4# You may obtain a copy of the License at5#     http://www.apache.org/licenses/LICENSE-2.06# Unless required by applicable law or agreed to in writing, software7# distributed under the License is distributed on an "AS IS" BASIS,8# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.9# See the License for the specific language governing permissions and10# limitations under the License.11 12import argparse13from utils.data_utils import get_loader14from utils.textswin_unetr import TextSwinUNETR15import os16import time17import torch18import torch.nn.parallel19import torch.utils.data.distributed20from utils.utils import AverageMeter21from monai.utils.enums import MetricReduction22from monai.metrics import DiceMetric, HausdorffDistanceMetric23 24 25parser = argparse.ArgumentParser(description="TextBraTS segmentation pipeline")26parser.add_argument("--data_dir", default="./data/TextBraTSData", type=str, help="dataset directory")27parser.add_argument("--exp_name", default="TextBraTS", type=str, help="experiment name")28parser.add_argument("--json_list", default="Test.json", type=str, help="dataset json file")29parser.add_argument("--fold", default=0, type=int, help="data fold")30parser.add_argument("--pretrained_model_name", default="model.pt", type=str, help="pretrained model name")31parser.add_argument("--feature_size", default=48, type=int, help="feature size")32parser.add_argument("--infer_overlap", default=0.6, type=float, help="sliding window inference overlap")33parser.add_argument("--in_channels", default=4, type=int, help="number of input channels")34parser.add_argument("--out_channels", default=3, type=int, help="number of output channels")35parser.add_argument("--a_min", default=-175.0, type=float, help="a_min in ScaleIntensityRanged")36parser.add_argument("--a_max", default=250.0, type=float, help="a_max in ScaleIntensityRanged")37parser.add_argument("--b_min", default=0.0, type=float, help="b_min in ScaleIntensityRanged")38parser.add_argument("--b_max", default=1.0, type=float, help="b_max in ScaleIntensityRanged")39parser.add_argument("--space_x", default=1.5, type=float, help="spacing in x direction")40parser.add_argument("--space_y", default=1.5, type=float, help="spacing in y direction")41parser.add_argument("--space_z", default=2.0, type=float, help="spacing in z direction")42parser.add_argument("--roi_x", default=128, type=int, help="roi size in x direction")43parser.add_argument("--roi_y", default=128, type=int, help="roi size in y direction")44parser.add_argument("--roi_z", default=128, type=int, help="roi size in z direction")45parser.add_argument("--dropout_rate", default=0.0, type=float, help="dropout rate")46parser.add_argument("--distributed", action="store_true", help="start distributed training")47parser.add_argument("--workers", default=8, type=int, help="number of workers")48parser.add_argument("--RandScaleIntensityd_prob", default=0.1, type=float, help="RandScaleIntensityd aug probability")49parser.add_argument("--RandShiftIntensityd_prob", default=0.1, type=float, help="RandShiftIntensityd aug probability")50parser.add_argument("--spatial_dims", default=3, type=int, help="spatial dimension of input data")51parser.add_argument("--use_checkpoint", action="store_true", help="use gradient checkpointing to save memory")52parser.add_argument(53    "--pretrained_dir",54    default="./runs/TextBraTS/",55    type=str,56    help="pretrained checkpoint directory",57)58 59 60def main():61    args = parser.parse_args()62    args.test_mode = True63    output_directory = "./outputs/" + args.exp_name64    if not os.path.exists(output_directory):65        os.makedirs(output_directory)66    test_loader = get_loader(args)67    pretrained_dir = args.pretrained_dir68    model_name = args.pretrained_model_name69    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")70    pretrained_pth = os.path.join(pretrained_dir, model_name)71    model = TextSwinUNETR(72        img_size=128,73        in_channels=args.in_channels,74        out_channels=args.out_channels,75        feature_size=args.feature_size,76        drop_rate=0.0,77        attn_drop_rate=0.0,78        dropout_path_rate=0.0,79        use_checkpoint=args.use_checkpoint,80        text_dim=768,81    )82    model_dict = torch.load(pretrained_pth)["state_dict"]83    model.load_state_dict(model_dict, strict=False)84    model.eval()85    model.to(device)86 87    def val_epoch(model, loader, acc_func,  hd95_func):88        model.eval()89        start_time = time.time()90        run_acc = AverageMeter()91        run_hd95 = AverageMeter()92 93        with torch.no_grad():94            for idx, batch_data in enumerate(loader):95                data, target, text = batch_data["image"], batch_data["label"], batch_data["text_feature"]96                data, target, text = data.cuda(), target.cuda(), text.cuda()97                logits = model(data,text)98                prob = torch.sigmoid(logits)99                prob = (prob > 0.5).int()100 101                acc_func(y_pred=prob, y=target)102                acc, not_nans = acc_func.aggregate()103                acc = acc.cuda()104 105                run_acc.update(acc.cpu().numpy(), n=not_nans.cpu().numpy())106 107                # HD95 Metric108                hd95_func(y_pred=prob, y=target)109                hd95 = hd95_func.aggregate()  # Assuming it returns a single value110                run_hd95.update(hd95.cpu().numpy())111 112 113                Dice_TC = run_acc.avg[0]114                Dice_WT = run_acc.avg[1]115                Dice_ET = run_acc.avg[2]116                HD95_TC = run_hd95.avg[0]117                HD95_WT = run_hd95.avg[1]118                HD95_ET = run_hd95.avg[2]119                print(120                    "Val  {}/{}".format(idx, len(loader)),121                    ", Dice_TC:", Dice_TC,122                    ", Dice_WT:", Dice_WT,123                    ", Dice_ET:", Dice_ET,124                    ", Avg Dice:", (Dice_ET + Dice_TC + Dice_WT) / 3,125                    ", HD95_TC:", HD95_TC,126                    ", HD95_WT:", HD95_WT,127                    ", HD95_ET:", HD95_ET,128                    ", Avg HD95:", (HD95_ET + HD95_TC + HD95_WT) / 3,129                    ", time {:.2f}s".format(time.time() - start_time),130                )131                start_time = time.time()132            with open(output_directory+'/log.txt', "a") as log_file:133                log_file.write(f"Experiment name:{args.pretrained_dir.split('/')[-2]}, "134            f"Final Validation Results - Dice_TC: {Dice_TC}, Dice_WT: {Dice_WT}, Dice_ET: {Dice_ET}, "135            f"Avg Dice: {(Dice_ET + Dice_TC + Dice_WT) / 3}, "136            f"HD95_TC: {HD95_TC}, HD95_WT: {HD95_WT}, HD95_ET: {HD95_ET}, "137            f"Avg HD95: {(HD95_ET + HD95_TC + HD95_WT) / 3}\n")138        return run_acc.avg139 140    dice_acc = DiceMetric(include_background=True, reduction=MetricReduction.MEAN_BATCH, get_not_nans=True)141    hd95_acc = HausdorffDistanceMetric(include_background=True, reduction=MetricReduction.MEAN_BATCH, percentile=95.0)142    val_epoch(model, test_loader, acc_func=dice_acc,hd95_func=hd95_acc)143 144if __name__ == "__main__":145    main()146