update models, output, examples

README.md

@@ -1,13 +0,0 @@
----
-title: Deep Learning Model for Pediatric Bone Age
-emoji: 💻
-colorFrom: red
-colorTo: blue
-sdk: gradio
-sdk_version: 3.8.2
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,106 +1,168 @@
+import cv2
 import gradio as gr
-import timm
-import torch 
+import json
+import numpy as np
+import torch
 import torch.nn as nn
 
+from einops import rearrange
+from importlib import import_module
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+from skimage.exposure import match_histograms
+from skp.utils import load_model_from_config, load_kfold_ensemble_as_list
 
-def change_num_input_channels(model, in_channels=1):
-    """
-    Assumes number of input channels in model is 3.
-    """
-    for i, m in enumerate(model.modules()):
-      if isinstance(m, (nn.Conv2d,nn.Conv3d)) and m.in_channels == 3:
-        m.in_channels = in_channels
-        # First, sum across channels
-        W = m.weight.sum(1, keepdim=True)
-        # Then, divide by number of channels
-        W = W / in_channels
-        # Then, repeat by number of channels
-        size = [1] * W.ndim
-        size[1] = in_channels
-        W = W.repeat(size)
-        m.weight = nn.Parameter(W)
-        break
-    return model
-
-
-class Net2D(nn.Module):
-
-    def __init__(self, weights):
-        super().__init__()
-        self.backbone = timm.create_model("tf_efficientnetv2_s", pretrained=False, global_pool="", num_classes=0)
-        self.backbone = change_num_input_channels(self.backbone, 2)
-        self.pool_layer = nn.AdaptiveAvgPool2d(1)
-        self.dropout = nn.Dropout(0.2)
-        self.classifier = nn.Linear(1280, 1) 
-        self.load_state_dict(weights)
-
-    def forward(self, x):
-        x = self.backbone(x)
-        x = self.pool_layer(x).view(x.size(0), -1)
-        x = self.dropout(x)
-        x = self.classifier(x)
-        return x[:, 0] if x.size(1) == 1 else x
 
-
-class Ensemble(nn.Module):
-
-    def __init__(self, model_list):
+class ModelForGradCAM(nn.Module):
+    def __init__(self, model):
         super().__init__()
-        self.model_list = nn.ModuleList(model_list)
-    
-    def forward(self, x):
-        return torch.stack([model(x) for model in self.model_list]).mean(0)
+        self.model = model
 
+    def forward(self, x):
+        return self.model({"x": x})["logits1"]
 
-checkpoints = ["fold0.ckpt", "fold1.ckpt", "fold2.ckpt"]
-weights = [torch.load(ckpt, map_location=torch.device("cpu"))["state_dict"] for ckpt in checkpoints]
-weights = [{k.replace("model.", "") : v for k, v in wt.items()} for wt in weights]
-models = [Net2D(wt) for wt in weights] 
-ensemble = Ensemble(models).eval()
 
