add necessary module

dataset/__init__.py

@@ -0,0 +1 @@
+# ProFound dataset package 

dataset/dataset_cls.py

@@ -0,0 +1,480 @@
+import pickle
+from monai.transforms import (
+    Compose,
+    RandCropByPosNegLabeld,
+    CropForegroundd,
+    SpatialPadd,
+    ScaleIntensityRanged,
+    RandShiftIntensityd,
+    RandFlipd,
+    RandAffined,
+    RandZoomd,
+    RandRotated,
+    RandBiasFieldd,
+    RandRotate90d,
+    RandGaussianNoised,
+    RandGaussianSmoothd,
+    NormalizeIntensityd,
+    MapTransform,
+    RandScaleIntensityd,
+    RandSpatialCropd,
+    CenterSpatialCropd,
+)
+
+from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler
+import torch
+import numpy as np
+import nibabel as nib
+import torch.nn.functional as F
+import os
+import pandas as pd
+from ast import literal_eval
+
+
+class RiskSet(Dataset):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__()
+        self.img_dict = pd.read_csv(image_paths)
+        if phase == 'train':
+            if args.data_num > 0:
+                # crop the dataset
+                self.img_dict = self.img_dict.iloc[: args.data_num]
+        print(f"Loading {phase} dataset with {len(self.img_dict)} samples")
+        self.root = args.root
+        self._set_dataset_stat()
+        self.transforms = transforms  # self.get_transforms()
+        if not args.demo:
+            self.set_sampler()
+
+    def set_sampler(self):
+        class_counts = self.img_dict["pirads"].value_counts().sort_index().values
+        class_weights = 1.0 / class_counts
+        values = self.img_dict["pirads"].values.astype(int) - 2
+        self.sampler_weight = class_weights[values]
+
+    def cal_weight(self):
+        class_counts = self.img_dict["pirads"].value_counts().sort_index().values
+        return class_counts
+
+    def _set_dataset_stat(self):
+        self.spacing = (0.5, 0.5, 1.0)
+        self.spatial_index = [2, 1, 0]  # index used to convert to DHW
+        self.target_class = 1
+
+    def __len__(self):
+        return len(self.img_dict)
+
+    def read(self, path):
+        vol = nib.load(os.path.join(self.root, path))
+        vol = vol.get_fdata().astype(np.float32).transpose(self.spatial_index)
+        vol = torch.from_numpy(vol)
+        return vol
+
+    def __getitem__(self, idx):
+        path = self.img_dict.iloc[idx]
+        t2w = self.read(path["t2w"])
+        dwi = self.read(path["highb"])
+        adc = self.read(path["adc"])
+        img = torch.stack([t2w, dwi, adc], 0)
+        label = torch.tensor(int(path["pirads"]) - 2, dtype=torch.long)
+        if self.transforms is not None:
+            trans_dict = self.transforms({"image": img})
+            if type(trans_dict) == list:
+                trans_dict = trans_dict[0]
+            img = trans_dict["image"]
+        return img, label, torch.tensor(idx, dtype=torch.long)
+
+
+class ScreeningSet(RiskSet):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__(args=args, image_paths=image_paths, phase = phase, transforms=transforms)
+
+    def set_sampler(self):
+        class_counts = self.img_dict["result"].value_counts().sort_index().values
+        class_weights = 1.0 / class_counts
+        self.sampler_weight = class_weights[self.img_dict["result"].values]
+
+    def cal_weight(self):
+        class_counts = self.img_dict["result"].value_counts().sort_index().values
+        return class_counts
+
+    def __getitem__(self, idx):
+        path = self.img_dict.iloc[idx]
+        t2w = self.read(path["t2w"])
+        dwi = self.read(path["dwi"])
+        adc = self.read(path["adc"])
+        img = torch.stack([t2w, dwi, adc], 0)
+        label = torch.tensor(int(path["result"]), dtype=torch.long)
+        if self.transforms is not None:
+            trans_dict = self.transforms({"image": img})
+            if type(trans_dict) == list:
+                trans_dict = trans_dict[0]
+            img = trans_dict["image"]
+        return img, label, torch.tensor(idx, dtype=torch.long)
+
+
+class PromisSet(RiskSet):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__(args=args, image_paths=image_paths, phase=phase, transforms=transforms)
+
+    def set_sampler(self):
+        class_counts = self.img_dict["patient_level"].value_counts().sort_index().values
+        class_weights = 1.0 / class_counts
+        self.sampler_weight = class_weights[self.img_dict["patient_level"].values.astype(int)]
+
+    def cal_weight(self):
+        class_counts = self.img_dict["patient_level"].value_counts().sort_index().values
+        return class_counts
+
+    def __getitem__(self, idx):
+        path = self.img_dict.iloc[idx]
+        t2w = self.read(path["t2w"])
+        dwi = self.read(path["dwi"])
+        adc = self.read(path["adc"])
+        img = torch.stack([t2w, dwi, adc], 0)
+        zone_level = literal_eval(path["zone_level"])
+        zone_level = torch.tensor(zone_level, dtype=torch.float32)
+        #patient_level  = torch.tensor(int(path["patient_level"]), dtype=torch.float32)
+        if self.transforms is not None:
+            trans_dict = self.transforms({"image": img})
+            if type(trans_dict) == list:
+                trans_dict = trans_dict[0]
+            img = trans_dict["image"]
+        return img, zone_level, torch.tensor(idx, dtype=torch.long)
+
+class Promis3HistSet(RiskSet):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__(args=args, image_paths=image_paths, phase=phase, transforms=transforms)
+
+    def set_sampler(self):
+        class_counts = self.img_dict["def"].value_counts().sort_index().values
+        class_weights = 1.0 / class_counts
+        self.sampler_weight = class_weights[self.img_dict["def"].values.astype(int)]
+
+    def cal_weight(self):
+        class_counts = self.img_dict["def"].value_counts().sort_index().values
+        return class_counts
+
+    def __getitem__(self, idx):
+        path = self.img_dict.iloc[idx]
+        t2w = self.read(path["t2w"])
+        dwi = self.read(path["dwi"])
+        adc = self.read(path["adc"])
+        img = torch.stack([t2w, dwi, adc], 0)
+        label = torch.tensor(int(path["def"]), dtype=torch.long)
+        if self.transforms is not None:
+            trans_dict = self.transforms({"image": img})
+            if type(trans_dict) == list:
+                trans_dict = trans_dict[0]
+            img = trans_dict["image"]
+        return img, label, torch.tensor(idx, dtype=torch.long)
+
+class Promis3GGSet(RiskSet):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__(args=args, image_paths=image_paths, phase=phase, transforms=transforms)
+
+    def set_sampler(self):
+        class_counts = self.img_dict["gleason"].value_counts().sort_index().values
+        class_weights = 1.0 / class_counts
+        self.sampler_weight = class_weights[self.img_dict["gleason"].values.astype(int)]
+
+    def cal_weight(self):
+        class_counts = self.img_dict["gleason"].value_counts().sort_index().values
+        return class_counts
+
+    def __getitem__(self, idx):
+        path = self.img_dict.iloc[idx]
+        t2w = self.read(path["t2w"])
+        dwi = self.read(path["dwi"])
+        adc = self.read(path["adc"])
+        img = torch.stack([t2w, dwi, adc], 0)
+        label = torch.tensor(int(path["gleason"]), dtype=torch.long)
+        if self.transforms is not None:
+            trans_dict = self.transforms({"image": img})
+            if type(trans_dict) == list:
+                trans_dict = trans_dict[0]
+            img = trans_dict["image"]
+        return img, label, torch.tensor(idx, dtype=torch.long)
+
+
+def get_transforms(args):
+    train_transforms = [
+        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
+        CenterSpatialCropd(keys="image", roi_size=(80, 300, 300)),
+        RandRotated(
+            keys="image",
+            prob=0.3,
+            range_x=10 / 180 * np.pi,
+            range_y=10 / 180 * np.pi,
+            range_z=10 / 180 * np.pi,
+            keep_size=False,
+            mode="bilinear",
+        ),
+        RandZoomd(
+            keys="image",
+            prob=0.3,
+            min_zoom=[0.9, 0.9, 0.9],
+            max_zoom=[1.1, 1.1, 1.1],
+            mode="trilinear",
+        ),
+        SpatialPadd(
+            keys="image",
+            spatial_size=[round(i * 1.2) for i in args.crop_spatial_size],
+        ),
+        RandSpatialCropd(
+            keys="image",
+            roi_size=args.crop_spatial_size,
+            random_size=False,
+        ),
+        RandFlipd(keys="image", prob=0.5, spatial_axis=2),
+        # BinarizeLabeld(keys=["label"])
+        RandScaleIntensityd(keys="image", factors=0.1, prob=0.8),
+        RandShiftIntensityd(keys="image", offsets=0.1, prob=0.8),
+        RandBiasFieldd(keys="image", prob=0.2),
+        RandGaussianSmoothd(keys="image", prob=1.0)
+    ]
+
+    train_transforms = Compose(train_transforms)
+    val_transforms = Compose(
+        [
+            NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
+            CenterSpatialCropd(keys="image", roi_size=args.crop_spatial_size),
+            SpatialPadd(keys="image", spatial_size=[i for i in args.crop_spatial_size]),
+            # BinarizeLabeld(keys=["label"])
+        ]
+    )
+    test_transforms = Compose(
+        [
+            NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
+            CenterSpatialCropd(keys="image", roi_size=args.crop_spatial_size),
+            SpatialPadd(keys="image", spatial_size=[i for i in args.crop_spatial_size]),
+            # BinarizeLabeld(keys=["label"])
+        ]
+    )
+    return train_transforms, val_transforms, test_transforms
+
+
+def build_Risk_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    
+    if args.demo:
+        test_set = RiskSet(args, "demo/data/risk/test.csv", 'test', test_transforms)
+        test_loader = DataLoader(
+            test_set,
+            batch_size=args.batch_size,
+            shuffle=False,
+            pin_memory=True,
+            num_workers=14,
+            drop_last=False,
+        )
+        args.in_channels = 3
+        args.num_classes = 4
+        return test_loader
+    else:
+        if args.data20:
+            train_set = RiskSet(args, "spilt/risk/train_16.csv", 'train', train_transforms)
+        else:
+            train_set = RiskSet(args, "spilt/risk/train.csv", 'train', train_transforms)
+        val_set = RiskSet(args, "spilt/risk/val.csv", 'val', val_transforms)
+        test_set = RiskSet(args, "spilt/risk/test.csv", 'test', test_transforms)
+
+        sampler = WeightedRandomSampler(
+            weights=train_set.sampler_weight, num_samples=len(train_set), replacement=True
+        )
+        train_loader = DataLoader(
+            train_set,
+            batch_size=args.batch_size,
+            sampler=sampler,
+            num_workers=args.num_workers,
+            drop_last=False,
+            pin_memory=True,
+        )
+        val_loader = DataLoader(
+            val_set,
+            batch_size=args.batch_size,
+            shuffle=False,
+            pin_memory=True,
+            num_workers=14,
+            drop_last=False,
+        )
+        test_loader = DataLoader(
+            test_set,
+            batch_size=args.batch_size,
+            shuffle=False,
+            pin_memory=True,
+            num_workers=14,
+            drop_last=False,
+        )
+        args.in_channels = 3
+        args.num_classes = 4
+        return train_loader, val_loader, test_loader
+
+
+def build_Screening_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    if args.kfold is None:
+        if args.data20:
+            train_set = ScreeningSet(
+                args, "spilt/screening/train_20.csv", 'train', train_transforms
+            )
+        else:
+            train_set = ScreeningSet(
+                args, "spilt/screening/train.csv", 'train', train_transforms
+            )
+        val_set = ScreeningSet(args, "spilt/screening/val.csv", 'val', val_transforms)
+        test_set = ScreeningSet(args, "spilt/screening/test.csv", 'test', test_transforms)
+        args.cls_account = train_set.cal_weight() / len(train_set)
+    else:
+        train_set = ScreeningSet(
+            args, f"spilt/screening/train_{args.kfold}.csv", train_transforms
+        )
+        args.cls_account = train_set.cal_weight() / len(train_set)
+        train_set, val_set = torch.utils.data.random_split(train_set, [0.9, 0.1])
+        val_set.transforms = val_transforms
+        test_set = ScreeningSet(
+            args, f"spilt/screening/test_{args.kfold}.csv", test_transforms
+        )
+
+    # sampler_weight = [train_set.dataset.sampler_weight[i] for i in train_set.indices]
+    sampler = WeightedRandomSampler(
+        weights=train_set.sampler_weight, num_samples=len(train_set), replacement=True
+    )
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        sampler=sampler,
+        num_workers=args.num_workers,
+        drop_last=True,
+        pin_memory=True,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    test_loader = DataLoader(
+        test_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    args.in_channels = 3
+    args.num_classes = 2
+    return train_loader, val_loader, test_loader
+
+
+# 4.0    453
+# 3.0    206
+# 5.0    195
+# 2.0    174
+
+
+def build_Promis_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    if args.data20:
+        train_set = PromisSet(args, "spilt/promis567_hist/train_20.csv", 'train', train_transforms)
+    else:
+        train_set = PromisSet(args, "spilt/promis567_hist/train.csv", 'train', train_transforms)
+    val_set = PromisSet(args, "spilt/promis567_hist/val.csv", 'val', val_transforms)
+    test_set = PromisSet(args, "spilt/promis567_hist/test.csv", 'test', test_transforms)
+
+    # sampler = WeightedRandomSampler(
+    #     weights=train_set.sampler_weight, num_samples=len(train_set), replacement=True
+    # )
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        drop_last=True,
+        pin_memory=True,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    test_loader = DataLoader(
+        test_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    args.in_channels = 3
+    args.num_classes = 20
+    return train_loader, val_loader, test_loader
+
+def build_Promis3_hist_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    train_set = Promis3HistSet(args, "spilt/promis_pirads3_hist/train.csv", 'train', train_transforms)
+    val_set = Promis3HistSet(args, "spilt/promis_pirads3_hist/val.csv", 'val', val_transforms)
+    test_set = Promis3HistSet(args, "spilt/promis_pirads3_hist/test.csv", 'test', test_transforms)
+
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        drop_last=True,
+        pin_memory=True,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    test_loader = DataLoader(
+        test_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    args.in_channels = 3
+    args.num_classes = 3
+    return train_loader, val_loader, test_loader
+
+def build_Promis3_gg_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    train_set = Promis3GGSet(args, "spilt/promis_pirads3_gg/train.csv", 'train', train_transforms)
+    val_set = Promis3GGSet(args, "spilt/promis_pirads3_gg/val.csv", 'val', val_transforms)
+    test_set = Promis3GGSet(args, "spilt/promis_pirads3_gg/test.csv", 'test', test_transforms)
+
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        drop_last=True,
+        pin_memory=True,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    test_loader = DataLoader(
+        test_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    args.in_channels = 3
+    args.num_classes = 5
+    return train_loader, val_loader, test_loader

dataset/dataset_seg.py

@@ -0,0 +1,556 @@
+import pickle
+from monai.transforms import (
+    Compose,
+    RandCropByPosNegLabeld,
+    CropForegroundd,
+    SpatialPadd,
+    ScaleIntensityRanged,
+    RandShiftIntensityd,
+    RandFlipd,
+    RandAffined,
+    RandZoomd,
+    RandRotated,
+    RandRotate90d,
+    RandGaussianNoised,
+    RandGaussianSmoothd,
+    NormalizeIntensityd,
+    RandBiasFieldd,
+    MapTransform,
+    RandScaleIntensityd,
+    RandSpatialCropd,
+    CenterSpatialCropd,
+)
+
+from torch.utils.data import DataLoader, Dataset
+import torch
+import numpy as np
+import nibabel as nib
+import torch.nn.functional as F
+import os
+import pandas as pd
+
+
+class BaseVolumeDataset(Dataset):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__()
+        self.img_dict = pd.read_csv(image_paths)
+        if phase == 'train':
+            if args.data_num > 0:
+                # crop the dataset
+                self.img_dict = self.img_dict.iloc[: args.data_num]
+        print(f"Loading {phase} dataset with {len(self.img_dict)} samples")
+        self.root = args.root
+        self._set_dataset_stat()
+        self.transforms = transforms  # self.get_transforms()
+
+    def _set_dataset_stat(self):
+        self.spacing = (0.5, 0.5, 1.0)
+        self.spatial_index = [2, 1, 0]  # index used to convert to DHW
+        self.target_class = 1
+
+    def __len__(self):
+        return len(self.img_dict)
+
+    def read(self, path):
+        vol = nib.load(os.path.join(self.root, path))
+        vol = vol.get_fdata().astype(np.float32).transpose(self.spatial_index)
+        vol = torch.from_numpy(vol)
+        return vol
+
+    def __getitem__(self, idx):
+        return NotImplemented
+
+
+class UCLSet(BaseVolumeDataset):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__(args=args, image_paths=image_paths, phase=phase, transforms=transforms)
+
+    def __getitem__(self, idx):
+        path = self.img_dict.iloc[idx]
+        t2w = self.read(path["t2w"])
+        dwi = self.read(path["dwi"])
+        adc = self.read(path["adc"])
+        img = torch.stack([t2w, dwi, adc], 0)
+        seg = self.read(path["lesion"]).unsqueeze(0)
+        seg = seg > 0
+        # print(img.shape)
+        # seg = (seg == self.target_class).float()
+        if self.transforms is not None:
+            trans_dict = self.transforms({"image": img, "label": seg})
+            if type(trans_dict) == list:
+                trans_dict = trans_dict[0]
+            img, seg = trans_dict["image"], trans_dict["label"]
+        return img, seg, torch.tensor(idx, dtype=torch.long)
+
+# TODO: need to update; unfinished
+"""
+class UCL2DSet(BaseVolumeDataset):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__(args=args, image_paths=image_paths, phase=phase, transforms=transforms)
+    
+    def __getitem__(self, idx):
+        path = self.img_dict.iloc[idx]
+        t2w = self.read(path["t2w"])
+        dwi = self.read(path["dwi"])
+        adc = self.read(path["adc"])
+
+        seg = self.read(path["lesion"]).unsqueeze(0)
+        seg = seg > 0
+
+        seg_mask = seg.squeeze(0).numpy()
+        non_zero_slices = np.where(seg_mask.any(axis=1,2))[0]
+        if len(non_zero_slices) > 0:
+            sampled_slices = np.random.choice(non_zero_slices, min(N, len(non_zero_slices)), replace=False)
+            filtered_seg = np.zeros_like(seg_mask)
+            filtered_seg[sampled_slices] = seg_mask[sampled_slices]
+        else:
+            filtered_seg = seg_mask
+        
+        img = torch.stack([t2w, dwi, adc], 0)
+        seg = torch.tensor(filtered_seg, dtype=torch.float32).unsqueeze(0)
+        if self.transforms is not None:
+            trans_dict = self.transforms({"image": img, "label": seg})
+            if type(trans_dict) == list:
+                trans_dict = trans_dict[0]
+            img, seg = trans_dict["image"], trans_dict["label"]
+        return img, seg, torch.tensor(idx, dtype=torch.long)
+"""
+
+class AnatomySet(BaseVolumeDataset):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__(args=args, image_paths=image_paths, phase=phase, transforms=transforms)
+    def __getitem__(self, idx):
+        path = self.img_dict.iloc[idx]
+        t2w = self.read(path["t2w"])
+        # img = t2w.unsqueeze(0)
+        zero = torch.zeros_like(t2w)
+        # modified to align img to 3 channel
+        img = torch.stack([t2w, zero, zero], 0)
+        seg = self.read(path["mask"]).unsqueeze(0)
+        if self.transforms is not None:
+            trans_dict = self.transforms({"image": img, "label": seg})
+            if type(trans_dict) == list:
+                trans_dict = trans_dict[0]
+            img, seg = trans_dict["image"], trans_dict["label"]
+        return img, seg, torch.tensor(idx, dtype=torch.long)
+
+
+class BpAnatomySet(BaseVolumeDataset):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__(args=args, image_paths=image_paths, phase=phase, transforms=transforms)
+
+    def __getitem__(self, idx):
+        path = self.img_dict.iloc[idx]
+        t2w = self.read(path["t2w"])
+        zero = torch.zeros_like(t2w)
+        img = torch.stack([t2w, zero, zero], 0)
+        seg = self.read(path["mask"]).unsqueeze(0)
+        if self.transforms is not None:
+            trans_dict = self.transforms({"image": img, "label": seg})
+            if type(trans_dict) == list:
+                trans_dict = trans_dict[0]
+            img, seg = trans_dict["image"], trans_dict["label"]
+        return img, seg, torch.tensor(idx, dtype=torch.long)        
+
+class PromisHist(BaseVolumeDataset):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__(args=args, image_paths=image_paths, phase=phase, transforms=transforms)
+
+    def __getitem__(self, idx):
+        path = self.img_dict.iloc[idx]
+        t2w = self.read(path["t2w"])
+        dwi = self.read(path["dwi"])
+        adc = self.read(path["adc"])
+        img = torch.stack([t2w, dwi, adc], 0)
+
+        zone_mask = self.read(path["gland"]).unsqueeze(0)
+
+        zone_level = list(map(int, path["zone_label"].split()))
+        zone_level = torch.tensor(zone_level)
+        
+        if self.transforms is not None:
+            trans_dict = self.transforms({"image": img, "label": zone_mask})
+            if type(trans_dict) == list:
+                trans_dict = trans_dict[0]
+            img, zone_mask = trans_dict["image"], trans_dict["label"]
+
+        return img, zone_mask, zone_level
+
+class PromisZone(BaseVolumeDataset):
+    def __init__(self, args, image_paths, phase, transforms=None):
+        super().__init__(args=args, image_paths=image_paths, phase=phase, transforms=transforms)
+
+    def __getitem__(self, idx):
+        path = self.img_dict.iloc[idx]
+        t2w = self.read(path["t2w"])
+        dwi = self.read(path["dwi"])
+        adc = self.read(path["adc"])
+        img = torch.stack([t2w, dwi, adc], 0)
+
+        zone_mask = self.read(path["zome_mask"]).unsqueeze(0)
+
+        zone_level = list(map(int, path["zone_label"].split()))
+        zone_level = torch.tensor(zone_level)
+        
+        if self.transforms is not None:
+            trans_dict = self.transforms({"image": img, "label": zone_mask})
+            if type(trans_dict) == list:
+                trans_dict = trans_dict[0]
+            img, zone_mask = trans_dict["image"], trans_dict["label"]
+
+        return img, zone_mask, zone_level
+
+def get_transforms(args):
+    train_transforms = [
+        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
+        RandRotated(
+            keys=["image", "label"],
+            prob=0.3,
+            range_x=30 / 180 * np.pi,
+            keep_size=False,
+            mode=["bilinear", "nearest"],
+        ),
+        RandZoomd(
+            keys=["image", "label"],
+            prob=0.3,
+            min_zoom=[1, 0.9, 0.9],
+            max_zoom=[1, 1.1, 1.1],
+            mode=["trilinear", "nearest"],
+        ),
+        SpatialPadd(
+            keys=["image", "label"],
+            spatial_size=[round(i * 1.2) for i in args.crop_spatial_size],
+        ),
+        # RandCropByPosNegLabeld(
+        #     keys=["image", "label"],
+        #     spatial_size=[round(i * 1.2) for i in args.crop_spatial_size],
+        #     label_key="label",
+        #     pos=2,
+        #     neg=1,
+        #     num_samples=1,
+        # ),
+        RandSpatialCropd(
+            keys=["image", "label"],
+            roi_size=args.crop_spatial_size,
+            random_size=False,
+        ),
+        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
+        # BinarizeLabeld(keys=["label"])
+        RandScaleIntensityd(keys="image", factors=0.1, prob=0.8),
+        RandShiftIntensityd(keys="image", offsets=0.1, prob=0.8),
+        RandBiasFieldd(keys="image", prob=0.2),
+        RandGaussianSmoothd(keys="image", prob=1.0)
+    ]
+
+    train_transforms = Compose(train_transforms)
+    val_transforms = Compose(
+        [
+            NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
+            CenterSpatialCropd(
+                keys=["image", "label"], roi_size=args.crop_spatial_size
+            ),
+            SpatialPadd(
+                keys=["image", "label"],
+                spatial_size=[i for i in args.crop_spatial_size],
+            ),
+            # BinarizeLabeld(keys=["label"])
+        ]
+    )
+    test_transforms = Compose(
+        [
+            NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
+            CenterSpatialCropd(
+                keys=["image", "label"], roi_size=args.crop_spatial_size
+            ),
+            SpatialPadd(
+                keys=["image", "label"],
+                spatial_size=[i for i in args.crop_spatial_size],
+            ),
+            # BinarizeLabeld(keys=["label"])
+        ]
+    )
+    return train_transforms, val_transforms, test_transforms
+
+
+def build_UCL_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    if args.demo:
+        test_set = UCLSet(args, "demo/data/UCL/test.csv", 'test', test_transforms)
+        test_loader = DataLoader(
+            test_set,
+            batch_size=1,
+            shuffle=False,
+            pin_memory=True,
+            num_workers=14,
+            drop_last=False,
+        )
+        args.in_channels = 3
+        args.out_channels = 1
+        args.num_classes = 1
+        return test_loader
+    else:
+        if args.data20:
+            train_set = UCLSet(args, "spilt/UCL/train_16.csv", 'train', train_transforms)
+        else:
+            train_set = UCLSet(args, "spilt/UCL/train.csv", 'train', train_transforms)
+        val_set = UCLSet(args, "spilt/UCL/val.csv", 'val', val_transforms)
+        test_set = UCLSet(args, "spilt/UCL/test.csv", 'test', test_transforms)
+        train_loader = DataLoader(
+            train_set,
+            batch_size=args.batch_size,
+            shuffle=True,
+            pin_memory=True,
+            num_workers=14,
+            drop_last=True,
+        )
+        val_loader = DataLoader(
+            val_set,
+            batch_size=args.batch_size,
+            shuffle=False,
+            pin_memory=True,
+            num_workers=14,
+            drop_last=False,
+        )
+        test_loader = DataLoader(
+            test_set,
+            batch_size=1,
+            shuffle=False,
+            pin_memory=True,
+            num_workers=14,
+            drop_last=False,
+        )
+        args.in_channels = 3
+        args.out_channels = 1
+        args.num_classes = 1
+        return train_loader, val_loader, test_loader
+
+
+def build_Promis_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    if args.data20:
+        train_set = UCLSet(args, "spilt/promis567/train_20.csv", 'train', train_transforms)
+    else:
+        train_set = UCLSet(args, "spilt/promis567/train.csv", 'train', train_transforms)
+    val_set = UCLSet(args, "spilt/promis567/val.csv", 'val', val_transforms)
+    test_set = UCLSet(args, "spilt/promis567/test.csv", 'test', test_transforms)
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        shuffle=True,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    test_loader = DataLoader(
+        test_set,
+        batch_size=1,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    args.in_channels = 3
+    args.out_channels = 1
+    args.num_classes = 1
+    return train_loader, val_loader, test_loader
+
+
+def build_Anatomy_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    if args.data20:
+        train_set = AnatomySet(args, "spilt/anatomy/train_20.csv", 'train', train_transforms)
+    else:
+        train_set = AnatomySet(args, "spilt/anatomy/train.csv", 'train', train_transforms)
+    val_set = AnatomySet(args, "spilt/anatomy/val.csv", 'val', val_transforms)
+    test_set = AnatomySet(
+        args,
+        "spilt/anatomy/test.csv",
+        'test',
+        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
+    )
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        shuffle=True,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    test_loader = DataLoader(
+        test_set,
+        batch_size=1,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    if args.prompt:
+        # TODO: need to update; currently not in use
+        args.in_channels = 3
+    else:
+        args.in_channels = 3
+    args.out_channels = 9
+    args.num_classes = 8
+    return train_loader, val_loader, test_loader
+
+
+def build_BpAnatomy_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    if args.data20:
+        train_set = BpAnatomySet(args, "spilt/anatomy/train_20.csv", 'train', train_transforms)
+    else:
+        train_set = BpAnatomySet(args, "spilt/anatomy/train.csv", 'train', train_transforms)
+    val_set = BpAnatomySet(args, "spilt/anatomy/val.csv", 'val', val_transforms)
+    test_set = BpAnatomySet(
+        args,
+        "spilt/anatomy/test.csv",
+        'test',
+        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
+    )
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=4,
+        drop_last=False,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=4,
+        drop_last=False,
+    )
+    test_loader = DataLoader(
+        test_set, batch_size=1, shuffle=False, num_workers=4, drop_last=False
+    )
+    args.in_channels = 3
+    args.out_channels = 9
+    args.num_classes = 8
+    return train_loader, val_loader, test_loader
+
+
+def build_PromisHist_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    if args.data20:
+        train_set = PromisHist(args, "spilt/promis567_hist/train_20.csv", 'train', train_transforms)
+    else:
+        train_set = PromisHist(args, "spilt/promis567_hist/train.csv", 'train', train_transforms)
+    val_set = PromisHist(args, "spilt/promis567_hist/val.csv", 'val', val_transforms)
+    test_set = PromisHist(args, "spilt/promis567_hist/test.csv", 'test', test_transforms)
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        shuffle=True,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=True,
+    )
+    test_loader = DataLoader(
+        test_set,
+        batch_size=1,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    args.in_channels = 3
+    args.out_channels = 1
+    args.num_classes = 1
+    return train_loader, val_loader, test_loader
+
+def build_PromisZone_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    train_set = PromisZone(args, "spilt/promis_zone/train.csv", 'train', train_transforms)
+    val_set = PromisZone(args, "spilt/promis_zone/val.csv", 'val', val_transforms)
+    test_set = PromisZone(args, "spilt/promis_zone/test.csv", 'test', test_transforms)
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        shuffle=True,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=True,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=True,
+    )
+    test_loader = DataLoader(
+        test_set,
+        batch_size=1,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    args.in_channels = 3
+    args.out_channels = 1
+    args.num_classes = 1
+    return train_loader, val_loader, test_loader
+
+
+def build_PromisPirads3_loader(args):
+    train_transforms, val_transforms, test_transforms = get_transforms(args)
+    if args.data20:
+        train_set = UCLSet(args, "spilt/promis_pirads3/train_15.csv", 'train', train_transforms)
+    else:
+        train_set = UCLSet(args, "spilt/promis_pirads3/train.csv", 'train', train_transforms)
+    val_set = UCLSet(args, "spilt/promis_pirads3/val.csv", 'val', val_transforms)
+    test_set = UCLSet(args, "spilt/promis_pirads3/test.csv", 'test', test_transforms)
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        shuffle=True,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=args.batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    test_loader = DataLoader(
+        test_set,
+        batch_size=1,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=14,
+        drop_last=False,
+    )
+    args.in_channels = 3
+    args.out_channels = 1
+    args.num_classes = 1
+    return train_loader, val_loader, test_loader

demo_classfication.py

@@ -0,0 +1,192 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import argparse
+import datetime
+import json
+import numpy as np
+import os
+import time
+from pathlib import Path
+from typing import Callable, List, Optional, Tuple
+import torch
+import torch.backends.cudnn as cudnn
+from models.classifier import Classifier
+from models.convnextv2 import convnextv2_tiny, remap_checkpoint_keys, load_state_dict
+from dataset.dataset_cls import build_Risk_loader, build_Screening_loader, build_Promis_loader, build_Promis3_hist_loader
+from engine.classification import test_risk
+
+def tuple_type(strings):
+    strings = strings.replace("(", "").replace(")", "")
+    mapped_int = map(int, strings.split(","))
+    return tuple(mapped_int)
+
+def get_args_parser():
+    parser = argparse.ArgumentParser("segmentation", add_help=False)
+    parser.add_argument(
+        "--batch_size",
+        default=1,
+        type=int,
+        help="Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus",
+    )
+    parser.add_argument("--epochs", default=400, type=int)
+    parser.add_argument(
+        "--root", default="./", type=str
+    )
+    parser.add_argument("--crop_spatial_size", default=(64, 256, 256), type=tuple_type)
+
+    # Model parameters
+    parser.add_argument("--model", help="model name")
+    parser.add_argument(
+        "--input_size", default=(64, 256, 256), type=tuple_type, help="images input size"
+    )
+    parser.add_argument(
+        "--train",
+        default="scratch",
+        choices=["fintune", "freeze", "scratch"],
+        help="train method",
+    )
+    parser.add_argument("--pretrain", default=None, type=str)
+    parser.add_argument("--tolerance", default=5, type=int)
+    parser.add_argument("--spacing", default=(1.0, 0.5, 0.5), type=tuple)
+    # Optimizer parameters
+    parser.add_argument(
+        "--weight_decay", type=float, default=1e-5, help="weight decay (default: 1e-5)"
+    )
+    parser.add_argument(
+        "--lr",
+        default=0.1,
+        type=float,
+        metavar="LR",
+        help="learning rate (absolute lr)",
+    )
+    parser.add_argument(
+        "--min_lr",
+        type=float,
+        default=0.0,
+        metavar="LR",
+        help="lower lr bound for cyclic schedulers that hit 0",
+    )
+    parser.add_argument(
+        "--warmup_epochs", type=int, default=40, metavar="N", help="epochs to warmup LR"
+    )
+
+    # Dataset parameters
+    parser.add_argument(
+        "--output_dir",
+        default="./outputcls",
+        help="path where to save, empty for no saving",
+    )
+    parser.add_argument("--file_name", default="")
+    parser.add_argument("--ckpt_dir", default="./outputcls")
+    parser.add_argument(
+        "--log_dir", default="./outputcls", help="path where to tensorboard log"
+    )
+    parser.add_argument("--dataset", default="UCL", help="dataset name")
+    parser.add_argument(
+        "--device", default="cuda", help="device to use for training / testing"
+    )
+    parser.add_argument("--seed", default=0, type=int)
+    parser.add_argument("--resume", default="", help="resume from checkpoint")
+
+    parser.add_argument(
+        "--start_epoch", default=0, type=int, metavar="N", help="start epoch"
+    )
+    parser.add_argument("--num_workers", default=10, type=int)
+    parser.add_argument(
+        "--pin_mem",
+        action="store_true",
+        help="Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.",
+    )
+    parser.add_argument("--no_pin_mem", action="store_false", dest="pin_mem")
+    parser.set_defaults(pin_mem=True)
+
+    parser.add_argument("--data20", action="store_true", help="Use 20 training data")
+    parser.set_defaults(data20=False)
+
+    parser.add_argument("--data_num", default=0, type=int, help="number of train data")
+
+    parser.add_argument("--save_fig", action="store_true")
+    parser.set_defaults(save_fig=False)
+
+    parser.add_argument(
+        "--prompt", action="store_true", help="Use visual prompt tuning"
+    )
+    parser.set_defaults(data20=False)
+
+    parser.add_argument(
+        "--world_size", default=1, type=int, help="number of distributed processes"
+    )
+    parser.add_argument("--local_rank", default=-1, type=int)
+    parser.add_argument("--dist_on_itp", action="store_true")
+    parser.add_argument(
+        "--dist_url", default="env://", help="url used to set up distributed training"
+    )
+    parser.add_argument("--kfold", type=int, default=None)
+    parser.add_argument("--demo", type=bool, default=True, help="Run in demo mode")
+    return parser
+
+
+def main(args):
+
+    device = "cuda"
+    # fix the seed for reproducibility
+    seed = args.seed
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    cudnn.benchmark = True
+    
+
+    if args.dataset == "risk":
+        data_loader_test = build_Risk_loader(args)
+    # elif args.dataset == "screening":
+    #     data_loader_train, data_loader_val, data_loader_test = build_Screening_loader(
+    #         args
+    #     )
+    # elif args.dataset == "promis":
+    #     data_loader_train, data_loader_val, data_loader_test = build_Promis_loader(args)
+    # elif args.dataset == "promis3hist":
+    #     data_loader_train, data_loader_val, data_loader_test = build_Promis3_hist_loader(args)
+    else:
+        raise NotImplementedError(f"unknown schedule sampler: {args.dataset}")
+    print(f"Loaded dataset: {args.dataset}, test set size: {len(data_loader_test.dataset)}")
+    
+    if args.model == "profound_conv":
+        convnext = convnextv2_tiny(in_chans=3)
+        model = Classifier(convnext, args.num_classes)
+    else:
+        raise NotImplementedError(f"unknown model: {args.model}")
+
+    args.output_dir = os.path.join(args.output_dir, args.dataset)
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    model.load_state_dict(torch.load(args.ckpt_dir, map_location='cpu', weights_only=False)["model"])
+    print(f"Loaded model from {args.ckpt_dir}")
+    model.to(device)
+    logits, gts = [], []
+    model.eval()
+    with torch.no_grad():
+        for idx, (img, gt, pid) in enumerate(data_loader_test):
+            img, gt = img.to(args.device), gt.to(args.device)
+            logit = model(img)
+            logits.append(logit)
+            gts.append(gt)
+
+        # if args.dataset == "risk":
+        #     test_risk(logits, gts)
+        logits = torch.cat(logits, 0).squeeze().cpu().numpy()
+        gts = torch.cat(gts, 0).squeeze().cpu().numpy()
+        print(f"test results: logits {logits}, gts {gts}")
+        np.savez(os.path.join(args.output_dir, f"{args.file_name}.npz"), logits = logits, gts=gts)
+
+if __name__ == "__main__":
+    args = get_args_parser()
+    args = args.parse_args()
+    main(args)

demo_segmentation.py

@@ -0,0 +1,250 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import argparse
+import datetime
+import json
+import numpy as np
+import os
+import time
+from pathlib import Path
+from typing import Callable, List, Optional, Tuple
+import torch
+import torch.backends.cudnn as cudnn
+from dataset.dataset_seg import (
+    build_UCL_loader,
+    build_Anatomy_loader,
+    build_BpAnatomy_loader,
+    build_Promis_loader,
+    build_PromisPirads3_loader
+)
+import monai
+from monai.inferers import sliding_window_inference
+from monai.metrics import compute_dice
+import SimpleITK as sitk
+from models.convnextv2 import convnextv2_tiny, remap_checkpoint_keys, load_state_dict
+from models.convnext_unter import ConvnextUNETR
+from models.upernet_module import UperNet
+
+
+def tuple_type(strings):
+    strings = strings.replace("(", "").replace(")", "")
+    mapped_int = map(int, strings.split(","))
+    return tuple(mapped_int)
+
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser("segmentation", add_help=False)
+    parser.add_argument(
+        "--batch_size",
+        default=1,
+        type=int,
+        help="Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus",
+    )
+    parser.add_argument("--epochs", default=400, type=int)
+    parser.add_argument(
+        "--root", default="./", type=str
+    )
+    parser.add_argument("--crop_spatial_size", default=(64, 256, 256), type=tuple_type)
+
+    # Model parameters
+    parser.add_argument("--model", help="model name")
+    parser.add_argument(
+        "--input_size", default=(64, 256, 256), type=tuple_type, help="images input size"
+    )
+    parser.add_argument(
+        "--train",
+        default="scratch",
+        choices=["fintune", "freeze", "scratch"],
+        help="train method",
+    )
+    parser.add_argument("--pretrain", default=None, type=str)
+    parser.add_argument("--tolerance", default=5, type=int)
+    parser.add_argument("--spacing", default=(1.0, 0.5, 0.5), type=tuple)
+    # Optimizer parameters
+    parser.add_argument(
+        "--weight_decay", type=float, default=1e-5, help="weight decay (default: 1e-5)"
+    )
+    parser.add_argument(
+        "--lr",
+        default=0.1,
+        type=float,
+        metavar="LR",
+        help="learning rate (absolute lr)",
+    )
+    parser.add_argument(
+        "--min_lr",
+        type=float,
+        default=0.0,
+        metavar="LR",
+        help="lower lr bound for cyclic schedulers that hit 0",
+    )
+    parser.add_argument(
+        "--warmup_epochs", type=int, default=40, metavar="N", help="epochs to warmup LR"
+    )
+
+    # Dataset parameters
+    parser.add_argument(
+        "--output_dir",
+        default="./outputseg",
+        help="path where to save, empty for no saving",
+    )
+    parser.add_argument("--file_name", default="")
+    parser.add_argument("--ckpt_dir", default="./outputseg")
+    parser.add_argument(
+        "--log_dir", default="./outputseg", help="path where to tensorboard log"
+    )
+    parser.add_argument("--dataset", default="UCL", help="dataset name")
+    parser.add_argument(
+        "--device", default="cuda", help="device to use for training / testing"
+    )
+    parser.add_argument("--seed", default=0, type=int)
+    parser.add_argument("--resume", default="", help="resume from checkpoint")
+
+    parser.add_argument(
+        "--start_epoch", default=0, type=int, metavar="N", help="start epoch"
+    )
+    parser.add_argument("--num_workers", default=10, type=int)
+    parser.add_argument(
+        "--pin_mem",
+        action="store_true",
+        help="Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.",
+    )
+    parser.add_argument("--no_pin_mem", action="store_false", dest="pin_mem")
+    parser.set_defaults(pin_mem=True)
+
+    parser.add_argument("--data20", action="store_true", help="Use 20 training data")
+    parser.set_defaults(data20=False)
+
+    parser.add_argument("--data_num", default=0, type=int, help="number of train data")
+
+    parser.add_argument("--save_fig", action="store_true")
+    parser.set_defaults(save_fig=False)
+
+    parser.add_argument(
+        "--prompt", action="store_true", help="Use visual prompt tuning"
+    )
+    parser.set_defaults(prompt=False)
+
+    parser.add_argument(
+        "--world_size", default=1, type=int, help="number of distributed processes"
+    )
+    parser.add_argument("--local_rank", default=-1, type=int)
+    parser.add_argument("--dist_on_itp", action="store_true")
+    parser.add_argument(
+        "--dist_url", default="env://", help="url used to set up distributed training"
+    )
+    parser.add_argument("--demo", type=bool, default=True, help="Run in demo mode")
+    return parser
+
+
+def main(args):
+
+    device = "cuda"
+    # fix the seed for reproducibility
+    seed = args.seed
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    cudnn.benchmark = True
+
+    if args.dataset == "UCL":
+        data_loader_test = build_UCL_loader(args)
+        args.sliding_window = False
+    
+    else:
+        raise NotImplementedError(f"unknown schedule sampler: {args.dataset}")
+    print(f"Loaded dataset: {args.dataset}, test set size: {len(data_loader_test)}")
+
+    if args.model == "profound_conv":
+        convnext = convnextv2_tiny(in_chans=3)
+        model = UperNet(
+            encoder=convnext,
+            in_channels=[96, 192, 384, 768],
+            out_channels=args.out_channels,
+        )
+        model = model.to(device)
+    
+    elif args.model == "profound_conv_unetr3d":
+        convnext = convnextv2_tiny(in_chans=3)
+        
+        model = ConvnextUNETR(
+            in_channels=3, out_channels=1, convnext=convnext, feature_size=32
+        )
+        model = model.to(device)
+
+    else:
+        raise NotImplementedError(f"unknown model: {args.model}")
+    
+    
+    args.output_dir = os.path.join(args.output_dir, args.dataset)
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    model.load_state_dict(torch.load(args.ckpt_dir, weights_only=False)["model"])
+    print(f"Loaded model: {args.ckpt_dir}")
+
+    dice_list = []
+    model.eval()
+    with torch.no_grad():
+        for idx, (img, gt, pid) in enumerate(data_loader_test):
+            img, gt = img.to(args.device), gt.to(args.device)
+            if args.sliding_window:
+                pred = sliding_window_inference(
+                    img, args.crop_spatial_size, 4, model, overlap=0.5
+                )
+            else:
+                pred = model(img)
+
+            if args.num_classes == 1:
+                pred = torch.sigmoid(pred) > 0.5
+                pred = pred.int()
+            else:
+                pred = torch.softmax(pred, dim=1)
+                pred = torch.argmax(pred, dim=1, keepdim=True)
+
+            dice = compute_dice(pred, gt)  # compute_dice(pred, gt, False,num_classes=9)
+            print(pid, dice.item())
+            if not torch.isnan(dice):
+                dice_list.append(dice)
+            # dice = int(dice.mean()*10000)
+            img = img.squeeze().cpu().numpy()
+            pred = pred.squeeze().cpu().numpy()
+            gt = gt.squeeze().cpu().numpy()
+            if args.save_fig:
+                if idx < 20:
+                    # print(img.shape,pred.shape, gt.shape )
+                    sitk.WriteImage(
+                        sitk.GetImageFromArray(img[0]),
+                        os.path.join(args.output_dir, f"{idx}_t2w.nii.gz"),
+                    )
+                    sitk.WriteImage(
+                        sitk.GetImageFromArray(img[1]),
+                        os.path.join(args.output_dir, f"{idx}_dwi.nii.gz"),
+                    )
+                    sitk.WriteImage(
+                        sitk.GetImageFromArray(pred),
+                        os.path.join(args.output_dir, f"{idx}_pred.nii.gz"),
+                    )
+                    sitk.WriteImage(
+                        sitk.GetImageFromArray(gt),
+                        os.path.join(args.output_dir, f"{idx}_gt.nii.gz"),
+                    )
+        dice_list = torch.stack(dice_list, 0)
+        np.save(
+            os.path.join(args.output_dir, f"{args.file_name}.npy"),
+            dice_list.cpu().numpy(),
+        )
+        print("dice mean: ", dice_list.mean().item())
+
+
+if __name__ == "__main__":
+    args = get_args_parser()
+    args = args.parse_args()
+    main(args)

engine/__init__.py

@@ -0,0 +1 @@
+# ProFound engine package 

engine/classification.py

@@ -0,0 +1,341 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import math
+import sys
+import torch
+import os
+import util.misc as misc
+import util.lr_sched as lr_sched
+import numpy as np
+from util.metric import accuracy, ConfusionMatrix, kappa
+from sklearn.metrics import (
+    roc_auc_score,
+    top_k_accuracy_score,
+    f1_score,
+    confusion_matrix,
+)
+from torchmetrics.classification import (
+    BinarySpecificityAtSensitivity,
+    BinarySensitivityAtSpecificity,
+)
+
+
+import pdb
+
+
+def train_one_epoch(
+    model,
+    data_loader,
+    optimizer,
+    device,
+    epoch: int,
+    loss_scaler,
+    log_writer=None,
+    args=None,
+):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter("lr", misc.SmoothedValue(window_size=1, fmt="{value:.6f}"))
+    header = "Epoch: [{}]".format(epoch)
+    print_freq = 20
+
+    if args.dataset == "promis":
+        loss_cal = torch.nn.BCEWithLogitsLoss()
+    else:
+        if args.num_classes > 1:
+            loss_cal = torch.nn.CrossEntropyLoss()
+        else:
+            loss_cal = torch.nn.BCEWithLogitsLoss()
+    
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print("log_dir: {}".format(log_writer.log_dir))
+    last_norm = 0.0
+    for data_iter_step, (img, gt, dataidx) in enumerate(
+        metric_logger.log_every(data_loader, print_freq, header)
+    ):
+        # we use a per iteration (instead of per epoch) lr scheduler
+        img, gt = img.to(device, non_blocking=True), gt.to(device, non_blocking=True)
+        lr_sched.adjust_learning_rate(
+            optimizer, data_iter_step / len(data_loader) + epoch, args
+        )
+        logit = model(img)
+        # print("logit: ", logit.shape, "gt: ", gt.shape, "image: ", img.shape)
+        loss = loss_cal(logit, gt)
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print(
+                "nan",
+                torch.isnan(logit).any(),
+                torch.isnan(img).any(),
+                dataidx,
+                last_norm,
+            )
+            print(
+                "inf",
+                torch.isinf(logit).any(),
+                torch.isinf(img).any(),
+                dataidx,
+                last_norm,
+            )
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        optimizer.zero_grad()
+        loss.backward()
+        # torch.nn.utils.clip_grad_norm_(model.parameters(),  1.0)
+        optimizer.step()
+
+        # last_norm = loss_scaler(loss, optimizer, parameters=model.parameters())
+        # optimizer.zero_grad()
+        # torch.cuda.synchronize()
+        metric_logger.update(loss=loss_value)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None:
+            """We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar("train_loss", loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar("lr", lr, epoch_1000x)
+
+    # gather the stats from all processes
+    # metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+def validation(model, data_loader_val, device, epoch, args):
+    model.eval()
+
+    if args.dataset == "promis":
+        loss_cal = torch.nn.BCEWithLogitsLoss()
+    else:
+        if args.num_classes > 1:
+            loss_cal = torch.nn.CrossEntropyLoss()
+        else:
+            loss_cal = torch.nn.BCEWithLogitsLoss()
+    
+    with torch.no_grad():
+        loss_summary = []
+        for idx, (img, gt, _) in enumerate(data_loader_val):
+            img, gt = img.to(device), gt.to(device)
+            mask = model(img)
+            loss = loss_cal(mask, gt)
+            loss_summary.append(loss.detach().cpu().numpy())
+            print(
+                "epoch: {}/{}, iter: {}/{}".format(
+                    epoch, args.epochs, idx, len(data_loader_val)
+                )
+                + " loss:"
+                + str(loss_summary[-1].flatten()[0])
+            )
+        avg_loss = np.mean(loss_summary)
+        print("Averaged stats:", str(avg_loss))
+    return avg_loss
+
+
+def test(model, test_loader, args):
+    filepath_best = os.path.join(args.output_dir, "best.pth.tar")
+    model.load_state_dict(torch.load(filepath_best)["model"], weights_only=False)
+    model.eval()
+    prob, gts = [], []
+    with torch.no_grad():
+        for idx, (img, gt, _) in enumerate(test_loader):
+            img, gt = img.to(args.device), gt.to(args.device)
+            logit = model(img)
+            prob.append(logit)
+            gts.append(gt)
+
+    if args.dataset == "risk":
+        return test_risk(prob, gts)
+    elif args.dataset == "screening":
+        return test_screening(prob, gts)
+    elif args.dataset == "promis":
+        return test_promis(prob, gts)
+    else:
+        raise NotImplementedError(f"unknown dataset: {args.dataset}")
+
+
+def test_risk(prob, gts):
+    log_stats = {}
+    prob = torch.cat(prob, 0)
+    prob = torch.softmax(prob, dim=-1).cpu().numpy()
+    gts = torch.cat(gts, 0).cpu().numpy()
+
+    score_acc = top_k_accuracy_score(gts, prob, k=1) * 100
+    score_qwk = kappa(gts, np.argmax(prob, 1))
+    score_auc = roc_auc_score(gts, prob, multi_class="ovr") * 100
+    score_f1 = f1_score(gts, np.argmax(prob, 1), average="macro") * 100
+
+    print("score")
+    print(f"acc\t auc \t qwk \t f1")
+    print(f"{score_acc:.2f} \t {score_auc:.2f} \t {score_qwk:.4f} \t {score_f1:.2f}")
+    log_stats["4-class_acc"] = f"{score_acc:.2f}"
+    log_stats["4-class_auc"] = f"{score_auc:.2f}"
+    log_stats["4-class_qwk"] = f"{score_qwk:.4f}"
+    log_stats["4-class_f1"] = f"{score_f1:.2f}"
+
+    # 2 3 4 5 four classes 0 1 2 3
+
+    sig_prob = np.sum(prob[:, 1:], -1)
+    sig_gts = (gts > 0).astype(int)
+    sig_acc = top_k_accuracy_score(sig_gts, sig_prob, k=1) * 100
+    sig_auc = roc_auc_score(sig_gts, sig_prob) * 100
+    sig_f1 = f1_score(sig_gts, sig_prob > 0.5) * 100
+
+    print("Pirads >=3")
+    print(f"auc \t f1 ")
+    print(f"{sig_auc:.2f} \t {sig_f1:.2f}")
+
+    log_stats["leq3_auc"]=f"{sig_auc:.2f}"
+    log_stats["leq3_f1"]=f"{sig_f1:.2f}"
+
+    for i in [0.8, 0.9]:
+        sig_spec = BinarySpecificityAtSensitivity(min_sensitivity=i, thresholds=None)
+        sig_specificity, _ = sig_spec(
+            torch.from_numpy(sig_prob), torch.from_numpy(sig_gts)
+        )
+        sig_specificity = sig_specificity * 100
+        sig_sens = BinarySensitivityAtSpecificity(min_specificity=i, thresholds=None)
+        sig_sensitivity, _ = sig_sens(
+            torch.from_numpy(sig_prob), torch.from_numpy(sig_gts)
+        )
+        sig_sensitivity = sig_sensitivity* 100
+
+        print(f"min: {i}")
+        print(f"Specificity at Sensitivity \t Sensitivity at Specificity")
+        print(f"{sig_specificity:.2f} \t {sig_sensitivity:.2f} ")
+        log_stats[f"leq3_specificity_at_{i}"]=f"{sig_specificity:.2f}"
+        log_stats[f"leq3_sensitivity_at_{i}"]=f"{sig_sensitivity:.2f}"
+
+    sig_prob = np.sum(prob[:, 2:], -1)
+    sig_gts = (gts > 1).astype(int)
+    sig_acc = top_k_accuracy_score(sig_gts, sig_prob, k=1) * 100
+    sig_auc = roc_auc_score(sig_gts, sig_prob) * 100
+    sig_f1 = f1_score(sig_gts, sig_prob > 0.5) * 100
+
+    print("Pirads >=4")
+    print(f"auc \t f1 ")
+    print(f"{sig_auc:.2f} \t {sig_f1:.2f}")
+
+    log_stats["leq4_auc"]=f"{sig_auc:.2f}"
+    log_stats["leq4_f1"]=f"{sig_f1:.2f}"
+
+    for i in [0.8, 0.9]:
+        sig_spec = BinarySpecificityAtSensitivity(min_sensitivity=i, thresholds=None)
+        sig_specificity, _ = sig_spec(
+            torch.from_numpy(sig_prob), torch.from_numpy(sig_gts)
+        )
+        sig_specificity = sig_specificity * 100
+        sig_sens = BinarySensitivityAtSpecificity(min_specificity=i, thresholds=None)
+        sig_sensitivity, _ = sig_sens(
+            torch.from_numpy(sig_prob), torch.from_numpy(sig_gts)
+        )
+        sig_sensitivity = sig_sensitivity* 100
+
+        print(f"min: {i}")
+        print(f"Specificity at Sensitivity \t Sensitivity at Specificity")
+        print(f"{sig_specificity:.2f} \t {sig_sensitivity:.2f} ")
+        log_stats[f"leq4_specificity_at_{i}"]=f"{sig_specificity:.2f}"
+        log_stats[f"leq4_sensitivity_at_{i}"]=f"{sig_sensitivity:.2f}"
+    return log_stats
+
+
+def test_screening(prob, gts):
+    prob = torch.cat(prob, 0)
+    prob = torch.sigmoid(prob).cpu().numpy()
+    gts = torch.cat(gts, 0).long().cpu().numpy()
+
+    np.savez("result.npz", gts=gts, prob=prob)
+    score_acc = top_k_accuracy_score(gts, prob, k=1) * 100
+    score_auc = roc_auc_score(gts, prob) * 100
+    score_f1 = f1_score(gts, np.argmax(prob, 1)) * 100
+
+    print(f"acc\t auc \t f1")
+    print(f"{score_acc:.2f} \t {score_auc:.2f} \t {score_f1:.2f}")
+
+    for i in [0.8, 0.9]:
+        sig_spec = BinarySpecificityAtSensitivity(min_sensitivity=i, thresholds=None)
+        sig_specificity, _ = sig_spec(torch.from_numpy(prob), torch.from_numpy(gts))
+        sig_sens = BinarySensitivityAtSpecificity(min_specificity=i, thresholds=None)
+        sig_sensitivity, _ = sig_sens(torch.from_numpy(prob), torch.from_numpy(gts))
+
+        print(f"min: {i}")
+        print(f"Specificity at Sensitivity \t Sensitivity at Specificity")
+        print(f"{sig_specificity* 100:.2f} \t {sig_sensitivity* 100:.2f} ")
+
+    log_stats = None
+    return log_stats
+
+
+
+def test_promis(prob, gts):
+    log_stats = {}
+
+    prob = torch.cat(prob, 0)
+    prob = torch.sigmoid(prob).cpu().numpy()
+    gts = torch.cat(gts, 0).cpu().numpy().astype(int)
+
+    #zone level
+    zone_prob = prob.reshape(-1)
+    zone_gt = gts.reshape(-1)
+    print(f"zone level performance")
+
+    auc = roc_auc_score(zone_prob, zone_gt) * 100
+    print(f"AUC: {auc:.2f}")
+    for i in [0.8, 0.9]:
+        sig_spec = BinarySpecificityAtSensitivity(min_sensitivity=i, thresholds=None)
+        sig_specificity, _ = sig_spec(
+            torch.from_numpy(zone_prob), torch.from_numpy(zone_gt)
+        )
+        sig_sens = BinarySensitivityAtSpecificity(min_specificity=i, thresholds=None)
+        sig_sensitivity, _ = sig_sens(
+            torch.from_numpy(zone_prob), torch.from_numpy(zone_gt)
+        )
+
+        print(f"min: {i}")
+        print(f"Specificity at Sensitivity \t Sensitivity at Specificity")
+        print(f"{sig_specificity* 100:.2f} \t {sig_sensitivity* 100:.2f} ")
+
+
+
+
+    #patient level
+    patient_prob = prob.max(-1)
+    patient_gt = gts.max(-1)
+    
+    print(f"patient level performance")
+
+    auc = roc_auc_score(patient_prob, patient_gt) * 100
+    print(f"AUC: {auc:.2f}")
+    for i in [0.8, 0.9]:
+        sig_spec = BinarySpecificityAtSensitivity(min_sensitivity=i, thresholds=None)
+        sig_specificity, _ = sig_spec(
+            torch.from_numpy(patient_prob), torch.from_numpy(patient_gt)
+        )
+        sig_sens = BinarySensitivityAtSpecificity(min_specificity=i, thresholds=None)
+        sig_sensitivity, _ = sig_sens(
+            torch.from_numpy(patient_prob), torch.from_numpy(patient_gt)
+        )
+
+        print(f"min: {i}")
+        print(f"Specificity at Sensitivity \t Sensitivity at Specificity")
+        print(f"{sig_specificity* 100:.2f} \t {sig_sensitivity* 100:.2f} ")
+
+    
+    return log_stats

engine/location.py

@@ -0,0 +1,206 @@
+import math
+import sys
+from typing import Iterable
+
+import torch
+import os
+import util.misc as misc
+import util.lr_sched as lr_sched
+from monai.losses import DiceCELoss, DiceLoss
+import numpy as np
+from monai.metrics import DiceHelper
+import surface_distance
+from surface_distance import metrics
+from util.meter import DiceMeter, HausdorffMeter, SurfaceDistanceMeter
+
+# from monai.data import ImageDataset, create_test_image_3d, decollate_batch, DataLoader
+from monai.inferers import sliding_window_inference
+from torchmetrics.classification import (
+    BinarySpecificityAtSensitivity,
+    BinarySensitivityAtSpecificity,
+)
+# from monai.metrics import DiceMetric
+# from monai.transforms import Activations
+import pdb
+from sklearn.metrics import (
+    roc_auc_score,
+    top_k_accuracy_score,
+    f1_score,
+    confusion_matrix,
+)
+
+
+def train_one_epoch(
+    model,
+    data_loader,
+    optimizer,
+    device,
+    epoch: int,
+    loss_scaler,
+    log_writer=None,
+    args=None,
+):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter("lr", misc.SmoothedValue(window_size=1, fmt="{value:.6f}"))
+    header = "Epoch: [{}]".format(epoch)
+    print_freq = 20
+
+    loss_cal = torch.nn.BCEWithLogitsLoss()
+    
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print("log_dir: {}".format(log_writer.log_dir))
+    last_norm = 0.0
+    for data_iter_step, (img, zone_mask, gt) in enumerate(
+        metric_logger.log_every(data_loader, print_freq, header)
+    ):
+        
+
+        # we use a per iteration (instead of per epoch) lr scheduler
+        img, zone_mask, gt = img.to(device, non_blocking=True), zone_mask.to(device, non_blocking=True), gt.to(device, non_blocking=True)
+        gt = gt.float()
+        lr_sched.adjust_learning_rate(
+            optimizer, data_iter_step / len(data_loader) + epoch, args
+        )
+        logit = model(img, zone_mask)
+        if isinstance(logit, list):
+            loss = loss_cal(logit[0], gt) + 0.4*loss_cal(logit[1], gt)
+        else:
+            loss = loss_cal(logit, gt)
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print(
+                "nan",
+                torch.isnan(logit).any(),
+                torch.isnan(img).any(),
+                last_norm,
+            )
+            print(
+                "inf",
+                torch.isinf(logit).any(),
+                torch.isinf(img).any(),
+                last_norm,
+            )
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        optimizer.zero_grad()
+        loss.backward()
+        # torch.nn.utils.clip_grad_norm_(model.parameters(),  1.0)
+        optimizer.step()
+
+        metric_logger.update(loss=loss_value)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None:
+            """We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar("train_loss", loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar("lr", lr, epoch_1000x)
+
+    # gather the stats from all processes
+    # metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+def validation(model, data_loader_val, device, epoch, args):
+    model.eval()
+    loss_cal = torch.nn.BCEWithLogitsLoss()
+
+    with torch.no_grad():
+        loss_summary = []
+        
+        for idx, (img, zone_mask, gt) in enumerate(data_loader_val):
+            img, zone_mask, gt = img.to(device, non_blocking=True), zone_mask.to(device, non_blocking=True), gt.to(device, non_blocking=True)
+            gt = gt.float()
+            logit = model(img, zone_mask)
+            loss = loss_cal(logit, gt)
+            loss_summary.append(loss.detach().cpu().numpy())
+            print(
+                "epoch: {}/{}, iter: {}/{}".format(
+                    epoch, args.epochs, idx, len(data_loader_val)
+                )
+                + " loss:"
+                + str(loss_summary[-1].flatten()[0])
+            )
+        avg_loss = np.mean(loss_summary)
+        print("Averaged stats:", str(avg_loss))
+    return avg_loss
+
+
+def test(model, test_loader, args, sliding_window=False):
+    model.eval()
+    filepath_best = os.path.join(args.output_dir, "best.pth.tar")
+    model.load_state_dict(torch.load(filepath_best)["model"], weights_only=False)
+   
+    log_stats = {}
+    with torch.no_grad():
+        prob, gts = [], []
+       
+        for idx, (img, zone_mask, gt) in enumerate(test_loader):
+            img, zone_mask, gt = img.to(args.device, non_blocking=True), zone_mask.to(args.device, non_blocking=True), gt.to(args.device, non_blocking=True)
+        
+            logit = model(img, zone_mask)
+            prob.append(logit)
+            gts.append(gt)
+
+            
+    prob = torch.cat(prob, 0)
+    prob = torch.sigmoid(prob).cpu()
+    gts = torch.cat(gts, 0).cpu()
+
+
+
+    print("- Zone level: ")
+    zone_prob = prob.reshape(-1, prob.shape[-1])
+    zone_gt = gts.reshape(-1, prob.shape[-1])
+    zone_auc = roc_auc_score(zone_prob, zone_gt) * 100
+    
+    for i in [0.8, 0.9]:
+        sig_spec = BinarySpecificityAtSensitivity(min_sensitivity=i, thresholds=None)
+        sig_specificity, _ = sig_spec(zone_prob, zone_gt)
+        sig_specificity = sig_specificity * 100
+        
+        sig_sens = BinarySensitivityAtSpecificity(min_specificity=i, thresholds=None)
+        sig_sensitivity, _ = sig_sens(zone_prob, zone_gt)
+        sig_sensitivity = sig_sensitivity* 100
+
+        print(f"min: {i}")
+        print(f"Specificity at Sensitivity \t Sensitivity at Specificity")
+        print(f"{sig_specificity:.2f} \t {sig_sensitivity:.2f} ")
+        log_stats[f"specificity_at_{i}"]=f"{sig_specificity:.2f}"
+        log_stats[f"sensitivity_at_{i}"]=f"{sig_sensitivity:.2f}"
+
+
+    print("- Patient level: ")
+    p_prob = prob.max(1).values
+    p_gt = gts.max(1).values
+
+    p_auc = roc_auc_score(p_prob, p_gt) * 100
+
+    for i in [0.8, 0.9]:
+        sig_spec = BinarySpecificityAtSensitivity(min_sensitivity=i, thresholds=None)
+        sig_specificity, _ = sig_spec(p_prob, p_gt)
+        sig_specificity = sig_specificity * 100
+        
+        sig_sens = BinarySensitivityAtSpecificity(min_specificity=i, thresholds=None)
+        sig_sensitivity, _ = sig_sens(p_prob, p_gt)
+        sig_sensitivity = sig_sensitivity* 100
+
+        print(f"min: {i}")
+        print(f"Specificity at Sensitivity \t Sensitivity at Specificity")
+        print(f"{sig_specificity:.2f} \t {sig_sensitivity:.2f} ")
+        log_stats[f"specificity_at_{i}"]=f"{sig_specificity:.2f}"
+        log_stats[f"sensitivity_at_{i}"]=f"{sig_sensitivity:.2f}"
+
+    return log_stats

engine/pretrain.py

@@ -0,0 +1,85 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import math
+import sys
+from typing import Iterable
+
+import torch
+
+import util.misc as misc
+import util.lr_sched as lr_sched
+
+
+def train_one_epoch(
+    model,
+    data_loader,
+    optimizer,
+    device,
+    epoch: int,
+    loss_scaler,
+    log_writer=None,
+    args=None,
+):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter("lr", misc.SmoothedValue(window_size=1, fmt="{value:.6f}"))
+    header = "Epoch: [{}]".format(epoch)
+    print_freq = 20
+
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print("log_dir: {}".format(log_writer.log_dir))
+
+    for data_iter_step, (samples, _) in enumerate(
+        metric_logger.log_every(data_loader, print_freq, header)
+    ):
+
+        # we use a per iteration (instead of per epoch) lr scheduler
+        samples = samples.to(device, non_blocking=True)
+        lr_sched.adjust_learning_rate(
+            optimizer, data_iter_step / len(data_loader) + epoch, args
+        )
+
+        # with torch.cuda.amp.autocast():
+        loss, _, _ = model(samples, mask_ratio=args.mask_ratio)
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        optimizer.zero_grad()
+        loss.backward()
+        # torch.nn.utils.clip_grad_norm_(model.parameters(),  1.0)
+        optimizer.step()
+        # loss_scaler(loss, optimizer, parameters=model.parameters(),clip_grad=1.0)
+        # optimizer.zero_grad()
+        torch.cuda.synchronize()
+        metric_logger.update(loss=loss_value)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None:
+            """We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar("train_loss", loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar("lr", lr, epoch_1000x)
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}

engine/pretrain_amp.py

@@ -0,0 +1,81 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import math
+import sys
+from typing import Iterable
+
+import torch
+
+import util.misc as misc
+import util.lr_sched as lr_sched
+
+
+def train_one_epoch(
+    model,
+    data_loader,
+    optimizer,
+    device,
+    epoch: int,
+    loss_scaler,
+    log_writer=None,
+    args=None,
+):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter("lr", misc.SmoothedValue(window_size=1, fmt="{value:.6f}"))
+    header = "Epoch: [{}]".format(epoch)
+    print_freq = 20
+
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print("log_dir: {}".format(log_writer.log_dir))
+
+    for data_iter_step, (samples, _) in enumerate(
+        metric_logger.log_every(data_loader, print_freq, header)
+    ):
+
+        # we use a per iteration (instead of per epoch) lr scheduler
+        samples = samples.to(device, non_blocking=True)
+        lr_sched.adjust_learning_rate(
+            optimizer, data_iter_step / len(data_loader) + epoch, args
+        )
+
+        with torch.cuda.amp.autocast():
+            loss, _, _ = model(samples, mask_ratio=args.mask_ratio)
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        loss_scaler(loss, optimizer, parameters=model.parameters(), clip_grad=1.0)
+        optimizer.zero_grad()
+        torch.cuda.synchronize()
+        metric_logger.update(loss=loss_value)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None:
+            """We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar("train_loss", loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar("lr", lr, epoch_1000x)
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}

engine/regression.py

@@ -0,0 +1,142 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import math
+import sys
+import torch
+import os
+import util.misc as misc
+import util.lr_sched as lr_sched
+import numpy as np
+
+
+def train_one_epoch(
+    model,
+    data_loader,
+    optimizer,
+    device,
+    epoch: int,
+    loss_scaler,
+    log_writer=None,
+    args=None,
+):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter("lr", misc.SmoothedValue(window_size=1, fmt="{value:.6f}"))
+    header = "Epoch: [{}]".format(epoch)
+    print_freq = 20
+
+    loss_cal = torch.nn.MSELoss()
+
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print("log_dir: {}".format(log_writer.log_dir))
+    last_norm = 0.0
+    for data_iter_step, (img, gt, dataidx) in enumerate(
+        metric_logger.log_every(data_loader, print_freq, header)
+    ):
+        # we use a per iteration (instead of per epoch) lr scheduler
+        img, gt = img.to(device, non_blocking=True), gt.to(device, non_blocking=True)
+        lr_sched.adjust_learning_rate(
+            optimizer, data_iter_step / len(data_loader) + epoch, args
+        )
+        logit = model(img)
+        loss = loss_cal(logit, gt)
+        loss_value = loss.item()
+        if not math.isfinite(loss_value):
+            print(
+                "nan",
+                torch.isnan(logit).any(),
+                torch.isnan(img).any(),
+                dataidx,
+                last_norm,
+            )
+            print(
+                "inf",
+                torch.isinf(logit).any(),
+                torch.isinf(img).any(),
+                dataidx,
+                last_norm,
+            )
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        optimizer.zero_grad()
+        loss.backward()
+        # torch.nn.utils.clip_grad_norm_(model.parameters(),  1.0)
+        optimizer.step()
+
+        # last_norm = loss_scaler(loss, optimizer, parameters=model.parameters())
+        # optimizer.zero_grad()
+        # torch.cuda.synchronize()
+        metric_logger.update(loss=loss_value)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None:
+            """We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar("train_loss", loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar("lr", lr, epoch_1000x)
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+def validation(model, data_loader_val, device, epoch, args):
+    model.eval()
+    loss_cal = torch.nn.MSELoss()
+    with torch.no_grad():
+        loss_summary = []
+        for idx, (img, gt, _) in enumerate(data_loader_val):
+            img, gt = img.to(device), gt.to(device)
+            loss = loss_cal(model(img), gt)
+            loss_summary.append(loss.detach().cpu().numpy())
+            print(
+                "epoch: {}/{}, iter: {}/{}".format(
+                    epoch, args.epochs, idx, len(data_loader_val)
+                )
+                + " loss:"
+                + str(loss_summary[-1].flatten()[0])
+            )
+        avg_loss = np.mean(loss_summary)
+        print("Averaged stats:", str(avg_loss))
+    return avg_loss
+
+
+def test(model, test_loader, args):
+    filepath_best = os.path.join(args.output_dir, "best.pth.tar")
+    model.load_state_dict(torch.load(filepath_best)["model"], weights_only=False)
+
+    model.eval()
+    log_stats = {}
+    pred, gts = [], []
+
+    with torch.no_grad():
+        for idx, (img, gt, _) in enumerate(test_loader):
+            img, gt = img.to(args.device), gt.to(args.device)
+            pred.append(model(img))
+            gts.append(gt)
+        pred = torch.cat(pred, 0)
+        gts = torch.cat(gts, 0)
+        pred = pred * 500000 + 70000
+        gts = gts * 500000 + 70000
+        mse = torch.nn.MSELoss()(pred, gts)
+        mae = torch.nn.L1Loss()(pred, gts)
+    print("MSE", mse.item(), "MAE", mae.item())
+    log_stats = {"MSE": mse.item(), "MAE": mae.item()}
+    return log_stats

engine/segment.py

@@ -0,0 +1,199 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import math
+import sys
+from typing import Iterable
+
+import torch
+import os
+import util.misc as misc
+import util.lr_sched as lr_sched
+from monai.losses import DiceCELoss, DiceLoss
+import numpy as np
+from monai.metrics import DiceHelper
+import surface_distance
+from surface_distance import metrics
+from util.meter import DiceMeter, HausdorffMeter, SurfaceDistanceMeter
+
+# from monai.data import ImageDataset, create_test_image_3d, decollate_batch, DataLoader
+from monai.inferers import sliding_window_inference
+
+# from monai.metrics import DiceMetric
+# from monai.transforms import Activations
+import pdb
+
+
+def train_one_epoch(
+    model,
+    data_loader,
+    optimizer,
+    device,
+    epoch: int,
+    loss_scaler,
+    log_writer=None,
+    args=None,
+):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter("lr", misc.SmoothedValue(window_size=1, fmt="{value:.6f}"))
+    header = "Epoch: [{}]".format(epoch)
+    print_freq = 20
+
+    if args.out_channels == 1:
+        loss_cal = DiceCELoss(sigmoid=True)
+    else:
+        loss_cal = DiceCELoss(to_onehot_y=True, softmax=True, include_background=False)
+
+    
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print("log_dir: {}".format(log_writer.log_dir))
+    last_norm = 0.0
+    for data_iter_step, (img, gt, dataidx) in enumerate(
+        metric_logger.log_every(data_loader, print_freq, header)
+    ):
+        # we use a per iteration (instead of per epoch) lr scheduler
+        img, gt = img.to(device, non_blocking=True), gt.to(device, non_blocking=True)
+        lr_sched.adjust_learning_rate(
+            optimizer, data_iter_step / len(data_loader) + epoch, args
+        )
+        # print(img.shape, img.mean(), img.std())
+        # with torch.cuda.amp.autocast():
+        logit = model(img)       
+        if isinstance(logit, list):
+            loss = loss_cal(logit[0], gt) + 0.4*loss_cal(logit[1], gt)
+        else:
+            loss = loss_cal(logit, gt)
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print(
+                "nan",
+                torch.isnan(logit).any(),
+                torch.isnan(img).any(),
+                dataidx,
+                last_norm,
+            )
+            print(
+                "inf",
+                torch.isinf(logit).any(),
+                torch.isinf(img).any(),
+                dataidx,
+                last_norm,
+            )
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        optimizer.zero_grad()
+        loss.backward()
+        # torch.nn.utils.clip_grad_norm_(model.parameters(),  1.0)
+        optimizer.step()
+
+        # last_norm = loss_scaler(loss, optimizer, parameters=model.parameters())
+        # optimizer.zero_grad()
+        # torch.cuda.synchronize()
+        metric_logger.update(loss=loss_value)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None:
+            """We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar("train_loss", loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar("lr", lr, epoch_1000x)
+
+    # gather the stats from all processes
+    # metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+def validation(model, data_loader_val, device, epoch, args):
+    model.eval()
+    if args.out_channels == 1:
+        dice_loss = DiceLoss(sigmoid=True)
+    else:
+        dice_loss = DiceLoss(to_onehot_y=True, softmax=True, include_background=False)
+
+    with torch.no_grad():
+        loss_summary = []
+        for idx, (img, gt, _) in enumerate(data_loader_val):
+            img, gt = img.to(device), gt.to(device)
+            mask = model(img)
+            loss = dice_loss(mask, gt)
+            loss_summary.append(loss.detach().cpu().numpy())
+            print(
+                "epoch: {}/{}, iter: {}/{}".format(
+                    epoch, args.epochs, idx, len(data_loader_val)
+                )
+                + " loss:"
+                + str(loss_summary[-1].flatten()[0])
+            )
+        avg_loss = np.mean(loss_summary)
+        print("Averaged stats:", str(avg_loss))
+    return avg_loss
+
+
+def test(model, test_loader, args, sliding_window=False):
+    model.eval()
+    filepath_best = os.path.join(args.output_dir, "best.pth.tar")
+    model.load_state_dict(torch.load(filepath_best)["model"], weights_only=False)
+    dice_meter = DiceMeter(args)
+    hausdorff_meter = HausdorffMeter(args)
+    sd_meter = SurfaceDistanceMeter(args)
+    log_stats = {}
+    with torch.no_grad():
+        for idx, (img, gt, _) in enumerate(test_loader):
+            img, gt = img.to(args.device), gt.to(args.device)
+            if sliding_window:
+                pred = sliding_window_inference(
+                    img, args.crop_spatial_size, 4, model, overlap=0.5
+                )
+            else:
+                pred = model(img)
+            if args.num_classes == 1:
+                pred = torch.sigmoid(pred) > 0.5
+            else:
+                pred = torch.softmax(pred, dim=1)
+                pred = torch.argmax(pred, dim=1, keepdim=True)
+            dice_meter.update(pred, gt)
+            hausdorff_meter.update(pred, gt)
+            sd_meter.update(pred, gt)
+
+    print("- Test metrics Dice: ")
+    dice_class_avg, dice_avg = dice_meter.get_average()
+    print("Class wise: ", dice_class_avg)
+    print("Avg.: ", dice_avg)
+
+    print("- Test metrics Hausdorff95: ")
+    hsd_class_avg, hsd_avg = hausdorff_meter.get_average()
+    print("Class wise: ", hsd_class_avg)
+    print("Avg.: ", hsd_avg)
+
+    print("- Test metrics SurfaceDistance: ")
+    sd_class_avg, sd_avg = sd_meter.get_average()
+    print("Class wise: ", sd_class_avg)
+    print("Avg.: ", sd_avg)
+    log_stats = {
+        "dice_class_avg": dice_class_avg.tolist() if isinstance(dice_class_avg, np.ndarray) else dice_class_avg,
+        "dice_avg": dice_avg.tolist() if isinstance(dice_avg, np.ndarray) else dice_avg,
+        "hsd_class_avg": hsd_class_avg.tolist() if isinstance(hsd_class_avg, np.ndarray) else hsd_class_avg,
+        "hsd_avg": hsd_avg.tolist() if isinstance(hsd_avg, np.ndarray) else hsd_avg,
+        "sd_class_avg": sd_class_avg.tolist() if isinstance(sd_class_avg, np.ndarray) else sd_class_avg,
+        "sd_avg": sd_avg.tolist() if isinstance(sd_avg, np.ndarray) else sd_avg,
+    }
+    return log_stats

models/__init__.py

@@ -0,0 +1 @@
+# ProFound models package 

models/build_classification.py

@@ -0,0 +1,83 @@
+from models.classifier import Classifier
+from models.convnextv2 import convnextv2_tiny, remap_checkpoint_keys, load_state_dict
+from util.lars import LARS
+import torch
+import os
+from util.convnext_optim import get_parameter_groups, LayerDecayValueAssigner
+
+def build_model(args, device):
+    if args.model == "profound_conv":
+        convnext = convnextv2_tiny(in_chans=3, drop_path_rate=0.1)
+        if args.pretrain is None:
+            raise NotImplementedError(f"No pretrained weight")
+        if not os.path.exists(args.pretrain):
+            raise FileExistsError(f"{args.pretrain} Not exists")
+        ckpt = torch.load(args.pretrain, map_location="cpu")
+        ckpt = remap_checkpoint_keys(ckpt)
+        load_state_dict(convnext, ckpt, weights_only=False)
+        model = Classifier(convnext, args.num_classes)
+        model = model.to(device)
+        if args.train == "freeze":
+            for key, value in model.encoder.named_parameters():
+                value.requires_grad = False
+            optimizer = LARS(model.head.parameters(), weight_decay=0, lr=args.lr)
+        else:
+            num_layers = sum(convnext.depths)
+            assigner = LayerDecayValueAssigner(
+                list(
+                    args.layer_decay ** (num_layers + 1 - i) for i in range(num_layers + 2)
+                ),
+                depths=convnext.depths,
+                layer_decay_type=args.layer_decay_type,
+            )
+
+            skip = {}
+            if hasattr(model.encoder, "no_weight_decay"):
+                skip = model.encoder.no_weight_decay()
+
+            backbone_param_groups = get_parameter_groups(
+                model.encoder,
+                args.weight_decay,
+                skip,
+                assigner.get_layer_id,
+                assigner.get_scale,
+            )
+            decoder_param_groups = [
+                {"params": model.head.parameters(), "weight_decay": 0.0, "lr": args.lr}
+            ]
+
+            optimizer = torch.optim.AdamW(
+                backbone_param_groups + decoder_param_groups, lr=args.lr
+            )
+
+    else:
+        raise NotImplementedError(f"unknown model: {args.model}")
+
+    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    print("Model = %s" % str(model))
+    print("number of params (M): %.2f" % (n_parameters / 1.0e6))
+
+    return model, optimizer
+
+
+def vit_backbone_parameters(
+    model: torch.nn.Module, weight_decay=1e-5, no_weight_decay_list=(), lr=1e-3
+):
+    no_weight_decay_list = set(no_weight_decay_list)
+    decay = []
+    no_decay = []
+
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+
+        if param.ndim <= 1 or name.endswith(".bias") or name in no_weight_decay_list:
+            no_decay.append(param)
+        else:
+            decay.append(param)
+
+    return [
+        {"params": no_decay, "weight_decay": 0.0, "lr": lr},
+        {"params": decay, "weight_decay": weight_decay, "lr": lr},
+    ]

models/classifier.py

@@ -0,0 +1,23 @@
+import torch
+import torch.nn as nn
+
+
+class Classifier(nn.Module):
+    def __init__(self, encoder, num_classes, bottleneck_dim=256):
+        super().__init__()
+        self.encoder = encoder
+        self.embed_dim = self.encoder.embed_dim
+        self.head = torch.nn.Sequential(
+            nn.Linear(self.embed_dim, bottleneck_dim),
+            nn.BatchNorm1d(bottleneck_dim),
+            nn.ReLU(),
+            nn.Linear(bottleneck_dim, num_classes)
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        if type(x) == tuple:
+            x = x[0]
+        x = self.head(x)
+        return x
+

models/convnext_unter.py

@@ -0,0 +1,182 @@
+import torch.nn.functional as F
+from typing import Sequence, Tuple, Union
+import torch
+import torch.nn as nn
+from monai.networks.blocks.dynunet_block import UnetOutBlock
+from monai.networks.blocks.unetr_block import (
+    UnetrBasicBlock,
+    UnetrPrUpBlock,
+    UnetrUpBlock,
+)
+from models.util import LayerNorm
+
+
+class ConvnextUNETR_Decoder(nn.Module):
+    """
+    UNETR based on: "Hatamizadeh et al.,
+    UNETR: Transformers for 3D Medical Image Segmentation <https://arxiv.org/abs/2103.10504>"
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        feature_size: int = 16,
+        norm_name: Union[Tuple, str] = "instance",
+        conv_block: bool = True,
+        res_block: bool = True,
+        spatial_dims: int = 3,
+        hidden_size = [96, 192, 384, 768]
+    ) -> None:
+
+        super().__init__()
+
+        self.encoder1 = UnetrBasicBlock(
+            spatial_dims=spatial_dims,
+            in_channels=in_channels,
+            out_channels=feature_size,
+            kernel_size=3,
+            stride=1,
+            norm_name=norm_name,
+            res_block=res_block,
+        )
+        self.encoder2 = UnetrPrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=hidden_size[0],
+            out_channels=feature_size * 2,
+            num_layer=0,
+            kernel_size=3,
+            stride=1,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            conv_block=conv_block,
+            res_block=res_block,
+        )
+        self.encoder3 = UnetrPrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=hidden_size[1],
+            out_channels=feature_size * 4,
+            num_layer=0,
+            kernel_size=3,
+            stride=1,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            conv_block=conv_block,
+            res_block=res_block,
+        )
+        self.encoder4 = UnetrPrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=hidden_size[2],
+            out_channels=feature_size * 8,
+            num_layer=0,
+            kernel_size=3,
+            stride=1,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            conv_block=conv_block,
+            res_block=res_block,
+        )
+        self.decoder5 = UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=hidden_size[3],
+            out_channels=feature_size * 8,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            res_block=res_block,
+        )
+        self.decoder4 = UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=feature_size * 8,
+            out_channels=feature_size * 4,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            res_block=res_block,
+        )
+        self.decoder3 = UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=feature_size * 4,
+            out_channels=feature_size * 2,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            res_block=res_block,
+        )
+        self.decoder2 = UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=feature_size * 2,
+            out_channels=feature_size,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name=norm_name,
+            res_block=res_block,
+        )
+        self.out = UnetOutBlock(
+            spatial_dims=spatial_dims,
+            in_channels=feature_size,
+            out_channels=out_channels,
+        )
+
+    def forward(self, x, x1, x2, x3, x4):
+        enc1 = self.encoder1(x)
+        enc2 = self.encoder2(x1)
+        enc3 = self.encoder3(x2)
+        enc4 = self.encoder4(x3)
+        dec3 = self.decoder5(x4, enc4)
+        dec2 = self.decoder4(dec3, enc3)
+        dec1 = self.decoder3(dec2, enc2)
+        out = self.decoder2(dec1, enc1)
+        mask = self.out(out)
+        return mask
+
+
+class ConvnextUNETR(nn.Module):
+    """
+    UNETR based on: "Hatamizadeh et al.,
+    UNETR: Transformers for 3D Medical Image Segmentation <https://arxiv.org/abs/2103.10504>"
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        convnext,
+        feature_size: int = 16,
+        norm_name: Union[Tuple, str] = "instance",
+        conv_block: bool = True,
+        res_block: bool = True,
+        spatial_dims: int = 3,
+        hidden_size = [96, 192, 384, 768]
+    ) -> None:
+        
+        super().__init__()
+
+        self.encoder = convnext
+
+        self.norm1 = LayerNorm(hidden_size[0], eps=1e-6, data_format="channels_first")
+        self.norm2 = LayerNorm(hidden_size[1], eps=1e-6, data_format="channels_first")
+        self.norm3 = LayerNorm(hidden_size[2], eps=1e-6, data_format="channels_first")
+
+        self.decoder = ConvnextUNETR_Decoder(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            feature_size=feature_size,
+            norm_name=norm_name,
+            conv_block=conv_block,
+            res_block=res_block,
+            spatial_dims=spatial_dims,
+            hidden_size=hidden_size
+        )
+
+    def forward(self, x):
+        _, hidden_states_out = self.encoder(x, ret_hids=True)
+        x1, x2, x3, x4 = hidden_states_out
+        x1 = self.norm1(x1)
+        x2 = self.norm2(x2)
+        x3 = self.norm3(x3)
+        x4 = x4.permute(0, 2, 3, 4, 1)  # (N, C, H, W, D) -> (N, H, W, D, C)
+        x4 = self.encoder.norm(x4)
+        x4 = x4.permute(0, 4, 1, 2, 3)
+        mask = self.decoder(x, x1, x2, x3, x4)
+        return mask

models/convnextv2.py

@@ -0,0 +1,311 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import trunc_normal_, DropPath
+from models.util import LayerNorm, GRN
+from collections import OrderedDict
+import math
+
+
+class Block(nn.Module):
+    """ConvNeXtV2 Block.
+
+    Args:
+        dim (int): Number of input channels.
+        drop_path (float): Stochastic depth rate. Default: 0.0
+    """
+
+    def __init__(self, dim, drop_path=0.0):
+        super().__init__()
+        self.dwconv = nn.Conv3d(
+            dim, dim, kernel_size=7, padding=3, groups=dim
+        )  # depthwise conv
+        self.norm = LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(
+            dim, 4 * dim
+        )  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.grn = GRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 4, 1)  # (N, C, H, W, D) -> (N, H, W, D, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        x = x.permute(0, 4, 1, 2, 3)  # (N, H, W, D, C) -> (N, C, H, W, D)
+        x = input + self.drop_path(x)
+        return x
+
+
+class ConvNeXtV2(nn.Module):
+    """ConvNeXt V2
+
+    Args:
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(
+        self,
+        in_chans=3,
+        depths=[3, 3, 9, 3],
+        dims=[96, 192, 384, 768],
+        drop_path_rate=0.0,
+    ):
+        super().__init__()
+        self.depths = depths
+        self.downsample_layers = (
+            nn.ModuleList()
+        )  # stem and 3 intermediate downsampling conv layers
+        stem = nn.Sequential(
+            nn.Conv3d(in_chans, dims[0], kernel_size=4, stride=4),
+            LayerNorm(dims[0], eps=1e-6, data_format="channels_first"),
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            if i == 2:
+                stride = 1
+            else:
+                stride = 2
+            downsample_layer = nn.Sequential(
+                LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
+                nn.Conv3d(dims[i], dims[i + 1], kernel_size=stride, stride=stride),
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = (
+            nn.ModuleList()
+        )  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(4):
+            stage = nn.Sequential(
+                *[
+                    Block(dim=dims[i], drop_path=dp_rates[cur + j])
+                    for j in range(depths[i])
+                ]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        self.norm = nn.LayerNorm(dims[-1], eps=1e-6)  # final norm layer
+        # self.head = nn.Linear(dims[-1], num_classes)
+
+        self.apply(self._init_weights)
+        # self.head.weight.data.mul_(head_init_scale)
+        # self.head.bias.data.mul_(head_init_scale)
+        self.embed_dim = dims[-1]
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv3d, nn.Linear)):
+            trunc_normal_(m.weight, std=0.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward_features(self, x):
+        hidden_states_out = []
+        for i in range(4):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+            hidden_states_out.append(x)
+        return self.norm(x.mean([-3, -2, -1])), hidden_states_out  # global average pooling, (N, C, H, W, D) -> (N, C)
+
+    def forward(self, x, ret_hids=False):
+        x, hidden_states_out = self.forward_features(x)
+        if ret_hids:
+            return x, hidden_states_out
+        else:
+            return x
+
+
+
+def convnextv2_atto(**kwargs):
+    model = ConvNeXtV2(depths=[2, 2, 6, 2], dims=[40, 80, 160, 320], **kwargs)
+    return model
+
+
+def convnextv2_femto(**kwargs):
+    model = ConvNeXtV2(depths=[2, 2, 6, 2], dims=[48, 96, 192, 384], **kwargs)
+    return model
+
+
+def convnext_pico(**kwargs):
+    model = ConvNeXtV2(depths=[2, 2, 6, 2], dims=[64, 128, 256, 512], **kwargs)
+    return model
+
+
+def convnextv2_nano(**kwargs):
+    model = ConvNeXtV2(depths=[2, 2, 8, 2], dims=[80, 160, 320, 640], **kwargs)
+    return model
+
+
+def convnextv2_tiny(**kwargs):
+    model = ConvNeXtV2(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
+    return model
+
+
+def convnextv2_base(**kwargs):
+    model = ConvNeXtV2(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
+    return model
+
+
+def convnextv2_large(**kwargs):
+    model = ConvNeXtV2(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
+    return model
+
+
+def convnextv2_huge(**kwargs):
+    model = ConvNeXtV2(depths=[3, 3, 27, 3], dims=[352, 704, 1408, 2816], **kwargs)
+    return model
+
+
+def remap_checkpoint_keys(ckpt):
+    new_ckpt = OrderedDict()
+    ckpt = ckpt["model"]
+
+    checkpoint_model_keys = list(ckpt.keys())
+    for k in checkpoint_model_keys:
+        if "decoder" in k or "mask_token" in k or "proj" in k or "pred" in k:
+            print(f"Removing key {k} from pretrained checkpoint")
+            del ckpt[k]
+
+    for k, v in ckpt.items():
+        if k.startswith("encoder"):
+            k = ".".join(k.split(".")[1:])  # remove encoder in the name
+        if k.endswith("kernel"):
+            k = ".".join(k.split(".")[:-1])  # remove kernel in the name
+            new_k = k + ".weight"
+            if len(v.shape) == 3:  # resahpe standard convolution
+                kv, in_dim, out_dim = v.shape
+                # ks = int(math.sqrt(kv))
+                # # pow(kv, 1/3)
+                # new_ckpt[new_k] = v.permute(2, 1, 0).\
+                #     reshape(out_dim, in_dim, ks, ks).transpose(3, 2)
+                ks = int(
+                    round(kv ** (1 / 3))
+                )  # calculate kernel size assuming cubic kernel
+                new_ckpt[new_k] = (
+                    v.permute(2, 1, 0)
+                    .reshape(out_dim, in_dim, ks, ks, ks)
+                    .permute(0, 1, 4, 3, 2)
+                )
+            elif len(v.shape) == 2:  # reshape depthwise convolution
+                kv, dim = v.shape
+                # ks = int(math.sqrt(kv))
+                # new_ckpt[new_k] = v.permute(1, 0).\
+                #     reshape(dim, 1, ks, ks).transpose(3, 2)
+                if new_k == "downsample_layers.3.1.weight":
+                    new_ckpt[new_k] = (
+                        v.permute(1, 0).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+                    )
+                else:
+                    ks = int(round(kv ** (1 / 3)))
+                    new_ckpt[new_k] = (
+                        v.permute(1, 0)
+                        .reshape(dim, 1, ks, ks, ks)
+                        .permute(0, 1, 4, 3, 2)
+                    )
+            continue
+        elif "ln" in k or "linear" in k:
+            k = k.split(".")
+            k.pop(-2)  # remove ln and linear in the name
+            new_k = ".".join(k)
+        else:
+            new_k = k
+        new_ckpt[new_k] = v
+
+    # reshape grn affine parameters and biases
+    for k, v in new_ckpt.items():
+        if k.endswith("bias") and len(v.shape) != 1:
+            new_ckpt[k] = v.reshape(-1)
+        elif "grn" in k:
+            new_ckpt[k] = v.unsqueeze(0).unsqueeze(1).unsqueeze(0)
+    return new_ckpt
+
+
+def load_state_dict(
+    model, state_dict, prefix="", ignore_missing="relative_position_index"
+):
+    missing_keys = []
+    unexpected_keys = []
+    error_msgs = []
+    # copy state_dict so _load_from_state_dict can modify it
+    metadata = getattr(state_dict, "_metadata", None)
+    state_dict = state_dict.copy()
+    if metadata is not None:
+        state_dict._metadata = metadata
+
+    def load(module, prefix=""):
+        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+        module._load_from_state_dict(
+            state_dict,
+            prefix,
+            local_metadata,
+            True,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
+        for name, child in module._modules.items():
+            if child is not None:
+                load(child, prefix + name + ".")
+
+    load(model, prefix=prefix)
+
+    warn_missing_keys = []
+    ignore_missing_keys = []
+    for key in missing_keys:
+        keep_flag = True
+        for ignore_key in ignore_missing.split("|"):
+            if ignore_key in key:
+                keep_flag = False
+                break
+        if keep_flag:
+            warn_missing_keys.append(key)
+        else:
+            ignore_missing_keys.append(key)
+
+    missing_keys = warn_missing_keys
+
+    if len(missing_keys) > 0:
+        print(
+            "Weights of {} not initialized from pretrained model: {}".format(
+                model.__class__.__name__, missing_keys
+            )
+        )
+    if len(unexpected_keys) > 0:
+        print(
+            "Weights from pretrained model not used in {}: {}".format(
+                model.__class__.__name__, unexpected_keys
+            )
+        )
+    if len(ignore_missing_keys) > 0:
+        print(
+            "Ignored weights of {} not initialized from pretrained model: {}".format(
+                model.__class__.__name__, ignore_missing_keys
+            )
+        )
+    if len(error_msgs) > 0:
+        print("\n".join(error_msgs))
+
+
+# if __name__ == 'main':
+#     model = convnextv2_base().cuda()
+#     x = torch.rand(1,3,256,256,32).cuda()
+#     print(model(x).shape)

models/upernet_module.py

@@ -0,0 +1,451 @@
+from typing import List, Optional, Tuple, Union
+import torch
+from torch import nn
+from models.util import LayerNorm, GRN
+
+class UperNetConvModule(nn.Module):
+    """
+    A convolutional block that bundles conv/norm/activation layers. This block simplifies the usage of convolution
+    layers, which are commonly used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU).
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, int]],
+        padding: Union[int, Tuple[int, int], str] = 0,
+        bias: bool = False,
+        dilation: Union[int, Tuple[int, int]] = 1,
+    ) -> None:
+        super().__init__()
+        self.conv = nn.Conv3d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            padding=padding,
+            bias=bias,
+            dilation=dilation,
+        )
+        self.batch_norm = LayerNorm(out_channels, eps=1e-6, data_format="channels_first") # nn.BatchNorm3d(out_channels)
+        self.activation = nn.GELU()
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        output = self.conv(input)
+        output = self.batch_norm(output)
+        output = self.activation(output)
+
+        return output
+
+
+class UperNetPyramidPoolingBlock(nn.Module):
+    def __init__(self, pool_scale: int, in_channels: int, channels: int) -> None:
+        super().__init__()
+        self.layers = [
+            nn.AdaptiveAvgPool3d(pool_scale),
+            UperNetConvModule(in_channels, channels, kernel_size=1),
+        ]
+        for i, layer in enumerate(self.layers):
+            self.add_module(str(i), layer)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        hidden_state = input
+        for layer in self.layers:
+            hidden_state = layer(hidden_state)
+        return hidden_state
+
+
+class UperNetPyramidPoolingModule(nn.Module):
+    """
+    Pyramid Pooling Module (PPM) used in PSPNet.
+
+    Args:
+        pool_scales (`Tuple[int]`):
+            Pooling scales used in Pooling Pyramid Module.
+        in_channels (`int`):
+            Input channels.
+        channels (`int`):
+            Channels after modules, before conv_seg.
+        align_corners (`bool`):
+            align_corners argument of F.interpolate.
+    """
+
+    def __init__(
+        self,
+        pool_scales: Tuple[int, ...],
+        in_channels: int,
+        channels: int,
+        align_corners: bool,
+    ) -> None:
+        super().__init__()
+        self.pool_scales = pool_scales
+        self.align_corners = align_corners
+        self.in_channels = in_channels
+        self.channels = channels
+        self.blocks = []
+        for i, pool_scale in enumerate(pool_scales):
+            block = UperNetPyramidPoolingBlock(
+                pool_scale=pool_scale, in_channels=in_channels, channels=channels
+            )
+            self.blocks.append(block)
+            self.add_module(str(i), block)
+
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        ppm_outs = []
+        for ppm in self.blocks:
+            ppm_out = ppm(x)
+            upsampled_ppm_out = nn.functional.interpolate(
+                ppm_out,
+                size=x.size()[2:],
+                mode="trilinear",
+                align_corners=self.align_corners,
+            )
+            ppm_outs.append(upsampled_ppm_out)
+        return ppm_outs
+
+
+class UperNetHead(nn.Module):
+    """
+    Unified Perceptual Parsing for Scene Understanding. This head is the implementation of
+    [UPerNet](https://arxiv.org/abs/1807.10221).
+    """
+
+    def __init__(self, in_channels, pool_scales, hidden_size, out_channels):
+        super().__init__()
+        self.pool_scales = pool_scales  # e.g. (1, 2, 3, 6)
+        self.in_channels = in_channels
+        self.channels = hidden_size
+        self.align_corners = False
+        self.classifier = nn.Conv3d(self.channels, out_channels, kernel_size=1)
+
+        # PSP Module
+        self.psp_modules = UperNetPyramidPoolingModule(
+            self.pool_scales,
+            self.in_channels[-1],
+            self.channels,
+            align_corners=self.align_corners,
+        )
+        self.bottleneck = UperNetConvModule(
+            self.in_channels[-1] + len(self.pool_scales) * self.channels,
+            self.channels,
+            kernel_size=3,
+            padding=1,
+        )
+        # FPN Module
+        self.lateral_convs = nn.ModuleList()
+        self.fpn_convs = nn.ModuleList()
+        for in_channels in self.in_channels[:-1]:  # skip the top layer
+            l_conv = UperNetConvModule(in_channels, self.channels, kernel_size=1)
+            fpn_conv = UperNetConvModule(
+                self.channels, self.channels, kernel_size=3, padding=1
+            )
+            self.lateral_convs.append(l_conv)
+            self.fpn_convs.append(fpn_conv)
+
+        self.fpn_bottleneck = UperNetConvModule(
+            len(self.in_channels) * self.channels,
+            self.channels,
+            kernel_size=3,
+            padding=1,
+        )
+
+    def init_weights(self):
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Conv3d):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+    def psp_forward(self, inputs):
+        x = inputs[-1]
+        psp_outs = [x]
+        psp_outs.extend(self.psp_modules(x))
+        psp_outs = torch.cat(psp_outs, dim=1)
+        output = self.bottleneck(psp_outs)
+
+        return output
+
+    def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+        # build laterals
+        laterals = [
+            lateral_conv(encoder_hidden_states[i])
+            for i, lateral_conv in enumerate(self.lateral_convs)
+        ]
+
+        laterals.append(self.psp_forward(encoder_hidden_states))
+
+        # build top-down path
+        used_backbone_levels = len(laterals)
+        for i in range(used_backbone_levels - 1, 0, -1):
+            prev_shape = laterals[i - 1].shape[2:]
+            laterals[i - 1] = laterals[i - 1] + nn.functional.interpolate(
+                laterals[i],
+                size=prev_shape,
+                mode="trilinear",
+                align_corners=self.align_corners,
+            )
+
+        # build outputs
+        fpn_outs = [
+            self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels - 1)
+        ]
+        # append psp feature
+        fpn_outs.append(laterals[-1])
+
+        for i in range(used_backbone_levels - 1, 0, -1):
+            fpn_outs[i] = nn.functional.interpolate(
+                fpn_outs[i],
+                size=fpn_outs[0].shape[2:],
+                mode="trilinear",
+                align_corners=self.align_corners,
+            )
+        fpn_outs = torch.cat(fpn_outs, dim=1)
+        output = self.fpn_bottleneck(fpn_outs)
+        output = self.classifier(output)
+
+        return output
+
+
+class UperNetFCNHead(nn.Module):
+    """
+    Fully Convolution Networks for Semantic Segmentation. This head is the implementation of
+    [FCNNet](https://arxiv.org/abs/1411.4038>).
+
+    Args:
+        in_channels (int):
+            Number of input channels.
+        kernel_size (int):
+            The kernel size for convs in the head. Default: 3.
+        dilation (int):
+            The dilation rate for convs in the head. Default: 1.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_size,
+        num_convs,
+        out_channels,
+        concat_input=False,
+        in_index: int = 2,
+        kernel_size: int = 3,
+        dilation: Union[int, Tuple[int, int]] = 1,
+    ) -> None:
+        super().__init__()
+
+        self.in_channels = in_channels[in_index]
+        self.channels = hidden_size
+        self.num_convs = num_convs
+        self.concat_input = concat_input
+        self.in_index = in_index
+
+        conv_padding = (kernel_size // 2) * dilation
+        convs = []
+        convs.append(
+            UperNetConvModule(
+                self.in_channels,
+                self.channels,
+                kernel_size=kernel_size,
+                padding=conv_padding,
+                dilation=dilation,
+            )
+        )
+        for i in range(self.num_convs - 1):
+            convs.append(
+                UperNetConvModule(
+                    self.channels,
+                    self.channels,
+                    kernel_size=kernel_size,
+                    padding=conv_padding,
+                    dilation=dilation,
+                )
+            )
+        if self.num_convs == 0:
+            self.convs = nn.Identity()
+        else:
+            self.convs = nn.Sequential(*convs)
+        if self.concat_input:
+            self.conv_cat = UperNetConvModule(
+                self.in_channels + self.channels,
+                self.channels,
+                kernel_size=kernel_size,
+                padding=kernel_size // 2,
+            )
+
+        self.classifier = nn.Conv3d(self.channels, out_channels, kernel_size=1)
+
+    def init_weights(self):
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Conv3d):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+    def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+        # just take the relevant feature maps
+        hidden_states = encoder_hidden_states[self.in_index]
+        output = self.convs(hidden_states)
+        if self.concat_input:
+            output = self.conv_cat(torch.cat([hidden_states, output], dim=1))
+        output = self.classifier(output)
+        return output
+
+
+class ViTAdapter(nn.Module):
+    def __init__(
+        self,
+        img_size=(64, 256, 256),
+        patch_size=(16, 32, 32),
+        embed_dim=768,
+        # out_indices=[3, 5, 7, 11],
+    ):
+        super().__init__()
+        # self.out_indices = out_indices
+
+        self.grid_size = tuple(img_d // p_d for img_d, p_d in zip(img_size, patch_size))
+        self.hidden_size = embed_dim
+
+        if patch_size == (16, 32, 32):
+            self.fpn1 = nn.Sequential(
+                nn.ConvTranspose3d(
+                    embed_dim, embed_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2)
+                ),
+                nn.BatchNorm3d(embed_dim),
+                nn.GELU(),
+                nn.ConvTranspose3d(embed_dim, embed_dim, kernel_size=2, stride=2),
+                nn.BatchNorm3d(embed_dim),
+                nn.GELU(),
+                nn.ConvTranspose3d(embed_dim, embed_dim, kernel_size=2, stride=2),
+            )
+
+            # 8
+            self.fpn2 = nn.Sequential(
+                nn.ConvTranspose3d(
+                    embed_dim, embed_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2)
+                ),
+                nn.BatchNorm3d(embed_dim),
+                nn.GELU(),
+                nn.ConvTranspose3d(embed_dim, embed_dim, kernel_size=2, stride=2),
+            )
+
+            # 16
+            self.fpn3 = nn.Sequential(
+                nn.ConvTranspose3d(
+                    embed_dim, embed_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2)
+                ),
+            )
+
+            # 32
+            self.fpn4 = nn.MaxPool3d(kernel_size=(2, 1, 1), stride=(2, 1, 1))
+
+            self.adapters = [self.fpn1, self.fpn2, self.fpn3, self.fpn4]
+
+    def proj_feat(self, x):
+        
+        new_view = (x.size(0), *self.grid_size, self.hidden_size)
+        # print(f"x.shape: {x.shape}, expected: {new_view}, grid_size: {self.grid_size}")
+        x = x.view(new_view)
+        new_axes = (0, len(x.shape) - 1) + tuple(
+            d + 1 for d in range(len(self.grid_size))
+        )
+        x = x.permute(new_axes).contiguous()
+        return x
+
+    def forward(self, encoder_hidden_states):
+        output = []
+        # print(f"len_encoder_hidden: {len(encoder_hidden_states)}")
+        for index, op in zip(range(len(encoder_hidden_states)), self.adapters):
+            output.append(op(self.proj_feat(encoder_hidden_states[index])))
+        return output
+
+
+class UperNet(nn.Module):
+    def __init__(
+        self,
+        encoder,
+        in_channels,
+        out_channels,
+        adapter=None,
+        out_indices=None,
+        pool_scales=[1, 2, 3, 6],
+        hidden_size=512,
+        auxiliary_channels=256,
+        use_auxiliary_head=True,
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.adapter = adapter
+        self.out_indices = out_indices
+        self.decode_head = UperNetHead(
+            in_channels=in_channels,
+            pool_scales=pool_scales,
+            hidden_size=hidden_size,
+            out_channels=out_channels,
+        )
+        self.auxiliary_head = (
+            UperNetFCNHead(
+                in_channels=in_channels,
+                hidden_size=auxiliary_channels,
+                num_convs=1,
+                out_channels=out_channels,
+            )
+            if use_auxiliary_head
+            else None
+        )
+
+        self.hidden_norm = nn.ModuleList()
+        for in_channel in in_channels:
+            norm = LayerNorm(in_channel, eps=1e-6, data_format="channels_first") # nn.BatchNorm3d(out_channels)        
+            self.hidden_norm.append(norm)
+
+    def forward(self, x):
+        # print(f"403 input x.shape: {x.shape}")
+        encoder_hidden_states = self.encoder(x, ret_hids=True)
+        # print(f"405 {type(encoder_hidden_states)}, encoder_hidden_states: {len(encoder_hidden_states)}")
+        # for i, hidden_state in enumerate(encoder_hidden_states):
+        #     print(f"407 encoder_hidden_states[{i}]: {type(hidden_state)}, {len(hidden_state)}")
+        if isinstance(encoder_hidden_states, list) or isinstance(
+            encoder_hidden_states, Tuple
+        ):
+            encoder_hidden_states = encoder_hidden_states[-1]
+        # print(f"410 {type(encoder_hidden_states)}, encoder_hidden_states: {len(encoder_hidden_states)}")
+        # for i, hidden_state in enumerate(encoder_hidden_states):
+        #     print(f"412 encoder_hidden_states[{i}]: {hidden_state.shape}")
+        if self.out_indices:
+            encoder_hidden_states = [
+                encoder_hidden_states[i] for i in self.out_indices
+            ]
+        
+        encoder_hidden_states = [
+            norm(encoder_hidden_states[i])
+            for i, norm in enumerate(self.hidden_norm)
+        ]
+        # print(f"415 encoder_hidden_states: {len(encoder_hidden_states)}")
+        # for i in range(len(encoder_hidden_states)):
+        #     print(f"417 encoder_hidden_states[{i}]: {encoder_hidden_states[i].shape}")
+
+        if self.adapter:
+            encoder_hidden_states = self.adapter(encoder_hidden_states)
+
+        logits = self.decode_head(encoder_hidden_states)
+        logits = nn.functional.interpolate(
+            logits, size=x.shape[2:], mode="trilinear", align_corners=False
+        )
+        if not self.training:
+            return logits
+
+        auxiliary_logits = None
+        if self.auxiliary_head is not None:
+            auxiliary_logits = self.auxiliary_head(encoder_hidden_states)
+            auxiliary_logits = nn.functional.interpolate(
+                auxiliary_logits,
+                size=x.shape[2:],
+                mode="trilinear",
+                align_corners=False,
+            )
+            return [logits, auxiliary_logits]
+        return logits

models/util.py

@@ -0,0 +1,258 @@
+import torch
+import torch.nn as nn
+from itertools import chain
+from typing import Callable
+from torch.utils.checkpoint import checkpoint
+
+import numpy.random as random
+import torch.nn.functional as F
+# from MinkowskiEngine import SparseTensor
+
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+# class MinkowskiGRN(nn.Module):
+#     """GRN layer for sparse tensors."""
+
+#     def __init__(self, dim):
+#         super().__init__()
+#         self.gamma = nn.Parameter(torch.zeros(1, dim))
+#         self.beta = nn.Parameter(torch.zeros(1, dim))
+
+#     def forward(self, x):
+#         cm = x.coordinate_manager
+#         in_key = x.coordinate_map_key
+
+#         Gx = torch.norm(x.F, p=2, dim=0, keepdim=True)
+#         Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+#         return SparseTensor(
+#             self.gamma * (x.F * Nx) + self.beta + x.F,
+#             coordinate_map_key=in_key,
+#             coordinate_manager=cm,
+#         )
+
+
+class MinkowskiDropPath(nn.Module):
+    """Drop Path for sparse tensors."""
+
+    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
+        super(MinkowskiDropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        if self.drop_prob == 0.0 or not self.training:
+            return x
+        cm = x.coordinate_manager
+        in_key = x.coordinate_map_key
+        keep_prob = 1 - self.drop_prob
+        mask = (
+            torch.cat(
+                [
+                    (
+                        torch.ones(len(_))
+                        if random.uniform(0, 1) > self.drop_prob
+                        else torch.zeros(len(_))
+                    )
+                    for _ in x.decomposed_coordinates
+                ]
+            )
+            .view(-1, 1)
+            .to(x.device)
+        )
+        if keep_prob > 0.0 and self.scale_by_keep:
+            mask.div_(keep_prob)
+        return SparseTensor(
+            x.F * mask, coordinate_map_key=in_key, coordinate_manager=cm
+        )
+
+
+class MinkowskiLayerNorm(nn.Module):
+    """Channel-wise layer normalization for sparse tensors."""
+
+    def __init__(
+        self,
+        normalized_shape,
+        eps=1e-6,
+    ):
+        super(MinkowskiLayerNorm, self).__init__()
+        self.ln = nn.LayerNorm(normalized_shape, eps=eps)
+
+    def forward(self, input):
+        output = self.ln(input.F)
+        return SparseTensor(
+            output,
+            coordinate_map_key=input.coordinate_map_key,
+            coordinate_manager=input.coordinate_manager,
+        )
+
+
+class LayerNorm(nn.Module):
+    """LayerNorm that supports two data formats: channels_last (default) or channels_first.
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
+    with shape (batch_size, channels, height, width).
+    """
+
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(
+                x, self.normalized_shape, self.weight, self.bias, self.eps
+            )
+        elif self.data_format == "channels_first":
+            if len(x.shape) == 3: # for vit adapter
+                u = x.mean(1, keepdim=True)
+                s = (x - u).pow(2).mean(1, keepdim=True)
+                x = (x - u) / torch.sqrt(s + self.eps)
+                x = self.weight * x + self.bias
+                return x
+            else:
+                u = x.mean(1, keepdim=True)
+                s = (x - u).pow(2).mean(1, keepdim=True)
+                x = (x - u) / torch.sqrt(s + self.eps)
+                x = self.weight[:, None, None, None] * x + self.bias[:, None, None, None]
+                return x
+
+
+class GRN(nn.Module):
+    """GRN (Global Response Normalization) layer"""
+
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=(1, 2, 3), keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+def get_tokens(embed_dim: int, n_tokens: int) -> nn.Parameter:
+    """Return a learnable token of shape (1, n_tokens, embed_dim).
+
+    Args:
+        embed_dim: number of embedding channels.
+        n_tokens: number of tokens.
+
+    Returns:
+        token: learnable token.
+    """
+    token = nn.Parameter(torch.zeros(1, n_tokens, embed_dim))
+    # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
+    nn.init.trunc_normal_(token, std=0.02, b=2.0)
+    return token
+
+
+def init_weights(m):
+    if isinstance(m, nn.Linear):
+        # we use xavier_uniform following official JAX ViT:
+        torch.nn.init.xavier_uniform_(m.weight)
+        if isinstance(m, nn.Linear) and m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.LayerNorm):
+        nn.init.constant_(m.bias, 0)
+        nn.init.constant_(m.weight, 1.0)
+
+
+"""Gradient checkpointing utilities.
+
+Copied from
+https://github.com/huggingface/pytorch-image-models/blob/f8979d4f50b7920c78511746f7315df8f1857bc5/timm/models/_manipulate.py
+and added use_reentrant=False following warnings in pytorch docs.
+"""
+
+
+def checkpoint_seq(
+    functions: nn.Sequential,
+    x: torch.Tensor,
+    every: int = 1,
+    flatten: bool = False,
+    skip_last: bool = False,
+    preserve_rng_state: bool = True,
+) -> torch.Tensor:
+    r"""A helper function for checkpointing sequential models.
+
+    Sequential models execute a list of modules/functions in order
+    (sequentially). Therefore, we can divide such a sequence into segments
+    and checkpoint each segment. All segments except run in :func:`torch.no_grad`
+    manner, i.e., not storing the intermediate activations. The inputs of each
+    checkpointed segment will be saved for re-running the segment in the backward pass.
+
+    See :func:`~torch.utils.checkpoint.checkpoint` on how checkpointing works.
+
+    .. warning::
+        Checkpointing currently only supports :func:`torch.autograd.backward`
+        and only if its `inputs` argument is not passed. :func:`torch.autograd.grad`
+        is not supported.
+
+    .. warning:
+        At least one of the inputs needs to have :code:`requires_grad=True` if
+        grads are needed for model inputs, otherwise the checkpointed part of the
+        model won't have gradients.
+
+    Args:
+        functions: A :class:`torch.nn.Sequential` or the list of modules or functions to run sequentially.
+        x: A Tensor that is input to :attr:`functions`
+        every: checkpoint every-n functions (default: 1)
+        flatten (bool): flatten nn.Sequential of nn.Sequentials
+        skip_last (bool): skip checkpointing the last function in the sequence if True
+        preserve_rng_state (bool, optional, default=True):  Omit stashing and restoring
+            the RNG state during each checkpoint.
+
+    Returns:
+        Output of running :attr:`functions` sequentially on :attr:`*inputs`
+
+    Example:
+        >>> model = nn.Sequential(...)
+        >>> input_var = checkpoint_seq(model, input_var, every=2)
+    """
+
+    def run_function(
+        start: int, end: int, functions: nn.Sequential
+    ) -> Callable[[torch.Tensor], torch.Tensor]:
+        def forward(_x: torch.Tensor) -> torch.Tensor:  # pylint: disable=invalid-name
+            for j in range(start, end + 1):
+                _x = functions[j](_x)
+            return _x
+
+        return forward
+
+    if isinstance(functions, torch.nn.Sequential):
+        functions = functions.children()
+    if flatten:
+        functions = chain.from_iterable(functions)
+    if not isinstance(functions, (tuple, list)):
+        functions = tuple(functions)
+
+    num_checkpointed = len(functions)
+    if skip_last:
+        num_checkpointed -= 1
+    end = -1
+    for start in range(0, num_checkpointed, every):
+        end = min(start + every - 1, num_checkpointed - 1)
+        x = checkpoint(
+            run_function(start, end, functions),
+            x,
+            use_reentrant=False,
+            preserve_rng_state=preserve_rng_state,
+        )
+    if skip_last:
+        return run_function(end + 1, len(functions) - 1, functions)(x)
+    return x

requirements.txt

@@ -34,6 +34,3 @@ tqdm==4.67.1
 # Additional dependencies for model architecture
 einops==0.8.1
 timm==1.0.15
-
-# ProFound package from GitHub
-git+https://github.com/pipiwang/ProFound.git@demo

util/__init__.py

@@ -0,0 +1 @@
+# ProFound utilities package 

util/convnext_optim.py

@@ -0,0 +1,127 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch
+from torch import optim as optim
+import json
+
+
+def get_num_layer_for_convnext_single(var_name, depths):
+    """
+    Each layer is assigned distinctive layer ids
+    """
+    if var_name.startswith("downsample_layers"):
+        stage_id = int(var_name.split(".")[1])
+        layer_id = sum(depths[:stage_id]) + 1
+        return layer_id
+
+    elif var_name.startswith("stages"):
+        stage_id = int(var_name.split(".")[1])
+        block_id = int(var_name.split(".")[2])
+        layer_id = sum(depths[:stage_id]) + block_id + 1
+        return layer_id
+
+    else:
+        return sum(depths) + 1
+
+
+def get_num_layer_for_convnext(var_name):
+    """
+    Divide [3, 3, 27, 3] layers into 12 groups; each group is three
+    consecutive blocks, including possible neighboring downsample layers;
+    adapted from https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py
+    """
+    num_max_layer = 12
+    if var_name.startswith("downsample_layers"):
+        stage_id = int(var_name.split(".")[1])
+        if stage_id == 0:
+            layer_id = 0
+        elif stage_id == 1 or stage_id == 2:
+            layer_id = stage_id + 1
+        elif stage_id == 3:
+            layer_id = 12
+        return layer_id
+
+    elif var_name.startswith("stages"):
+        stage_id = int(var_name.split(".")[1])
+        block_id = int(var_name.split(".")[2])
+        if stage_id == 0 or stage_id == 1:
+            layer_id = stage_id + 1
+        elif stage_id == 2:
+            layer_id = 3 + block_id // 3
+        elif stage_id == 3:
+            layer_id = 12
+        return layer_id
+    else:
+        return num_max_layer + 1
+
+
+class LayerDecayValueAssigner(object):
+    def __init__(self, values, depths=[3, 3, 27, 3], layer_decay_type="single"):
+        self.values = values
+        self.depths = depths
+        self.layer_decay_type = layer_decay_type
+
+    def get_scale(self, layer_id):
+        return self.values[layer_id]
+
+    def get_layer_id(self, var_name):
+        if self.layer_decay_type == "single":
+            return get_num_layer_for_convnext_single(var_name, self.depths)
+        else:
+            return get_num_layer_for_convnext(var_name)
+
+
+def get_parameter_groups(
+    model, weight_decay=1e-5, skip_list=(), get_num_layer=None, get_layer_scale=None
+):
+    parameter_group_names = {}
+    parameter_group_vars = {}
+
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue  # frozen weights
+        if (
+            len(param.shape) == 1
+            or name.endswith(".bias")
+            or name in skip_list
+            or name.endswith(".gamma")
+            or name.endswith(".beta")
+        ):
+            group_name = "no_decay"
+            this_weight_decay = 0.0
+        else:
+            group_name = "decay"
+            this_weight_decay = weight_decay
+        if get_num_layer is not None:
+            layer_id = get_num_layer(name)
+            group_name = "layer_%d_%s" % (layer_id, group_name)
+        else:
+            layer_id = None
+
+        if group_name not in parameter_group_names:
+            if get_layer_scale is not None:
+                scale = get_layer_scale(layer_id)
+            else:
+                scale = 1.0
+
+            parameter_group_names[group_name] = {
+                "weight_decay": this_weight_decay,
+                "params": [],
+                "lr_scale": scale,
+            }
+            parameter_group_vars[group_name] = {
+                "weight_decay": this_weight_decay,
+                "params": [],
+                "lr_scale": scale,
+            }
+
+        parameter_group_vars[group_name]["params"].append(param)
+        parameter_group_names[group_name]["params"].append(name)
+    print("Param groups = %s" % json.dumps(parameter_group_names, indent=2))
+    return list(parameter_group_vars.values())

util/lars.py

@@ -0,0 +1,59 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# LARS optimizer, implementation from MoCo v3:
+# https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+
+import torch
+
+class LARS(torch.optim.Optimizer):
+    """
+    LARS optimizer, no rate scaling or weight decay for parameters <= 1D.
+    """
+
+    def __init__(
+        self, params, lr=0, weight_decay=0, momentum=0.9, trust_coefficient=0.001
+    ):
+        defaults = dict(
+            lr=lr,
+            weight_decay=weight_decay,
+            momentum=momentum,
+            trust_coefficient=trust_coefficient,
+        )
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        for g in self.param_groups:
+            for p in g["params"]:
+                dp = p.grad
+
+                if dp is None:
+                    continue
+
+                if p.ndim > 1:  # if not normalization gamma/beta or bias
+                    dp = dp.add(p, alpha=g["weight_decay"])
+                    param_norm = torch.norm(p)
+                    update_norm = torch.norm(dp)
+                    one = torch.ones_like(param_norm)
+                    q = torch.where(
+                        param_norm > 0.0,
+                        torch.where(
+                            update_norm > 0,
+                            (g["trust_coefficient"] * param_norm / update_norm),
+                            one,
+                        ),
+                        one,
+                    )
+                    dp = dp.mul(q)
+
+                param_state = self.state[p]
+                if "mu" not in param_state:
+                    param_state["mu"] = torch.zeros_like(p)
+                mu = param_state["mu"]
+                mu.mul_(g["momentum"]).add_(dp)
+                p.add_(mu, alpha=-g["lr"])

util/lr_sched.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+
+
+def adjust_learning_rate(optimizer, epoch, args):
+    """Decay the learning rate with half-cycle cosine after warmup"""
+    if epoch < args.warmup_epochs:
+        lr = args.lr * epoch / args.warmup_epochs
+    else:
+        lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * (
+            1.0
+            + math.cos(
+                math.pi
+                * (epoch - args.warmup_epochs)
+                / (args.epochs - args.warmup_epochs)
+            )
+        )
+    for param_group in optimizer.param_groups:
+        if "lr_scale" in param_group:
+            param_group["lr"] = lr * param_group["lr_scale"]
+        else:
+            param_group["lr"] = lr
+    return lr

util/metric.py

@@ -0,0 +1,340 @@
+import torch
+import prettytable
+import copy
+import sys
+from importlib import import_module
+from inspect import signature
+from pathlib import Path
+from typing import Optional, Union
+
+import numpy as np
+from scipy.stats import kendalltau, pearsonr, spearmanr
+from sklearn.metrics import (
+    confusion_matrix,
+    f1_score,
+    fbeta_score,
+    get_scorer,
+    get_scorer_names,
+    make_scorer,
+)
+
+
+def binary_accuracy(output: torch.Tensor, target: torch.Tensor) -> float:
+    """Computes the accuracy for binary classification"""
+    with torch.no_grad():
+        batch_size = target.size(0)
+        pred = (output >= 0.5).float().t().view(-1)
+        correct = pred.eq(target.view(-1)).float().sum()
+        correct.mul_(100.0 / batch_size)
+        return correct
+
+
+def accuracy(output, target, topk=(1,)):
+    r"""
+    Computes the accuracy over the k top predictions for the specified values of k
+
+    Args:
+        output (tensor): Classification outputs, :math:`(N, C)` where `C = number of classes`
+        target (tensor): :math:`(N)` where each value is :math:`0 \leq \text{targets}[i] \leq C-1`
+        topk (sequence[int]): A list of top-N number.
+
+    Returns:
+        Top-N accuracies (N :math:`\in` topK).
+    """
+    with torch.no_grad():
+        maxk = max(topk)
+        batch_size = target.size(0)
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target[None])
+
+        res = []
+        for k in topk:
+            correct_k = correct[:k].flatten().sum(dtype=torch.float32)
+            res.append(correct_k * (100.0 / batch_size))
+        return res
+
+
+class ConfusionMatrix(object):
+    def __init__(self, num_classes):
+        self.num_classes = num_classes
+        self.mat = None
+
+    def update(self, target, output):
+        """
+        Update confusion matrix.
+
+        Args:
+            target: ground truth
+            output: predictions of models
+
+        Shape:
+            - target: :math:`(minibatch, C)` where C means the number of classes.
+            - output: :math:`(minibatch, C)` where C means the number of classes.
+        """
+        n = self.num_classes
+        if self.mat is None:
+            self.mat = torch.zeros((n, n), dtype=torch.int64, device=target.device)
+        with torch.no_grad():
+            k = (target >= 0) & (target < n)
+            inds = n * target[k].to(torch.int64) + output[k]
+            self.mat += torch.bincount(inds, minlength=n**2).reshape(n, n)
+
+    def reset(self):
+        self.mat.zero_()
+
+    def compute(self):
+        """compute global accuracy, per-class accuracy and per-class IoU"""
+        h = self.mat.float()
+        acc_global = torch.diag(h).sum() / h.sum()
+        acc = torch.diag(h) / h.sum(1)
+        iu = torch.diag(h) / (h.sum(1) + h.sum(0) - torch.diag(h))
+        return acc_global, acc, iu
+
+    # def reduce_from_all_processes(self):
+    #     if not torch.distributed.is_available():
+    #         return
+    #     if not torch.distributed.is_initialized():
+    #         return
+    #     torch.distributed.barrier()
+    #     torch.distributed.all_reduce(self.mat)
+
+    def __str__(self):
+        acc_global, acc, iu = self.compute()
+        return (
+            "global correct: {:.1f}\n"
+            "average row correct: {}\n"
+            "IoU: {}\n"
+            "mean IoU: {:.1f}"
+        ).format(
+            acc_global.item() * 100,
+            ["{:.1f}".format(i) for i in (acc * 100).tolist()],
+            ["{:.1f}".format(i) for i in (iu * 100).tolist()],
+            iu.mean().item() * 100,
+        )
+
+    def format(self, classes: list):
+        """Get the accuracy and IoU for each class in the table format"""
+        acc_global, acc, iu = self.compute()
+
+        table = prettytable.PrettyTable(["class", "acc", "iou"])
+        for i, class_name, per_acc, per_iu in zip(
+            range(len(classes)), classes, (acc * 100).tolist(), (iu * 100).tolist()
+        ):
+            table.add_row([class_name, per_acc, per_iu])
+
+        return (
+            "global correct: {:.1f}\nmean correct:{:.1f}\nmean IoU: {:.1f}\n{}".format(
+                acc_global.item() * 100,
+                acc.mean().item() * 100,
+                iu.mean().item() * 100,
+                table.get_string(),
+            )
+        )
+
+
+def kappa(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    weights: Optional[Union[str, np.ndarray]] = None,
+    allow_off_by_one: bool = False,
+) -> float:
+    """
+    Calculate the kappa inter-rater agreement.
+
+    The agreement is calculated between the gold standard and the predicted
+    ratings. Potential values range from -1 (representing complete disagreement)
+    to 1 (representing complete agreement).  A kappa value of 0 is expected if
+    all agreement is due to chance.
+
+    In the course of calculating kappa, all items in ``y_true`` and ``y_pred`` will
+    first be converted to floats and then rounded to integers.
+
+    It is assumed that y_true and y_pred contain the complete range of possible
+    ratings.
+
+    This function contains a combination of code from yorchopolis's kappa-stats
+    and Ben Hamner's Metrics projects on Github.
+
+    Parameters
+    ----------
+    y_true : numpy.ndarray
+        The true/actual/gold labels for the data.
+    y_pred : numpy.ndarray
+        The predicted/observed labels for the data.
+    weights : Optional[Union[str, numpy.ndarray]], default=None
+        Specifies the weight matrix for the calculation.
+        Possible values are: ``None`` (unweighted-kappa), ``"quadratic"``
+        (quadratically weighted kappa), ``"linear"`` (linearly weighted kappa),
+        and a two-dimensional numpy array (a custom matrix of weights). Each
+        weight in this array corresponds to the :math:`w_{ij}` values in the
+        Wikipedia description of how to calculate weighted Cohen's kappa.
+    allow_off_by_one : bool, default=False
+        If true, ratings that are off by one are counted as
+        equal, and all other differences are reduced by
+        one. For example, 1 and 2 will be considered to be
+        equal, whereas 1 and 3 will have a difference of 1
+        for when building the weights matrix.
+
+    Returns
+    -------
+    float
+        The weighted or unweighted kappa score.
+
+    Raises
+    ------
+    AssertionError
+        If ``y_true`` != ``y_pred``.
+    ValueError
+        If labels cannot be converted to int.
+    ValueError
+        If invalid weight scheme.
+    """
+    # Ensure that the lists are both the same length
+    assert len(y_true) == len(y_pred)
+
+    # This rather crazy looking typecast is intended to work as follows:
+    # If an input is an int, the operations will have no effect.
+    # If it is a float, it will be rounded and then converted to an int
+    # because the ml_metrics package requires ints.
+    # If it is a str like "1", then it will be converted to a (rounded) int.
+    # If it is a str that can't be typecast, then the user is
+    # given a hopefully useful error message.
+    try:
+        y_true = np.array([int(np.round(float(y))) for y in y_true])
+        y_pred = np.array([int(np.round(float(y))) for y in y_pred])
+    except ValueError:
+        raise ValueError(
+            "For kappa, the labels should be integers or strings"
+            " that can be converted to ints (E.g., '4.0' or "
+            "'3')."
+        )
+
+    # Figure out normalized expected values
+    min_rating = min(min(y_true), min(y_pred))
+    max_rating = max(max(y_true), max(y_pred))
+
+    # shift the values so that the lowest value is 0
+    # (to support scales that include negative values)
+    y_true = y_true - min_rating
+    y_pred = y_pred - min_rating
+
+    # Build the observed/confusion matrix
+    num_ratings = max_rating - min_rating + 1
+    observed = confusion_matrix(y_true, y_pred, labels=list(range(num_ratings)))
+    num_scored_items = float(len(y_true))
+
+    # Build weight array if weren't passed one
+    if isinstance(weights, str):
+        wt_scheme = weights
+        weights = None
+    else:
+        wt_scheme = ""
+
+    if weights is None:
+        kappa_weights = np.empty((num_ratings, num_ratings))
+        for i in range(num_ratings):
+            for j in range(num_ratings):
+                diff = abs(i - j)
+                if allow_off_by_one and diff:
+                    diff -= 1
+                if wt_scheme == "linear":
+                    kappa_weights[i, j] = diff
+                elif wt_scheme == "quadratic":
+                    kappa_weights[i, j] = diff**2
+                elif not wt_scheme:  # unweighted
+                    kappa_weights[i, j] = bool(diff)
+                else:
+                    raise ValueError(
+                        "Invalid weight scheme specified for " f"kappa: {wt_scheme}"
+                    )
+    else:
+        kappa_weights = weights
+
+    hist_true: np.ndarray = np.bincount(y_true, minlength=num_ratings)
+    hist_true = hist_true[:num_ratings] / num_scored_items
+    hist_pred: np.ndarray = np.bincount(y_pred, minlength=num_ratings)
+    hist_pred = hist_pred[:num_ratings] / num_scored_items
+    expected = np.outer(hist_true, hist_pred)
+
+    # Normalize observed array
+    observed = observed / num_scored_items
+
+    # If all weights are zero, that means no disagreements matter.
+    k = 1.0
+    if np.count_nonzero(kappa_weights):
+        observed_sum = np.sum(kappa_weights * observed)
+        expected_sum = np.sum(kappa_weights * expected)
+        k -= np.sum(observed_sum) / np.sum(expected_sum)
+
+    return k
+
+
+def correlation(
+    y_true: np.ndarray, y_pred: np.ndarray, corr_type: str = "pearson"
+) -> float:
+    """
+    Calculate given correlation type between ``y_true`` and ``y_pred``.
+
+    ``y_pred`` can be multi-dimensional. If ``y_pred`` is 1-dimensional, it
+    may either contain probabilities, most-likely classification labels, or
+    regressor predictions. In that case, we simply return the correlation
+    between ``y_true`` and ``y_pred``. If ``y_pred`` is multi-dimensional,
+    it contains probabilties for multiple classes in which case, we infer the
+    most likely labels and then compute the correlation between those and
+    ``y_true``.
+
+    Parameters
+    ----------
+    y_true : numpy.ndarray
+        The true/actual/gold labels for the data.
+    y_pred : numpy.ndarray
+        The predicted/observed labels for the data.
+    corr_type : str, default="pearson"
+        Which type of correlation to compute. Possible
+        choices are "pearson", "spearman", and "kendall_tau".
+
+    Returns
+    -------
+    float
+        correlation value if well-defined, else 0.0
+    """
+    # get the correlation function to use based on the given type
+    corr_func = pearsonr
+    if corr_type == "spearman":
+        corr_func = spearmanr
+    elif corr_type == "kendall_tau":
+        corr_func = kendalltau
+
+    # convert to numpy array in case we are passed a list
+    y_pred = np.array(y_pred)
+
+    # multi-dimensional -> probability array -> get label
+    if y_pred.ndim > 1:
+        labels = np.argmax(y_pred, axis=1)
+        ret_score = corr_func(y_true, labels)[0]
+    # 1-dimensional -> probabilities/labels -> use as is
+    else:
+        ret_score = corr_func(y_true, y_pred)[0]
+    return ret_score
+
+
+def f1_score_least_frequent(y_true: np.ndarray, y_pred: np.ndarray) -> float:
+    """
+    Calculate F1 score of the least frequent label/class.
+
+    Parameters
+    ----------
+    y_true : numpy.ndarray
+        The true/actual/gold labels for the data.
+    y_pred : numpy.ndarray
+        The predicted/observed labels for the data.
+
+    Returns
+    -------
+    float
+        F1 score of the least frequent label.
+    """
+    least_frequent = np.bincount(y_true).argmin()
+    return f1_score(y_true, y_pred, average=None)[least_frequent]

util/misc.py

@@ -0,0 +1,455 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+
+import builtins
+import datetime
+import os
+import time
+from collections import defaultdict, deque
+from pathlib import Path
+import shutil
+import torch
+import torch.distributed as dist
+from torch import inf
+import json
+
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value,
+        )
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError(
+            "'{}' object has no attribute '{}'".format(type(self).__name__, attr)
+        )
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append("{}: {}".format(name, str(meter)))
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ""
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(len(iterable)))) + "d"
+        log_msg = [
+            header,
+            "[{0" + space_fmt + "}/{1}]",
+            "eta: {eta}",
+            "{meters}",
+            "time: {time}",
+            "data: {data}",
+        ]
+        if torch.cuda.is_available():
+            log_msg.append("max mem: {memory:.0f}")
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
+                else:
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                        )
+                    )
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print(
+            "{} Total time: {} ({:.4f} s / it)".format(
+                header, total_time_str, total_time / len(iterable)
+            )
+        )
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    builtin_print = builtins.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        force = force or (get_world_size() > 8)
+        if is_master or force:
+            now = datetime.datetime.now().time()
+            builtin_print("[{}] ".format(now), end="")  # print with time stamp
+            builtin_print(*args, **kwargs)
+
+    builtins.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if args.dist_on_itp:
+        args.rank = int(os.environ["OMPI_COMM_WORLD_RANK"])
+        args.world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"])
+        args.gpu = int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])
+        args.dist_url = "tcp://%s:%s" % (
+            os.environ["MASTER_ADDR"],
+            os.environ["MASTER_PORT"],
+        )
+        os.environ["LOCAL_RANK"] = str(args.gpu)
+        os.environ["RANK"] = str(args.rank)
+        os.environ["WORLD_SIZE"] = str(args.world_size)
+        # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+    elif "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print("Not using distributed mode")
+        setup_for_distributed(is_master=True)  # hack
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = "nccl"
+    print(
+        "| distributed init (rank {}): {}, gpu {}".format(
+            args.rank, args.dist_url, args.gpu
+        ),
+        flush=True,
+    )
+    torch.distributed.init_process_group(
+        backend=args.dist_backend,
+        init_method=args.dist_url,
+        world_size=args.world_size,
+        rank=args.rank,
+    )
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+class NativeScalerWithGradNormCount:
+    state_dict_key = "amp_scaler"
+
+    def __init__(self):
+        self._scaler = torch.cuda.amp.GradScaler()
+
+    def __call__(
+        self,
+        loss,
+        optimizer,
+        clip_grad=None,
+        parameters=None,
+        create_graph=False,
+        update_grad=True,
+    ):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(
+                    optimizer
+                )  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm_(parameters)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = [p for p in parameters if p.grad is not None]
+    norm_type = float(norm_type)
+    if len(parameters) == 0:
+        return torch.tensor(0.0)
+    device = parameters[0].grad.device
+    if norm_type == inf:
+        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+    else:
+        total_norm = torch.norm(
+            torch.stack(
+                [torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]
+            ),
+            norm_type,
+        )
+    return total_norm
+
+
+def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler):
+    output_dir = Path(args.output_dir)
+    epoch_name = str(epoch)
+    if loss_scaler is not None:
+        checkpoint_paths = [output_dir / ("checkpoint-%s.pth" % epoch_name)]
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                "model": model_without_ddp.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "epoch": epoch,
+                "scaler": loss_scaler.state_dict(),
+                "args": args,
+            }
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {"epoch": epoch}
+        model.save_checkpoint(
+            save_dir=args.output_dir,
+            tag="checkpoint-%s" % epoch_name,
+            client_state=client_state,
+        )
+
+
+def save_best_model(
+    args, epoch, model, model_without_ddp, optimizer, loss_scaler, is_best
+):
+    output_dir = Path(args.output_dir)
+    epoch_name = str(epoch)
+    if loss_scaler is not None:
+        checkpoint_path = output_dir / ("last.pth.tar")
+        to_save = {
+            "model": model_without_ddp.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "epoch": epoch,
+            "scaler": loss_scaler.state_dict(),
+            "args": args,
+        }
+        save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {"epoch": epoch}
+        model.save_checkpoint(
+            save_dir=args.output_dir,
+            tag="checkpoint-%s" % epoch_name,
+            client_state=client_state,
+        )
+    if is_best:
+        filepath_best = output_dir / ("best.pth.tar")
+        shutil.copyfile(checkpoint_path, filepath_best)
+
+
+def save_current_best_model(
+    args, epoch, model, model_without_ddp, optimizer, loss_scaler, is_best, current_interval
+):
+    output_dir = Path(args.output_dir)
+    epoch_name = str(epoch)
+    if loss_scaler is not None:
+        checkpoint_paths = [output_dir / (f"{current_interval}_last.pth.tar")]
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                "model": model_without_ddp.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "epoch": epoch,
+                "scaler": loss_scaler.state_dict(),
+                "args": args,
+            }
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {"epoch": epoch}
+        model.save_checkpoint(
+            save_dir=args.output_dir,
+            tag="checkpoint-%s" % epoch_name,
+            client_state=client_state,
+        )
+    if is_best:
+        filepath_best = output_dir / (f"{current_interval}_best.pth.tar")
+        shutil.copyfile(checkpoint_path, filepath_best)
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler):
+    if args.resume:
+        if args.resume.startswith("https"):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location="cpu", check_hash=True, weights_only=False
+            )
+        else:
+            checkpoint = torch.load(args.resume, map_location="cpu")
+        model_without_ddp.load_state_dict(checkpoint["model"], weights_only=False)
+        print("Resume checkpoint %s" % args.resume)
+        if (
+            "optimizer" in checkpoint
+            and "epoch" in checkpoint
+            and not (hasattr(args, "eval") and args.eval)
+        ):
+            optimizer.load_state_dict(checkpoint["optimizer"], weights_only=False)
+            args.start_epoch = checkpoint["epoch"] + 1
+            if "scaler" in checkpoint:
+                loss_scaler.load_state_dict(checkpoint["scaler"], weights_only=False)
+            print("With optim & sched!")
+
+
+def all_reduce_mean(x):
+    world_size = get_world_size()
+    if world_size > 1:
+        x_reduce = torch.tensor(x).cuda()
+        dist.all_reduce(x_reduce)
+        x_reduce /= world_size
+        return x_reduce.item()
+    else:
+        return x
+
+
+def write_log(log_writer, log_stats, args):
+    if args.output_dir and is_main_process():
+        if log_writer is not None:
+            log_writer.flush()
+        with open(
+            os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8"
+        ) as f:
+            f.write(json.dumps(log_stats) + "\n")