import os import torch import wandb from dataset import BreastMRI_ABMIL_NII, MammographyDataset, BreastMRI, BreastMRI_ABMIL #from src.models_all.models import create_soat_cnn_models import torch.optim as optim import pandas as pd import argparse import torch.nn as nn from pathlib import Path from torch.utils.data import WeightedRandomSampler from torchinfo import summary from src.trainer_tester.supervised_trainer import FullySupervisedMultiViewTrainer, MultiInstanceTrainer from src.utils.utils_functions import compute_sample_weights, log_scalar_values, set_seed from src.data_loader.mammo_transforms import BasicTransforms, TrainTransform, TrainTransformBaseline from src.utils.plotter import display_confusion_matrix from sklearn.metrics import classification_report import numpy as np import pandas as pd from pathlib import Path from src.utils.dataset_utils import extract_target_names #import utils import torch import torch.nn as nn import timm from focal_loss.focal_loss import FocalLoss class MultiViewSwinModel(nn.Module): def __init__(self, model_name="swin_tiny_patch4_window7_224.ms_in22k", pretrained=True, num_classes=3, attn_dim=1024, num_heads=4): super(MultiViewSwinModel, self).__init__() # Load backbone without classification head self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0) self.feature_dim = self.backbone.num_features # should be 1024 for swin_base # Self-attention across views (3 views → sequence length 3) self.attn = nn.MultiheadAttention(embed_dim=self.feature_dim, num_heads=num_heads, batch_first=True) # Final classifier self.classifier = nn.Linear(self.feature_dim, num_classes) def extract_feat(self, x): # Get unpooled features (B, H*W, C) and then pooled vector (B, C) features = self.backbone.forward_features(x) # (B, H*W, C) pooled = self.backbone.forward_head(features, pre_logits=True) # (B, C) return pooled def forward(self, img1, img2, img3): # Per-view embeddings f1 = self.extract_feat(img1) # [B, C] f2 = self.extract_feat(img2) f3 = self.extract_feat(img3) # Stack views to form a sequence [B, 3, C] views = torch.stack([f1, f2, f3], dim=1) # Self-attention across views attn_out, _ = self.attn(views, views, views) # [B, 3, C] # Aggregate (mean pooling across 3 views) fused = attn_out.mean(dim=1) # [B, C] return self.classifier(fused) class MultiViewSwinCrossAttn(nn.Module): """ Cross‑attention Swin for 3‑view breast‑MRI (pre / post / subtraction). """ def __init__( self, model_name: str = "swin_tiny_patch4_window7_224.ms_in22k", pretrained: bool = True, num_classes: int = 3, num_xattn_layers: int = 2, num_heads: int = 4, dropout: float = 0.1, ): super().__init__() # 1) Swin backbone WITHOUT final linear head self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0) self.embed_dim = self.backbone.num_features # 1024 for swin‑base # 2) Learnable CLS token (like ViT) and view‑type embeddings self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim)) self.view_embed = nn.Parameter(torch.zeros(3, 1, self.embed_dim)) # pre/post/sub # 3) Cross‑attention encoder (TransformerEncoder) encoder_layer = nn.TransformerEncoderLayer( d_model=self.embed_dim, nhead=num_heads, dim_feedforward=self.embed_dim * 4, dropout=dropout, activation="gelu", batch_first=True, norm_first=True, ) self.xattn = nn.TransformerEncoder(encoder_layer, num_layers=num_xattn_layers) # 4) Classification head self.norm = nn.LayerNorm(self.embed_dim) self.classifier = nn.Linear(self.embed_dim, num_classes) # init nn.init.trunc_normal_(self.cls_token, std=0.02) nn.init.trunc_normal_(self.view_embed, std=0.02) # ------------------------------------------------------------------ # def _img_embed(self, x: torch.Tensor) -> torch.Tensor: """ Forward a single image through Swin and return pooled vector (B, C). """ feats = self.backbone.forward_features(x) # (B, H*W, C) pooled = self.backbone.forward_head(feats, pre_logits=True) # (B, C) return pooled def forward(self, pre, post, sub): """ Args: pre, post, sub: tensors (B, 3, 224, 224) Returns: logits (B, num_classes) """ # 1) Per‑view embeddings z_pre = self._img_embed(pre) # (B, C) z_post = self._img_embed(post) z_sub = self._img_embed(sub) # 2) Assemble sequence: [CLS] + three view tokens B = z_pre.size(0) cls = self.cls_token.expand(B, -1, -1) # (B, 1, C) views = torch.stack([z_pre, z_post, z_sub], dim=1) # (B, 3, C) # add view‑type embeddings (broadcast over batch) views = views + self.view_embed.transpose(0, 1) # (B, 3, C) tokens = torch.cat([cls, views], dim=1) # (B, 4, C) # 3) Cross‑attention tokens = self.xattn(tokens) # (B, 4, C) # 4) CLS pooling → head out = self.classifier(self.norm(tokens[:, 0])) # (B, num_classes) return out import torch import torch.nn as nn import timm import torch import torch.nn as nn import timm import math import numpy as np class CrossModalAttentionABMIL_Swin(nn.Module): """ Swin‑based Multiple‑Instance model with per‑slice cross‑modal attention over (pre, post, sub) features, followed by ABMIL pooling. Input : (B, 32, 3, 224, 224) # 3 modalities stacked in channel dim Output : logits (B, num_classes), attention‑per‑slice (B, 32) """ def __init__( self, model_name: str = "swin_tiny_patch4_window7_224.ms_in22k", pretrained: bool = True, num_classes: int = 3, hidden_dim: int = 256, cross_attn_heads: int = 4, dropout: float = 0.1, ): super().__init__() # ------------------------------------------------------------ # 1) Shared Swin backbone (no classifier head) # ------------------------------------------------------------ self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0) self.embed_dim = self.backbone.num_features # 768 for swin_tiny # ------------------------------------------------------------ # 2) Per‑slice cross‑modal fusion # Input: three embeddings (pre, post, sub) -> fused embedding # ------------------------------------------------------------ encoder_layer = nn.TransformerEncoderLayer( d_model=self.embed_dim, nhead=cross_attn_heads, dim_feedforward=self.embed_dim * 4, dropout=dropout, activation="gelu", batch_first=True, norm_first=True, ) # One small transformer (2 layers) that attends across 3 tokens self.slice_fuser = nn.TransformerEncoder(encoder_layer, num_layers=2) # Token type embedding to differentiate modalities self.mod_embed = nn.Parameter(torch.randn(3, 1, self.embed_dim)) # ------------------------------------------------------------ # 3) ABMIL Attention across 32 fused slices # ------------------------------------------------------------ self.attn_V = nn.Linear(self.embed_dim, hidden_dim) self.attn_U = nn.Linear(hidden_dim, 1) # ------------------------------------------------------------ # 4) Classifier head # ------------------------------------------------------------ self.norm = nn.LayerNorm(self.embed_dim) self.dropout = nn.Dropout(dropout) self.classifier = nn.Linear(self.embed_dim, num_classes) # Init nn.init.trunc_normal_(self.mod_embed, std=0.02) nn.init.trunc_normal_(self.attn_V.weight, std=0.02) nn.init.trunc_normal_(self.attn_U.weight, std=0.02) # ------------------------------------------------------------ def _encode_slices(self, x: torch.Tensor) -> torch.Tensor: """ x : (B*N, 1, 224, 224) -> Swin pooled feats (B*N, embed_dim) """ feats = self.backbone.forward_features(x) return self.backbone.forward_head(feats, pre_logits=True) # ------------------------------------------------------------ def forward(self, volume: torch.Tensor): """ volume : (B, 32, 3, 224, 224) """ B, N, C_img, H, W = volume.shape # C_img = 3 modalities assert C_img == 3, "Expect channel dim = 3 (pre, post, sub)" # -------------------------------------------------------- # 1) Split modalities, flatten, and encode with Swin # -------------------------------------------------------- # Each is (B*N, 1, H, W) def expand3(x): return x.expand(-1, 3, -1, -1) pre = expand3(volume[:, :, 0, :, :].contiguous().view(B * N, 1, H, W)) post = expand3(volume[:, :, 1, :, :].contiguous().view(B * N, 1, H, W)) sub = expand3(volume[:, :, 2, :, :].contiguous().view(B * N, 1, H, W)) feat_pre = self._encode_slices(pre).view(B, N, -1) # (B, N, C) feat_post = self._encode_slices(post).view(B, N, -1) feat_sub = self._encode_slices(sub).view(B, N, -1) # -------------------------------------------------------- # 2) Cross‑modal attention fusion (per slice) # -------------------------------------------------------- # Build token sequence [pre, post, sub] for each slice # Shape before fuser: (B, N, 3, C) -> we fuse along dim=2 slice_tokens = torch.stack([feat_pre, feat_post, feat_sub], dim=2) # Add modality embeddings slice_tokens = slice_tokens + self.mod_embed.transpose(0, 1) # (B, N, 3, C) slice_tokens = slice_tokens.view(B * N, 3, self.embed_dim) # (B*N, 3, C) # Transformer encoder attends across the 3 tokens fused = self.slice_fuser(slice_tokens)[:, 0] # take CLS‑like first token fused = fused.view(B, N, self.embed_dim) # (B, 32, C) # -------------------------------------------------------- # 3) ABMIL attention over 32 fused slices # -------------------------------------------------------- A = torch.tanh(self.attn_V(fused)) # (B, N, hidden) A = self.attn_U(A) # (B, N, 1) A = torch.softmax(A, dim=1) # (B, N, 1) patient_feat = (A * fused).sum(dim=1) # (B, C) # -------------------------------------------------------- # 4) Head # -------------------------------------------------------- logits = self.classifier(self.dropout(self.norm(patient_feat))) return logits, A.squeeze(-1) # (B, num_classes), (B, 32) class ABMIL_Swin(nn.Module): """ Attention‑based Multiple‑Instance Learning (ABMIL) model for 3‑channel breast MRI slice triplets (pre / post / sub). Input shape : (B, 32, 3, 224, 224) Output : logits (B, num_classes) + attention weights (B, 32) """ def __init__( self, model_name: str = "swin_tiny_patch4_window7_224.ms_in22k", pretrained: bool = True, num_classes: int = 2, hidden_dim: int = 256, dropout: float = 0.1, ): super().__init__() # 1) Swin backbone WITHOUT final linear head self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0) self.embed_dim = self.backbone.num_features # e.g. 768 (tiny), 1024 (base) # 2) Attention network (ABMIL) self.attn_V = nn.Linear(self.embed_dim, hidden_dim) self.attn_U = nn.Linear(hidden_dim, 1) # 3) Patient‑level classifier self.norm = nn.LayerNorm(self.embed_dim) self.dropout = nn.Dropout(dropout) self.classifier = nn.Linear(self.embed_dim, num_classes) # Optional init nn.init.trunc_normal_(self.attn_V.weight, std=0.02) nn.init.trunc_normal_(self.attn_U.weight, std=0.02) # ------------------------------------------------------------------ # def _slice_embed(self, x: torch.Tensor) -> torch.Tensor: """ Forward one or more images through Swin and return pooled features. Args: x : (N, 3, 224, 224) Returns: pooled : (N, C) """ feats = self.backbone.forward_features(x) # (N, H*W, C) pooled = self.backbone.forward_head(feats, pre_logits=True) # (N, C) return pooled # ------------------------------------------------------------------ # def forward(self, volume: torch.Tensor): """ Args: volume : (B, 32, 3, 224, 224) Returns: logits : (B, num_classes) attn_w : (B, 32) -- attention per slice (sums to 1) """ B, N, C, H, W = volume.shape x = volume.view(B * N, C, H, W) # flatten slices # (B*N, 3, 224, 224) → (B*N, embed_dim) → reshape slice_feats = self._slice_embed(x).view(B, N, -1) # (B, 32, embed_dim) # ---------------- ABMIL attention ---------------- # A = torch.tanh(self.attn_V(slice_feats)) # (B, 32, hidden) A = self.attn_U(A) # (B, 32, 1) A = torch.softmax(A, dim=1) # attention weights # Weighted sum → patient embedding patient_feat = (A * slice_feats).sum(dim=1) # (B, embed_dim) # ---------------- Head ---------------- # out = self.classifier(self.dropout(self.norm(patient_feat))) # (B, num_classes) return out, A.squeeze(-1) # logits, attention weights def main(args): read_excel = pd.read_csv(args.configuration_file, skip_blank_lines=True, na_values=['NaN']) max_f1_scores = [] model_files = [] for index, row in read_excel.iterrows(): if args.log_wandb: wandb.login() run = wandb.init(project=args.project_name, reinit=True, entity="adarshbhandary") with run: print('Experiment Index', index) torch.cuda.empty_cache() print("Reading Configuration File..") config = { "cv_run": row['no'], "task": row['task'], "model": row['model'], "method": row['method'], "view": row['view'], "pretrained": row['pretrained'], "loss": row['loss'], "optimizer": row['optimizer'], "use_sampler": row['use_sampler'], "height": int(row['height']), "width": int(row['width']), "background_crop": row['background_crop'], "use_clahe": row['use_clahe'], "batch_size": int(row['batch_size']), "multi_gpu": bool(row['multi_gpu']), "num_worker": row['num_worker'], "learning_rate": row['lr'], "weight_decay": row['weight_decay'], "drop_out": row['drop_out'], "num_epochs": int(row['epochs']), "patience": row['patience'], "probability": row['probability'], "aug_mix_p":row['aug_mix_p'], "erasing":row['erasing'], "attention_head": row["attention_head"], "rage":row["rage"], "get_gradcam": bool(row['get_gradcam']), "out_folder": row['out_folder'], "folder_name": row['folder_name'], "excel_file_train": row['train'], "excel_file_validation": row['valid'], "excel_file_test": row['test'], "data_folder": row['data_folder'], "pretrained_weights": row["pretrained_weights"], "results_path": row['csv_results'] } set_seed(int(config["cv_run"])) wandb.config.update(config, allow_val_change=True) os.makedirs(os.path.join(config["out_folder"], str(config["folder_name"])), exist_ok=True) save_out_folder = config["out_folder"] + str(config["folder_name"]) ## config method is MRI target_names = extract_target_names(method = config["method"]) num_classes = len(target_names) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print('Device : ', device) weight, samples_weight = compute_sample_weights(config["excel_file_train"], class_type = config["method"], view = config["view"]) if config["use_sampler"]: sampler = WeightedRandomSampler(samples_weight, num_samples=len(samples_weight), replacement=True) shuffle = False else: sampler = None shuffle = True if config["model"] == "vit_small": train_transforms = TrainTransform(aug=True, aug_mix_p = config["aug_mix_p"]) else: train_transforms = TrainTransformBaseline(aug=True, aug_mix_p = config["aug_mix_p"], erasing = config["erasing"]) valid_transforms = BasicTransforms() if config["method"] == "MRI_MIL": data_train = BreastMRI_ABMIL_NII( split='train' ) data_valid = BreastMRI_ABMIL_NII( split='val' ) data_test = BreastMRI_ABMIL_NII( split='test' ) else: data_train = BreastMRI(csv_file = config["excel_file_train"], root_dir=None, split='train', transform = valid_transforms) data_valid = BreastMRI(csv_file = config["excel_file_train"], root_dir=None, split='val', transform = valid_transforms) data_test = BreastMRI(csv_file = config["excel_file_train"], root_dir=None, split='test', transform = valid_transforms) print("Train dataset size:", len(data_train)) print("Val dataset size:", len(data_valid)) print("Test dataset size:", len(data_test)) dataloader = { 'train': torch.utils.data.DataLoader(data_train, batch_size=config["batch_size"], shuffle=shuffle, sampler=sampler, num_workers=config["num_worker"], pin_memory=True, prefetch_factor=2, drop_last=True ), 'valid': torch.utils.data.DataLoader(data_valid, batch_size=1, shuffle=False, sampler=None, num_workers=config["num_worker"], pin_memory=True, prefetch_factor=2, drop_last=False ), 'test': torch.utils.data.DataLoader(data_test, batch_size=1, shuffle=False, sampler=None, num_workers=config["num_worker"], pin_memory=True, prefetch_factor=2, drop_last=False ) } if config["method"]=="MRI": if config["model"] == "swin_cross": model = MultiViewSwinCrossAttn(num_classes=num_classes) else: model = MultiViewSwinModel(num_classes=num_classes) elif config["method"] == "MRI_MIL": if config["model"] == "swin_cross": model = CrossModalAttentionABMIL_Swin(num_classes=num_classes) else: model = ABMIL_Swin(num_classes=num_classes) model = model.to(device) #summary(model, (config["batch_size"], 3, config["height"], config["width"])) if config["loss"] == "F": criterion = FocalLoss(gamma=config["erasing"], reduction='mean').to(device) elif config["loss"] == "CE": criterion = nn.CrossEntropyLoss().to(device) optimizer_ft = optim.Adam(model.parameters(), lr=config["learning_rate"], weight_decay=config["weight_decay"]) if config["method"] == 'MRI': train = FullySupervisedMultiViewTrainer(device, config["method"], model, criterion, optimizer_ft, dataloader, target_names, save_out_folder, config["num_epochs"], config["patience"]) train.main_loop() model.load_state_dict(torch.load(save_out_folder + 'max_metrics_epoch.pth')) train = FullySupervisedMultiViewTrainer(device, config["method"], model, criterion, optimizer_ft, dataloader, target_names, save_out_folder, config["num_epochs"], config["patience"]) test_accuracy, test_f1, true, pred, test_auc, averaged_results = train.test_loop() elif config["method"] == 'MRI_MIL': #model.load_state_dict(torch.load(save_out_folder + 'max_metrics_epoch.pth')) train = MultiInstanceTrainer(device, config["method"], model, criterion, config["loss"], optimizer_ft, dataloader, target_names, save_out_folder, config["num_epochs"], config["patience"]) train.main_loop() model.load_state_dict(torch.load(save_out_folder + 'max_metrics_epoch.pth')) train = MultiInstanceTrainer(device, config["method"], model, criterion, config["loss"], optimizer_ft, dataloader, target_names, save_out_folder, config["num_epochs"], config["patience"]) test_accuracy, test_f1, true, pred, test_auc, averaged_results = train.test_loop() print('Test Results') print(classification_report(true, pred, target_names=target_names, digits=4)) display_confusion_matrix(true, pred, target_names, save_out_folder+'confusion_matrix.png', use_tta=False, label=config["task"], normalize=False) config["model location"] = save_out_folder + 'least_validation_loss.pth' config["test accuracy"] = test_accuracy config["test f1"] = test_f1 config["test auc"] = test_auc #log_scalar_values(config["results_path"], config) max_f1_scores.append(test_f1) model_files.append(save_out_folder + 'least_validation_loss.pth') wandb.log({ "test_accuracy" : test_accuracy, "test_f1" : test_f1, "test_auc" : test_auc, "test_metrics" : averaged_results, }) #log_scalar_values(config["results_path"], config) wandb.config.update(config, allow_val_change=True) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--configuration_file", type=Path, default='/cluster/panambur/DINO/config/transfer_baseline.csv', help="""Change hyperparameters and input the filepath of CSV file""") parser.add_argument("--log_wandb", type=bool, default=True, help="""Use FALSE for Siemens Code.""") parser.add_argument("--project_name", type=str, default='BreastMRI Tiny', help="""For Wandb logging""") args = parser.parse_args() main(args)