-def predict_bone_age(Radiograph, Sex):
-    img = torch.from_numpy(Radiograph)
-    img = img.unsqueeze(0).unsqueeze(0)
-    img = img / img.max()
-    img = img - 0.5
-    img = img * 2.0
-    if Sex == 1:
-        img = torch.cat([img, torch.zeros_like(img) + 1], dim=1)
-    else:
-        img = torch.cat([img, torch.zeros_like(img) - 1], dim=1)
-    with torch.no_grad():
-        bone_age = ensemble(img.float())[0].item()
-    total_months = bone_age * 12
-    years = int(total_months // 12)
-    months = total_months - years * 12
+def convert_bone_age_to_string(bone_age: float):
+    # bone_age in months
+    years = round(bone_age // 12)
+    months = bone_age - (years * 12)
     months = round(months)
-    if months == 12: 
+
+    if months == 12:
         years += 1
         months = 0
+
     if years == 0:
         str_output = f"{months} months" if months != 1 else "1 month"
     else:
-        months = round(months)
         if months == 0:
             str_output = f"{years} years" if years != 1 else "1 year"
         else:
-            str_output = f"{years} years, {months} months" if months != 1 else f"{years} years, 1 month"
-    return f"Estimated Bone Age: {str_output}"
+            str_output = (
+                f"{years} years, {months} months"
+                if months != 1
+                else f"{years} years, 1 month"
+            )
+    return str_output
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+cfg_crop = import_module("skp.configs.boneage.cfg_crop_simple_resize").cfg
+crop_model = load_model_from_config(
+    cfg_crop, weights_path="crop.pt", device=device, eval_mode=True
+)
 
+cfg = import_module("skp.configs.boneage.cfg_female_channel_reg_cls_match_hist").cfg
+cfg.backbone = "convnextv2_tiny"
 
-image = gr.Image(shape=(512, 512), image_mode="L")
+model_list = load_kfold_ensemble_as_list(
+    cfg, [f"net{i}.pt" for i in range(3)], device=device, eval_mode=True
+)
+
+ref_img = rearrange(cv2.imread("ref_img.png", 0), "h w -> h w 1 ")
+
+with open("greulich_and_pyle_ages.json", "r") as f:
+    greulich_and_pyle_ages = json.load(f)["bone_ages"]
+
+greulich_and_pyle_ages = {k: np.asarray(v) for k, v in greulich_and_pyle_ages.items()}
+
+model_grad_cam = ModelForGradCAM(model_list[0])
+target_layers = [model_grad_cam.model.backbone.stages[-1]]
+
+def predict_bone_age(Radiograph, Sex):
+    x0 = rearrange(Radiograph, "h w -> h w 1")
+    x = cfg_crop.val_transforms(image=x0)["image"]
+    x = torch.from_numpy(x)
+    x = rearrange(x, "h w c -> 1 c h w")
+    # crop
+    with torch.inference_mode():
+        box = crop_model({"x": x.to(device).float()}, return_loss=False)["logits"][
+            0
+        ].cpu()
+        box[[0, 2]] = box[[0, 2]] * x0.shape[1]
+        box[[1, 3]] = box[[1, 3]] * x0.shape[0]
+    box = box.numpy().astype("int")
+    x, y, w, h = box
+    x0 = x0[y : y + h, x : x + w]
+    # histogram matching
+    x0 = match_histograms(x0, ref_img)
+    x = cfg.val_transforms(image=x0)["image"]
+    # create image channel for female/male
+    ch = np.zeros_like(x)
+    if Sex:  # 0- male, 1- female
+        ch[...] = 255
+    x = np.concatenate([x, ch], axis=-1)
+    x = torch.from_numpy(x)
+    x = rearrange(x, "h w c -> 1 c h w")
+    with torch.inference_mode():
+        bone_age = []
+        for each_model in model_list:
+            pred = each_model({"x": x.to(device).float()}, return_loss=False)[
+                "logits1"
+            ][0].cpu()
+            pred = (pred.softmax(0) * torch.arange(240)).sum().numpy()
+            bone_age.append(pred)
+        bone_age = np.mean(bone_age)
+
+    gp_ages = greulich_and_pyle_ages["female" if Sex else "male"]
+    diffs_gp = np.abs(bone_age - gp_ages)
+    diffs_gp = np.argsort(diffs_gp)
+    closest1 = gp_ages[diffs_gp[0]]
+    closest2 = gp_ages[diffs_gp[1]]
+
+    bone_age_str = convert_bone_age_to_string(bone_age)
+    closest1 = convert_bone_age_to_string(closest1)
+    closest2 = convert_bone_age_to_string(closest2)
+
+    targets = [ClassifierOutputTarget(round(bone_age))]
+    with GradCAM(model=model_grad_cam, target_layers=target_layers) as cam:
+        grayscale_cam = cam(input_tensor=x.to(device).float(), targets=targets, eigen_smooth=True)
+    
+    heatmap = cv2.applyColorMap((grayscale_cam[0] * 255).astype("uint8"), cv2.COLORMAP_JET)
+    image = cv2.cvtColor(x[0, 0].cpu().numpy().astype("uint8"), cv2.COLOR_GRAY2RGB)
+    image_weight = 0.6
+    grad_cam_image = (1 - image_weight) * heatmap[..., ::-1] + image_weight * image
+    grad_cam_image = grad_cam_image.astype("uint8")
+
+    return f"Predicted bone age: {bone_age_str}\n\nThe closest Greulich & Pyle bone ages are:\n  1) {closest1}\n  2) {closest2}", grad_cam_image
+
+
+image = gr.Image(image_mode="L")
 sex = gr.Radio(["Male", "Female"], type="index")
-label = gr.Label(show_label=True, label="Result")
+textbox = gr.Textbox(show_label=True, label="Result")
+grad_cam_image = gr.Image(image_mode="RGB", label="Heatmap")
 
-demo = gr.Interface(
-    fn=predict_bone_age,
-    inputs=[image, sex],
-    outputs=label,
-    )
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+    # Deep Learning Model for Pediatric Bone Age
 
+    This model predicts the bone age from a single frontal view hand radiograph. 
+    The model was trained on the publicly available 
+    [RSNA Pediatric Bone Age Challenge](https://www.rsna.org/rsnai/ai-image-challenge/rsna-pediatric-bone-age-challenge-2017) dataset. 
+    The model achieves a mean absolute error of 4.26 months on the original test set comprising 200 multi-annotated hand radiographs, 
+    which is competitive with [top solutions](https://pubs.rsna.org/doi/10.1148/radiol.2018180736) from the original challenge.
 
-if __name__ == "__main__":
-    demo.launch()
+    This model is for demonstration purposes only and has NOT been approved by any regulatory agency for clinical use. The user assumes
+    any and all responsibility regarding their own use of this model and its outputs. Do NOT upload any images containing protected
+    health information, as this demonstration is not compliant with patient privacy laws.
 
+    Created by: Ian Pan, <https://ianpan.me>
+    
+    Last updated: December 15, 2024
+    """
+    )
+    gr.Interface(
+        fn=predict_bone_age,
+        inputs=[image, sex],
+        outputs=[textbox, grad_cam_image],
+        examples=[
+            ["examples/2639.png", "Female"],
+            ["examples/10043.png", "Female"],
+            ["examples/8888.png", "Female"],
+        ],
+    )
 
+if __name__ == "__main__":
+    demo.launch(share=True)

crop.pt

@@ -0,0 +1 @@
+../../experiments/boneage/boneage.cfg_crop_simple_resize/8b59fed7/fold0/checkpoints/last.ckpt

fold0.ckpt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2db6d3fb26a05b916341574c83683017e4a04a1c0df8fda4a97ad2314b33f109
-size 81642981

fold1.ckpt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8c806c2ccd21cb4f1d1102e86d8716ed67583f561d4eea6a1761ac4f9bf6a60b
-size 81642981

fold2.ckpt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:cabdc105bb4c3239d1a57ceaaca4306096a017763c1ec1d23adacf6d8c0713ab
-size 81642981

greulich_and_pyle_ages.json

@@ -0,0 +1,65 @@
+{
+    "bone_ages": {
+        "female": [
+            0,
+            3,
+            6,
+            9,
+            12,
+            15,
+            18,
+            24,
+            36,
+            42,
+            50,
+            60,
+            69,
+            82,
+            94,
+            106,
+            120,
+            132,
+            144,
+            156,
+            162,
+            168,
+            180,
+            192,
+            204,
+            216
+        ],
+        "male": [
+            0,
+            3,
+            6,
+            9,
+            12,
+            15,
+            18,
+            24,
+            30,
+            32,
+            36,
+            42,
+            48,
+            54,
+            60,
+            72,
+            84,
+            96,
+            108,
+            120,
+            132,
+            138,
+            150,
+            156,
+            162,
+            168,
+            180,
+            192,
+            204,
+            216,
+            228
+        ]
+    }
+}

net0.pt

@@ -0,0 +1 @@
+../../experiments/boneage/boneage.cfg_female_channel_reg_cls_match_hist/fa77ff59/fold0/checkpoints/last.ckpt

net1.pt

@@ -0,0 +1 @@
+../../experiments/boneage/boneage.cfg_female_channel_reg_cls_match_hist/fa77ff59/fold1/checkpoints/last.ckpt

net2.pt

@@ -0,0 +1 @@
+../../experiments/boneage/boneage.cfg_female_channel_reg_cls_match_hist/fa77ff59/fold2/checkpoints/last.ckpt

requirements.txt

@@ -1,6 +1,7 @@
-gradio==3.45.0
-numpy
-omegaconf
+albumentations
+einops
+grad-cam
+gradio
+scikit-image
 timm
-torch
-
+torch

skp/configs/__init__.py

@@ -0,0 +1,21 @@
+from types import SimpleNamespace
+
+
+class Config(SimpleNamespace):
+
+    def __getattribute__(self, value):
+        # If attribute not specified in config,
+        # return None instead of raise error
+        try:
+            return super().__getattribute__(value)
+        except AttributeError:
+            return None  
+
+    def __str__(self):
+        # pretty print
+        string = ["config"]
+        string.append("=" * len(string[0]))
+        longest_param_name = max([len(k) for k in [*self.__dict__]])
+        for k, v in self.__dict__.items():
+            string.append(f"{k.ljust(longest_param_name)} : {v}")
+        return "\n".join(string)

skp/configs/base.py

@@ -0,0 +1,21 @@
+from types import SimpleNamespace
+
+
+class Config(SimpleNamespace):
+
+    def __getattribute__(self, value):
+        # If attribute not specified in config,
+        # return None instead of raise error
+        try:
+            return super().__getattribute__(value)
+        except AttribuateError:
+            return None  
+
+    def __str__(self):
+        # pretty print
+        string = ["config"]
+        string.append("=" * len(string[0]))
+        longest_param_name = max([len(k) for k in [*self.__dict__]])
+        for k, v in self.__dict__.items():
+            string.append(f"{k.ljust(longest_param_name)} : {v}")
+        return "\n".join(string)

skp/configs/boneage/cfg_baseline.py

@@ -0,0 +1,117 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "classification.net2d_var_embed"
+cfg.backbone = "tf_efficientnetv2_s"
+cfg.embed_num_classes = 2
+cfg.embed_dim = 32
+cfg.pretrained = True
+cfg.num_input_channels = 1
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = 1
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+
+cfg.fold = 0
+cfg.dataset = "simple2d"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/cropped_train_plus_valid/"
+cfg.annotations_file = "/mnt/stor/datasets/bone-age/train_plus_valid_kfold.csv"
+cfg.inputs = "imgfile0"
+cfg.targets = ["bone_age_years"]
+cfg.vars = "female"
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.num_workers = 16
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 1000
+
+cfg.loss = "classification.L1Loss"
+cfg.loss_params = {}
+
+cfg.batch_size = 32
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.MAE", "classification.MSE"]
+cfg.val_metric = "mae_mean"
+cfg.val_track = "min"
+
+cfg.image_height = 512
+cfg.image_width = 512
+
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+    A.PadIfNeeded(
+        min_height=cfg.image_height,
+        min_width=cfg.image_width,
+        border_mode=cv2.BORDER_CONSTANT,
+        p=1,
+    ),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ]
+)
+
+cfg.val_transforms = A.Compose(resize_transforms)

skp/configs/boneage/cfg_crop.py

@@ -0,0 +1,123 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "classification.net2d"
+cfg.backbone = "mobilenetv3_small_100"
+cfg.pretrained = True
+cfg.num_input_channels = 1
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = 4
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+cfg.model_activation_fn = "sigmoid"
+
+cfg.fold = 0
+cfg.dataset = "crop2d"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/train/"
+cfg.annotations_file = (
+    "/mnt/stor/datasets/bone-age/train_with_bounding_box_crop_coords_kfold.csv"
+)
+cfg.inputs = "imgfile"
+cfg.targets = ["x1", "y1", "w", "h"]
+cfg.normalize_crop_coords = True
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.num_workers = 16
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 100
+
+cfg.loss = "classification.L1Loss"
+cfg.loss_params = {}
+
+cfg.batch_size = 16
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.MAE", "classification.MSE"]
+cfg.val_metric = "mae_mean"
+cfg.val_track = "min"
+
+cfg.image_height = 512
+cfg.image_width = 512
+
+bbox_params = A.BboxParams(format="coco")
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+    A.PadIfNeeded(
+        min_height=cfg.image_height,
+        min_width=cfg.image_width,
+        border_mode=cv2.BORDER_CONSTANT,
+        p=1,
+    ),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ],
+    bbox_params=bbox_params,
+)
+
+cfg.val_transforms = A.Compose(
+    resize_transforms,
+    bbox_params=bbox_params,
+)

skp/configs/boneage/cfg_crop_simple_resize.py

@@ -0,0 +1,117 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "classification.net2d"
+cfg.backbone = "mobilenetv3_small_100"
+cfg.pretrained = True
+cfg.num_input_channels = 1
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = 4
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+cfg.model_activation_fn = "sigmoid"
+
+cfg.fold = 0
+cfg.dataset = "crop2d"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/train/"
+cfg.annotations_file = (
+    "/mnt/stor/datasets/bone-age/train_with_bounding_box_crop_coords_kfold.csv"
+)
+cfg.inputs = "imgfile"
+cfg.targets = ["x1", "y1", "w", "h"]
+cfg.normalize_crop_coords = True
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.num_workers = 16
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 200
+
+cfg.loss = "classification.L1Loss"
+cfg.loss_params = {}
+
+cfg.batch_size = 16
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.MAE", "classification.MSE"]
+cfg.val_metric = "mae_mean"
+cfg.val_track = "min"
+
+cfg.image_height = 512
+cfg.image_width = 512
+
+bbox_params = A.BboxParams(format="coco")
+resize_transforms = [
+    A.Resize(height=cfg.image_height, width=cfg.image_width, p=1)
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ],
+    bbox_params=bbox_params,
+)
+
+cfg.val_transforms = A.Compose(
+    resize_transforms,
+    bbox_params=bbox_params,
+)

skp/configs/boneage/cfg_female_channel.py

@@ -0,0 +1,114 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "classification.net2d"
+cfg.backbone = "tf_efficientnetv2_s"
+cfg.pretrained = True
+cfg.num_input_channels = 2
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = 1
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+
+cfg.fold = 0
+cfg.dataset = "boneage.female_channel"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/cropped_train_plus_valid/"
+cfg.annotations_file = "/mnt/stor/datasets/bone-age/train_plus_valid_kfold.csv"
+cfg.inputs = "imgfile0"
+cfg.targets = ["bone_age_years"]
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.num_workers = 16
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 1000
+
+cfg.loss = "classification.L1Loss"
+cfg.loss_params = {}
+
+cfg.batch_size = 32
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.MAE", "classification.MSE"]
+cfg.val_metric = "mae_mean"
+cfg.val_track = "min"
+
+cfg.image_height = 512
+cfg.image_width = 512
+
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+    A.PadIfNeeded(
+        min_height=cfg.image_height,
+        min_width=cfg.image_width,
+        border_mode=cv2.BORDER_CONSTANT,
+        p=1,
+    ),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ]
+)
+
+cfg.val_transforms = A.Compose(resize_transforms)

skp/configs/boneage/cfg_female_channel_MIL.py

@@ -0,0 +1,113 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "MIL.net2d_basic_attn"
+cfg.backbone = "tf_efficientnetv2_s"
+cfg.pretrained = True
+cfg.num_input_channels = 2
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = 1
+cfg.attn_dropout = 0.0
+cfg.attn_version = "v1"
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+
+cfg.fold = 0
+cfg.dataset = "boneage.female_channel_grid_patch"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/cropped_train_plus_valid/"
+cfg.annotations_file = "/mnt/stor/datasets/bone-age/train_plus_valid_kfold.csv"
+cfg.inputs = "imgfile0"
+cfg.targets = ["bone_age_years"]
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.patch_size = 224
+cfg.patch_num_rows = 5
+cfg.patch_num_cols = 3
+cfg.num_workers = 16
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 1000
+
+cfg.loss = "classification.L1Loss"
+cfg.loss_params = {}
+
+cfg.batch_size = 16
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.MAE", "classification.MSE"]
+cfg.val_metric = "mae_mean"
+cfg.val_track = "min"
+
+cfg.image_height = 560
+cfg.image_width = cfg.image_height # not used 
+
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ]
+)
+
+cfg.val_transforms = A.Compose(resize_transforms)

skp/configs/boneage/cfg_female_channel_MIL_lstm.py

@@ -0,0 +1,116 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "MIL.net2d_attn"
+cfg.backbone = "tf_efficientnetv2_s"
+cfg.pretrained = True
+cfg.num_input_channels = 2
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = 1
+cfg.add_lstm = True
+cfg.lstm_dropout = 0.0
+cfg.lstm_num_layers = 1
+cfg.attn_dropout = 0.0
+cfg.attn_version = "v1"
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+
+cfg.fold = 0
+cfg.dataset = "boneage.female_channel_grid_patch"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/cropped_train_plus_valid/"
+cfg.annotations_file = "/mnt/stor/datasets/bone-age/train_plus_valid_kfold.csv"
+cfg.inputs = "imgfile0"
+cfg.targets = ["bone_age_years"]
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.patch_size = 224
+cfg.patch_num_rows = 5
+cfg.patch_num_cols = 3
+cfg.num_workers = 16
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 1000
+
+cfg.loss = "classification.L1Loss"
+cfg.loss_params = {}
+
+cfg.batch_size = 16
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.MAE", "classification.MSE"]
+cfg.val_metric = "mae_mean"
+cfg.val_track = "min"
+
+cfg.image_height = 560
+cfg.image_width = cfg.image_height # not used 
+
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ]
+)
+
+cfg.val_transforms = A.Compose(resize_transforms)

skp/configs/boneage/cfg_female_channel_MIL_transformer.py

@@ -0,0 +1,117 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "MIL.net2d_attn"
+cfg.backbone = "tf_efficientnetv2_s"
+cfg.pretrained = True
+cfg.num_input_channels = 2
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = 1
+cfg.reduce_feature_dim = 256
+cfg.add_transformer = True
+cfg.transformer_dropout = 0.0
+cfg.transformer_num_layers = 1
+cfg.attn_dropout = 0.0
+cfg.attn_version = "v1"
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+
+cfg.fold = 0
+cfg.dataset = "boneage.female_channel_grid_patch"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/cropped_train_plus_valid/"
+cfg.annotations_file = "/mnt/stor/datasets/bone-age/train_plus_valid_kfold.csv"
+cfg.inputs = "imgfile0"
+cfg.targets = ["bone_age_years"]
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.patch_size = 224
+cfg.patch_num_rows = 5
+cfg.patch_num_cols = 3
+cfg.num_workers = 16
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 1000
+
+cfg.loss = "classification.L1Loss"
+cfg.loss_params = {}
+
+cfg.batch_size = 16
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.MAE", "classification.MSE"]
+cfg.val_metric = "mae_mean"
+cfg.val_track = "min"
+
+cfg.image_height = 560
+cfg.image_width = cfg.image_height # not used 
+
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ]
+)
+
+cfg.val_transforms = A.Compose(resize_transforms)

skp/configs/boneage/cfg_female_channel_reg_cls.py

@@ -0,0 +1,115 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "classification.net2d_multihead"
+cfg.backbone = "tf_efficientnetv2_s"
+cfg.pretrained = True
+cfg.num_input_channels = 2
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = [1, 240]
+cfg.num_heads = 2
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+
+cfg.fold = 0
+cfg.dataset = "boneage.female_channel"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/cropped_train_plus_valid/"
+cfg.annotations_file = "/mnt/stor/datasets/bone-age/train_plus_valid_kfold.csv"
+cfg.inputs = "imgfile0"
+cfg.targets = ["bone_age"]
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.num_workers = 16
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 1000
+
+cfg.loss = "classification.DoubleL1Loss"
+cfg.loss_params = {"reg_weight": 1.0, "cls_weight": 0.4}
+
+cfg.batch_size = 32
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.DoubleMAE"]
+cfg.val_metric = "mae_reg"
+cfg.val_track = "min"
+
+cfg.image_height = 512
+cfg.image_width = 512
+
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+    A.PadIfNeeded(
+        min_height=cfg.image_height,
+        min_width=cfg.image_width,
+        border_mode=cv2.BORDER_CONSTANT,
+        p=1,
+    ),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ]
+)
+
+cfg.val_transforms = A.Compose(resize_transforms)

skp/configs/boneage/cfg_female_channel_reg_cls_clip_outliers_aug.py

@@ -0,0 +1,119 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "classification.net2d_multihead"
+cfg.backbone = "tf_efficientnetv2_s"
+cfg.pretrained = True
+cfg.num_input_channels = 2
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = [1, 240]
+cfg.num_heads = 2
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+
+cfg.fold = 0
+cfg.dataset = "boneage.female_channel"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/cropped_train_plus_valid/"
+cfg.annotations_file = "/mnt/stor/datasets/bone-age/train_plus_valid_kfold.csv"
+cfg.inputs = "imgfile0"
+cfg.targets = ["bone_age"]
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.num_workers = 16
+cfg.clip_outlier_pixels_and_rescale = True
+cfg.clip_as_data_aug = True
+cfg.clip_proba = 0.5
+cfg.clip_bounds = (1, 99)
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 1000
+
+cfg.loss = "classification.DoubleL1Loss"
+cfg.loss_params = {"reg_weight": 1.0, "cls_weight": 0.4}
+
+cfg.batch_size = 32
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.DoubleMAE"]
+cfg.val_metric = "mae_reg"
+cfg.val_track = "min"
+
+cfg.image_height = 512
+cfg.image_width = 512
+
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+    A.PadIfNeeded(
+        min_height=cfg.image_height,
+        min_width=cfg.image_width,
+        border_mode=cv2.BORDER_CONSTANT,
+        p=1,
+    ),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ]
+)
+
+cfg.val_transforms = A.Compose(resize_transforms)

skp/configs/boneage/cfg_female_channel_reg_cls_match_hist.py

@@ -0,0 +1,116 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "classification.net2d_multihead"
+cfg.backbone = "tf_efficientnetv2_s"
+cfg.pretrained = True
+cfg.num_input_channels = 2
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = [1, 240]
+cfg.num_heads = 2
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+
+cfg.fold = 0
+cfg.dataset = "boneage.female_channel_match_hist"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/cropped_train_plus_valid/"
+cfg.annotations_file = "/mnt/stor/datasets/bone-age/train_plus_valid_kfold.csv"
+cfg.ref_image_match_hist = "/mnt/stor/datasets/bone-age/reference_cropped_image_for_histogram_matching.png"
+cfg.inputs = "imgfile0"
+cfg.targets = ["bone_age"]
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.num_workers = 16
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 1000
+
+cfg.loss = "classification.DoubleL1Loss"
+cfg.loss_params = {"reg_weight": 1.0, "cls_weight": 0.4}
+
+cfg.batch_size = 32
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.DoubleMAE"]
+cfg.val_metric = "mae_reg"
+cfg.val_track = "min"
+
+cfg.image_height = 512
+cfg.image_width = 512
+
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+    A.PadIfNeeded(
+        min_height=cfg.image_height,
+        min_width=cfg.image_width,
+        border_mode=cv2.BORDER_CONSTANT,
+        p=1,
+    ),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ]
+)
+
+cfg.val_transforms = A.Compose(resize_transforms)

skp/configs/boneage/cfg_female_channel_with_cls.py

@@ -0,0 +1,115 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "classification.net2d_multihead"
+cfg.backbone = "tf_efficientnetv2_s"
+cfg.pretrained = True
+cfg.num_input_channels = 2
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = [1, 24]
+cfg.num_heads = 2
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+
+cfg.fold = 0
+cfg.dataset = "boneage.female_channel_with_cls"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/cropped_train_plus_valid/"
+cfg.annotations_file = "/mnt/stor/datasets/bone-age/train_plus_valid_kfold.csv"
+cfg.inputs = "imgfile0"
+cfg.targets = ["bone_age_years", "bone_age_categorical"]
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.num_workers = 16
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 1000
+
+cfg.loss = "classification.L1CELoss"
+cfg.loss_params = {"l1_weight": 1.0, "ce_weight": 0.2}
+
+cfg.batch_size = 32
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.MAE_Accuracy"]
+cfg.val_metric = "mae_mean"
+cfg.val_track = "min"
+
+cfg.image_height = 512
+cfg.image_width = 512
+
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+    A.PadIfNeeded(
+        min_height=cfg.image_height,
+        min_width=cfg.image_width,
+        border_mode=cv2.BORDER_CONSTANT,
+        p=1,
+    ),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ]
+)
+
+cfg.val_transforms = A.Compose(resize_transforms)

skp/configs/boneage/cfg_female_channel_with_cls_clip_outliers.py

@@ -0,0 +1,117 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "classification.net2d_multihead"
+cfg.backbone = "tf_efficientnetv2_s"
+cfg.pretrained = True
+cfg.num_input_channels = 2
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = [1, 24]
+cfg.num_heads = 2
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+
+cfg.fold = 0
+cfg.dataset = "boneage.female_channel_with_cls"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/cropped_train_plus_valid/"
+cfg.annotations_file = "/mnt/stor/datasets/bone-age/train_plus_valid_kfold.csv"
+cfg.inputs = "imgfile0"
+cfg.targets = ["bone_age_years", "bone_age_categorical"]
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.num_workers = 16
+cfg.clip_outlier_pixels_and_rescale = True
+cfg.clip_bounds = (1, 99)
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 1000
+
+cfg.loss = "classification.L1CELoss"
+cfg.loss_params = {"l1_weight": 1.0, "ce_weight": 0.2}
+
+cfg.batch_size = 32
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.MAE_Accuracy"]
+cfg.val_metric = "mae_mean"
+cfg.val_track = "min"
+
+cfg.image_height = 512
+cfg.image_width = 512
+
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+    A.PadIfNeeded(
+        min_height=cfg.image_height,
+        min_width=cfg.image_width,
+        border_mode=cv2.BORDER_CONSTANT,
+        p=1,
+    ),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ]
+)
+
+cfg.val_transforms = A.Compose(resize_transforms)

skp/configs/boneage/cfg_female_channel_with_cls_clip_outliers_aug.py

@@ -0,0 +1,119 @@
+import albumentations as A
+import cv2
+
+from skp.configs import Config
+
+
+cfg = Config()
+cfg.neptune_mode = "async"
+
+cfg.save_dir = "/home/ian/projects/SKP/experiments/boneage/"
+cfg.project = "gradientecho/SKP"
+
+cfg.task = "classification"
+
+cfg.model = "classification.net2d_multihead"
+cfg.backbone = "tf_efficientnetv2_s"
+cfg.pretrained = True
+cfg.num_input_channels = 2
+cfg.pool = "gem"
+cfg.pool_params = {"p": 3}
+cfg.dropout = 0.1
+cfg.num_classes = [1, 24]
+cfg.num_heads = 2
+cfg.normalization = "-1_1"
+cfg.normalization_params = {"min": 0, "max": 255}
+cfg.backbone_img_size = False
+
+cfg.fold = 0
+cfg.dataset = "boneage.female_channel_with_cls"
+cfg.data_dir = "/mnt/stor/datasets/bone-age/cropped_train_plus_valid/"
+cfg.annotations_file = "/mnt/stor/datasets/bone-age/train_plus_valid_kfold.csv"
+cfg.inputs = "imgfile0"
+cfg.targets = ["bone_age_years", "bone_age_categorical"]
+cfg.cv2_load_flag = cv2.IMREAD_GRAYSCALE
+cfg.num_workers = 16
+cfg.clip_outlier_pixels_and_rescale = True
+cfg.clip_as_data_aug = True
+cfg.clip_proba = 0.5
+cfg.clip_bounds = (1, 99)
+cfg.pin_memory = True
+cfg.persistent_workers = True
+cfg.sampler = "IterationBasedSampler"
+cfg.num_iterations_per_epoch = 1000
+
+cfg.loss = "classification.L1CELoss"
+cfg.loss_params = {"l1_weight": 1.0, "ce_weight": 0.2}
+
+cfg.batch_size = 32
+cfg.num_epochs = 10
+cfg.optimizer = "AdamW"
+cfg.optimizer_params = {"lr": 3e-4}
+
+cfg.scheduler = "LinearWarmupCosineAnnealingLR"
+cfg.scheduler_params = {"pct_start": 0.1, "div_factor": 100, "final_div_factor": 1_000}
+cfg.scheduler_interval = "step"
+
+cfg.val_batch_size = cfg.batch_size * 2
+cfg.metrics = ["classification.MAE_Accuracy"]
+cfg.val_metric = "mae_mean"
+cfg.val_track = "min"
+
+cfg.image_height = 512
+cfg.image_width = 512
+
+resize_transforms = [
+    A.LongestMaxSize(max_size=cfg.image_height, p=1),
+    A.PadIfNeeded(
+        min_height=cfg.image_height,
+        min_width=cfg.image_width,
+        border_mode=cv2.BORDER_CONSTANT,
+        p=1,
+    ),
+]
+
+cfg.train_transforms = A.Compose(
+    resize_transforms
+    + [
+        A.VerticalFlip(p=0.5),
+        A.HorizontalFlip(p=0.5),
+        A.SomeOf(
+            [
+                A.ShiftScaleRotate(
+                    shift_limit=0.2,
+                    scale_limit=0.0,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.2,
+                    rotate_limit=0,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.ShiftScaleRotate(
+                    shift_limit=0.0,
+                    scale_limit=0.0,
+                    rotate_limit=30,
+                    border_mode=cv2.BORDER_CONSTANT,
+                    p=1,
+                ),
+                A.GaussianBlur(p=1),
+                A.GaussNoise(p=1),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.3, brightness_limit=0.0, p=1
+                ),
+                A.RandomBrightnessContrast(
+                    contrast_limit=0.0, brightness_limit=0.3, p=1
+                ),
+            ],
+            n=3,
+            p=0.9,
+            replace=False,
+        ),
+    ]
+)
+
+cfg.val_transforms = A.Compose(resize_transforms)

skp/models/MIL/net2d_attn.py

@@ -0,0 +1,286 @@
+"""
+2D model for multiple instance learning (MIL)
+Performs attention over bag of features (i.e., attention-weighted mean of features)
+Option to add LSTM or Transformer before attention aggregation
+Uses timm backbones
+"""
+
+import re
+import torch
+import torch.nn as nn
+
+from einops import rearrange
+from timm import create_model
+from typing import Dict, Optional, Tuple
+
+from skp.configs.base import Config
+from skp.models.modules import FeatureReduction
+from skp.models.pooling import get_pool_layer
+
+
+class Attention(nn.Module):
+    """
+    Given a batch containing bags of features (B, N, D),
+    generate attention scores over the features in a bag, N,
+    and perform an attention-weighted mean of the features (B, D)
+    """
+
+    def __init__(self, embed_dim: int, dropout: float = 0.0, version: str = "v1"):
+        super().__init__()
+        version = version.lower()
+        if version == "v1":
+            self.mlp = nn.Sequential(
+                nn.Tanh(), nn.Dropout(dropout), nn.Linear(embed_dim, 1)
+            )
+        elif version == "v2":
+            self.mlp = nn.Sequential(
+                nn.Linear(embed_dim, embed_dim),
+                nn.Tanh(),
+                nn.Dropout(dropout),
+                nn.Linear(embed_dim, 1),
+            )
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor]:
+        a = self.mlp(x)
+        a = a.softmax(dim=1)
+        x = (x * a).sum(dim=1)
+        return x, a
+
+
+class BiLSTM(nn.Module):
+    def __init__(self, embed_dim: int, dropout: float = 0.0, num_layers: int = 1):
+        super().__init__()
+        self.lstm = nn.LSTM(
+            input_size=embed_dim,
+            hidden_size=embed_dim // 2,
+            num_layers=num_layers,
+            bias=True,
+            batch_first=True,
+            dropout=dropout,
+            bidirectional=True,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.lstm(x)
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        nhead: int = 16,
+        activation: str = "gelu",
+    ):
+        super().__init__()
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim,
+            nhead=nhead,
+            dim_feedforward=embed_dim,
+            dropout=dropout,
+            activation=activation,
+            batch_first=True,
+            norm_first=False,
+            bias=True,
+        )
+        self.T = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+
+    def forward(
+        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        return self.T(x, mask=mask)
+
+
+class Net(nn.Module):
+    def __init__(self, cfg: Config):
+        super().__init__()
+        self.cfg = cfg
+        backbone_args = {
+            "pretrained": self.cfg.pretrained,
+            "num_classes": 0,
+            "global_pool": "",
+            "features_only": self.cfg.features_only,
+            "in_chans": self.cfg.num_input_channels,
+        }
+        if self.cfg.backbone_img_size:
+            # some models require specifying image size (e.g., coatnet)
+            if "efficientvit" in self.cfg.backbone:
+                backbone_args["img_size"] = self.cfg.image_height
+            else:
+                backbone_args["img_size"] = (
+                    self.cfg.image_height,
+                    self.cfg.image_width,
+                )
+        self.backbone = create_model(self.cfg.backbone, **backbone_args)
+        # get feature dim by passing sample through net
+        self.feature_dim = self.backbone(
+            torch.randn(
+                (
+                    2,
+                    self.cfg.num_input_channels,
+                    self.cfg.image_height,
+                    self.cfg.image_width,
+                )
+            )
+        ).size(
+            -1 if "xcit" in self.cfg.backbone else 1
+        )  # xcit models are channels-last
+
+        self.feature_dim = self.feature_dim * (2 if self.cfg.pool == "catavgmax" else 1)
+        self.pooling = get_pool_layer(self.cfg, dim=2)
+
+        if isinstance(self.cfg.reduce_feature_dim, int):
+            self.backbone = nn.Sequential(
+                self.backbone,
+                FeatureReduction(self.feature_dim, self.cfg.reduce_feature_dim),
+            )
+            self.feature_dim = self.cfg.reduce_feature_dim
+
+        if self.cfg.add_lstm:
+            self.pre_attn = BiLSTM(
+                embed_dim=self.feature_dim,
+                dropout=self.cfg.lstm_dropout or 0.0,
+                num_layers=self.cfg.lstm_num_layers or 1,
+            )
+        elif self.cfg.add_transformer:
+            self.pre_attn = Transformer(
+                embed_dim=self.feature_dim,
+                dropout=self.cfg.transformer_dropout or 0.0,
+                num_layers=self.cfg.transformer_num_layers or 1,
+                nhead=self.cfg.transformer_nhead or 16,
+                activation=self.cfg.transformer_act or "gelu",
+            )
+        else:
+            self.pre_attn = nn.Identity()
+
+        self.attn = Attention(
+            self.feature_dim,
+            dropout=self.cfg.attn_dropout,
+            version=self.cfg.attn_version or "v1",
+        )
+        self.dropout = nn.Dropout(p=self.cfg.dropout)
+        self.linear = nn.Linear(self.feature_dim, self.cfg.num_classes)
+
+        if self.cfg.load_pretrained_backbone:
+            print(
+                f"Loading pretrained backbone from {self.cfg.load_pretrained_backbone} ..."
+            )
+            weights = torch.load(
+                self.cfg.load_pretrained_backbone,
+                map_location=lambda storage, loc: storage,
+            )["state_dict"]
+            # Replace model prefix as this does not exist in Net
+            weights = {re.sub(r"^model.", "", k): v for k, v in weights.items()}
+            # Get backbone only
+            weights = {
+                re.sub(r"^backbone.", "", k): v
+                for k, v in weights.items()
+                if "backbone" in k
+            }
+            self.backbone.load_state_dict(weights)
+
+        self.criterion = None
+
+        self.backbone_frozen = False
+        if self.cfg.freeze_backbone:
+            self.freeze_backbone()
+
+    def normalize(self, x: torch.Tensor) -> torch.Tensor:
+        if self.cfg.normalization == "-1_1":
+            mini, maxi = (
+                self.cfg.normalization_params["min"],
+                self.cfg.normalization_params["max"],
+            )
+            x = x - mini
+            x = x / (maxi - mini)
+            x = x - 0.5
+            x = x * 2.0
+        elif self.cfg.normalization == "0_1":
+            mini, maxi = (
+                self.cfg.normalization_params["min"],
+                self.cfg.normalization_params["max"],
+            )
+            x = x - mini
+            x = x / (maxi - mini)
+        elif self.cfg.normalization == "mean_sd":
+            mean, sd = (
+                self.cfg.normalization_params["mean"],
+                self.cfg.normalization_params["sd"],
+            )
+            x = (x - mean) / sd
+        elif self.cfg.normalization == "per_channel_mean_sd":
+            mean, sd = (
+                self.cfg.normalization_params["mean"],
+                self.cfg.normalization_params["sd"],
+            )
+            assert len(mean) == len(sd) == x.size(1)
+            mean, sd = torch.tensor(mean).unsqueeze(0), torch.tensor(sd).unsqueeze(0)
+            for i in range(x.ndim - 2):
+                mean, sd = mean.unsqueeze(-1), sd.unsqueeze(-1)
+            x = (x - mean) / sd
+        elif self.cfg.normalization == "none":
+            x = x
+        return x
+
+    def forward(
+        self,
+        batch: Dict,
+        return_loss: bool = False,
+        return_features: bool = False,
+        return_attn_scores: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        x = batch["x"]
+        y = batch.get("y", None)
+
+        if return_loss:
+            assert y is not None
+
+        b, n = x.shape[:2]
+        x = rearrange(x, "b n c h w -> (b n) c h w")
+        features = self.extract_features(x, normalize=True)
+        features = rearrange(features, "(b n) d -> b n d", b=b, n=n)
+        if isinstance(self.pre_attn, Transformer):
+            features = self.pre_attn(features, mask=batch.get("mask", None))
+        else:
+            features = self.pre_attn(features)
+        features, attn_scores = self.attn(features)
+
+        if self.cfg.multisample_dropout:
+            logits = torch.stack(
+                [self.linear(self.dropout(features)) for _ in range(5)]
+            ).mean(0)
+        else:
+            logits = self.linear(self.dropout(features))
+
+        if self.cfg.model_activation_fn == "sigmoid":
+            logits = logits.sigmoid()
+        elif self.cfg.model_activation_fn == "softmax":
+            logits = logits.softmax(dim=1)
+
+        out = {"logits": logits}
+        if return_features:
+            out["features"] = features
+        if return_attn_scores:
+            out["attn_scores"] = attn_scores
+        if return_loss:
+            loss = self.criterion(out, batch)
+            if isinstance(loss, dict):
+                out.update(loss)
+            else:
+                out["loss"] = loss
+
+        return out
+
+    def extract_features(self, x: torch.Tensor, normalize: bool = True) -> torch.Tensor:
+        x = self.normalize(x) if normalize else x
+        return self.pooling(self.backbone(x))
+
+    def freeze_backbone(self) -> None:
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+        self.backbone_frozen = True
+
+    def set_criterion(self, loss: nn.Module) -> None:
+        self.criterion = loss

skp/models/MIL/net2d_basic_attn.py

@@ -0,0 +1,284 @@
+"""
+2D model for multiple instance learning (MIL)
+Performs attention over bag of features (i.e., attention-weighted mean of features)
+Uses timm backbones
+"""
+
+import re
+import torch
+import torch.nn as nn
+
+from einops import rearrange
+from timm import create_model
+from typing import Dict, Optional, Tuple
+
+from skp.configs.base import Config
+from skp.models.modules import FeatureReduction
+from skp.models.pooling import get_pool_layer
+
+
+class Attention(nn.Module):
+    """
+    Given a batch containing bags of features (B, N, D),
+    generate attention scores over the features in a bag, N,
+    and perform an attention-weighted mean of the features (B, D)
+    """
+
+    def __init__(self, embed_dim: int, dropout: float = 0.0, version: str = "v1"):
+        super().__init__()
+        version = version.lower()
+        if version == "v1":
+            self.mlp = nn.Sequential(
+                nn.Tanh(), nn.Dropout(dropout), nn.Linear(embed_dim, 1)
+            )
+        elif version == "v2":
+            self.mlp = nn.Sequential(
+                nn.Linear(embed_dim, embed_dim),
+                nn.Tanh(),
+                nn.Dropout(dropout),
+                nn.Linear(embed_dim, 1),
+            )
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor]:
+        a = self.mlp(x)
+        a = a.softmax(dim=1)
+        x = (x * a).sum(dim=1)
+        return x, a
+
+
+class BiLSTM(nn.Module):
+    def __init__(self, embed_dim: int, dropout: float = 0.0, num_layers: int = 1):
+        super().__init__()
+        self.lstm = nn.LSTM(
+            input_size=embed_dim,
+            hidden_size=embed_dim // 2,
+            num_layers=num_layers,
+            bias=True,
+            batch_first=True,
+            dropout=dropout,
+            bidirectional=True,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.lstm(x)
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        nheads: int = 16,
+        activation: str = "gelu",
+    ):
+        super().__init__()
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim,
+            dim_feedforward=embed_dim,
+            dropout=dropout,
+            activation=activation,
+            batch_first=True,
+            norm_first=False,
+            bias=True,
+        )
+        self.T = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+
+    def forward(
+        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        return self.T(x, mask=mask)
+
+
+class Net(nn.Module):
+    def __init__(self, cfg: Config):
+        super().__init__()
+        self.cfg = cfg
+        backbone_args = {
+            "pretrained": self.cfg.pretrained,
+            "num_classes": 0,
+            "global_pool": "",
+            "features_only": self.cfg.features_only,
+            "in_chans": self.cfg.num_input_channels,
+        }
+        if self.cfg.backbone_img_size:
+            # some models require specifying image size (e.g., coatnet)
+            if "efficientvit" in self.cfg.backbone:
+                backbone_args["img_size"] = self.cfg.image_height
+            else:
+                backbone_args["img_size"] = (
+                    self.cfg.image_height,
+                    self.cfg.image_width,
+                )
+        self.backbone = create_model(self.cfg.backbone, **backbone_args)
+        # get feature dim by passing sample through net
+        self.feature_dim = self.backbone(
+            torch.randn(
+                (
+                    2,
+                    self.cfg.num_input_channels,
+                    self.cfg.image_height,
+                    self.cfg.image_width,
+                )
+            )
+        ).size(
+            -1 if "xcit" in self.cfg.backbone else 1
+        )  # xcit models are channels-last
+
+        self.feature_dim = self.feature_dim * (2 if self.cfg.pool == "catavgmax" else 1)
+        self.pooling = get_pool_layer(self.cfg, dim=2)
+
+        if isinstance(self.cfg.reduce_feature_dim, int):
+            self.backbone = nn.Sequential(
+                self.backbone,
+                FeatureReduction(self.feature_dim, self.cfg.reduce_feature_dim),
+            )
+            self.feature_dim = self.cfg.reduce_feature_dim
+
+        if self.cfg.add_lstm:
+            self.pre_attn = BiLSTM(
+                embed_dim=self.feature_dim,
+                dropout=self.cfg.lstm_dropout or 0.0,
+                num_layers=self.cfg.lstm_num_layers or 1,
+            )
+        elif self.cfg.add_transformer:
+            self.pre_attn = Transformer(
+                embed_dim=self.feature_dim,
+                dropout=self.transformer_dropout or 0.0,
+                num_layers=self.transformer_num_layers or 1,
+                nheads=self.transformer_nheads or 16,
+                activation=self.transformer_act or "gelu",
+            )
+        else:
+            self.pre_attn = nn.Identity()
+
+        self.attn = Attention(
+            self.feature_dim,
+            dropout=self.cfg.attn_dropout,
+            version=self.cfg.attn_version or "v1",
+        )
+        self.dropout = nn.Dropout(p=self.cfg.dropout)
+        self.linear = nn.Linear(self.feature_dim, self.cfg.num_classes)
+
+        if self.cfg.load_pretrained_backbone:
+            print(
+                f"Loading pretrained backbone from {self.cfg.load_pretrained_backbone} ..."
+            )
+            weights = torch.load(
+                self.cfg.load_pretrained_backbone,
+                map_location=lambda storage, loc: storage,
+            )["state_dict"]
+            # Replace model prefix as this does not exist in Net
+            weights = {re.sub(r"^model.", "", k): v for k, v in weights.items()}
+            # Get backbone only
+            weights = {
+                re.sub(r"^backbone.", "", k): v
+                for k, v in weights.items()
+                if "backbone" in k
+            }
+            self.backbone.load_state_dict(weights)
+
+        self.criterion = None
+
+        self.backbone_frozen = False
+        if self.cfg.freeze_backbone:
+            self.freeze_backbone()
+
+    def normalize(self, x: torch.Tensor) -> torch.Tensor:
+        if self.cfg.normalization == "-1_1":
+            mini, maxi = (
+                self.cfg.normalization_params["min"],
+                self.cfg.normalization_params["max"],
+            )
+            x = x - mini
+            x = x / (maxi - mini)
+            x = x - 0.5
+            x = x * 2.0
+        elif self.cfg.normalization == "0_1":
+            mini, maxi = (
+                self.cfg.normalization_params["min"],
+                self.cfg.normalization_params["max"],
+            )
+            x = x - mini
+            x = x / (maxi - mini)
+        elif self.cfg.normalization == "mean_sd":
+            mean, sd = (
+                self.cfg.normalization_params["mean"],
+                self.cfg.normalization_params["sd"],
+            )
+            x = (x - mean) / sd
+        elif self.cfg.normalization == "per_channel_mean_sd":
+            mean, sd = (
+                self.cfg.normalization_params["mean"],
+                self.cfg.normalization_params["sd"],
+            )
+            assert len(mean) == len(sd) == x.size(1)
+            mean, sd = torch.tensor(mean).unsqueeze(0), torch.tensor(sd).unsqueeze(0)
+            for i in range(x.ndim - 2):
+                mean, sd = mean.unsqueeze(-1), sd.unsqueeze(-1)
+            x = (x - mean) / sd
+        elif self.cfg.normalization == "none":
+            x = x
+        return x
+
+    def forward(
+        self,
+        batch: Dict,
+        return_loss: bool = False,
+        return_features: bool = False,
+        return_attn_scores: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        x = batch["x"]
+        y = batch.get("y", None)
+
+        if return_loss:
+            assert y is not None
+
+        b, n = x.shape[:2]
+        x = rearrange(x, "b n c h w -> (b n) c h w")
+        features = self.extract_features(x, normalize=True)
+        features = rearrange(features, "(b n) d -> b n d", b=b, n=n)
+        if isinstance(self.pre_attn, Transformer):
+            features = self.pre_attn(features, mask=batch.get("mask", None))
+        else:
+            features = self.pre_attn(features)
+        features, attn_scores = self.attn(features)
+
+        if self.cfg.multisample_dropout:
+            logits = torch.stack(
+                [self.linear(self.dropout(features)) for _ in range(5)]
+            ).mean(0)
+        else:
+            logits = self.linear(self.dropout(features))
+
+        if self.cfg.model_activation_fn == "sigmoid":
+            logits = logits.sigmoid()
+        elif self.cfg.model_activation_fn == "softmax":
+            logits = logits.softmax(dim=1)
+
+        out = {"logits": logits}
+        if return_features:
+            out["features"] = features
+        if return_attn_scores:
+            out["attn_scores"] = attn_scores
+        if return_loss:
+            loss = self.criterion(out, batch)
+            if isinstance(loss, dict):
+                out.update(loss)
+            else:
+                out["loss"] = loss
+
+        return out
+
+    def extract_features(self, x: torch.Tensor, normalize: bool = True) -> torch.Tensor:
+        x = self.normalize(x) if normalize else x
+        return self.pooling(self.backbone(x))
+
+    def freeze_backbone(self) -> None:
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+        self.backbone_frozen = True
+
+    def set_criterion(self, loss: nn.Module) -> None:
+        self.criterion = loss