test.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import pandas as pd
import numpy as np
import argparse
from collections import defaultdict
from sklearn.preprocessing import label_binarize
from sklearn.metrics import f1_score, accuracy_score, roc_auc_score, roc_curve
import cv2
from src.data_loader.mammo_transforms import BasicTransforms
import os
import torch
import wandb
from dataset import MammographyDataset, BreastMRI
#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
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 sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt

class BreastMRI(torch.utils.data.Dataset):
    def __init__(self, csv_file, root_dir=None, split=None, transform=None):
        self.data = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def _load_image(self, img_path):
        image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        image = cv2.resize(image, (224, 224))
        image = np.stack([image] * 3, axis=-1)
        if self.transform:
            image = self.transform(image)
        else:
            image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
        return image

    def __getitem__(self, idx):
        row = self.data.iloc[idx]
        pre_img = self._load_image(row['pre_png'])
        post_img = self._load_image(row['post_png'])
        sub_img = self._load_image(row['sub_png'])
        label = torch.tensor(row['Lesion']).long()
        patient_id = row['Patient_ID']
        return {
            'pre': pre_img,
            'post': post_img,
            'sub': sub_img,
            'label': label,
            'patient_id': patient_id
        }


class MultiViewSwinModel(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



def has_valid_lesion(pred_classes, target_class):
        # Check for first or last slice being lesion
        if pred_classes[0] == target_class or pred_classes[-1] == target_class:
            return True
        # Check for 2 or more consecutive slices
        for i in range(len(pred_classes) - 1):
            if pred_classes[i] == target_class and pred_classes[i + 1] == target_class :
                return True
        return False



@torch.no_grad()
def evaluate(model, dataloader, device):
    read_ann = pd.read_csv("/anvme/workspace/b268dc11-breastt/BreastMRI/DataExtraction/annotated_odelia_slice_level.csv")
    read_ann = read_ann[read_ann["Split"].str.lower() == "test"]

    model.eval()
    patient_probs = defaultdict(list)
    patient_labels = {}

    for batch in dataloader:
        pre = batch['pre'].to(device)
        post = batch['post'].to(device)
        sub = batch['sub'].to(device)
        labels = batch['label']
        patient_ids = batch['patient_id']

        logits = model(pre, post, sub)
        probs = F.softmax(logits, dim=1).cpu().numpy()

        for i, pid in enumerate(patient_ids):
            patient_probs[pid].append(probs[i])
            patient_labels[pid] = labels[i].item()

    y_true = []
    y_pred = []
    y_true_onehot = []
    y_pred_onehot = []

    for pid in patient_probs:
        gt_label = patient_labels[pid]
        if pid in read_ann["UID"].values:
            rows = read_ann[read_ann["UID"] == pid]
            print("Tri-Slice Number(s):", rows["lesion_slice"].dropna().tolist())
            print("")
        else:
            print(f"{pid} not found in annotation.")
        slice_probs = np.array(patient_probs[pid])  # shape: (num_slices, num_classes)
        print("Patient ID:", pid, "GT label:", gt_label, "Slice probs shape:", np.round(slice_probs, 2))
        print("")
        predicted_classes = np.argmax(slice_probs, axis=1)
        print("")
        print("Predicted Classes:", predicted_classes)
        print("")
        # Check lesion logic
        has_malignant = has_valid_lesion(predicted_classes, target_class=2)
        has_benign = has_valid_lesion(predicted_classes, target_class=1)

        if has_malignant:
            idx = np.argmax(slice_probs[:, 2])  # slice with highest malignant prob
        elif has_benign:
            idx = np.argmax(slice_probs[:, 1])  # slice with highest benign prob
        else:
            idx = np.argmax(slice_probs[:, 0])  # slice with highest normal prob

        final_prob = slice_probs[idx]
        final_pred = np.argmax(final_prob)

        print(f"Evaluating patient: {pid} | GT label: {gt_label} | Pred: {final_pred} | MaxProbSlice: {idx}")
        print("")
        print("")
        print("")
        y_true.append(gt_label)
        y_pred.append(final_pred)
        y_true_onehot.append(label_binarize([gt_label], classes=[0, 1, 2])[0])
        y_pred_onehot.append(final_prob)

    # Convert to arrays
    y_true = np.array(y_true)
    y_pred = np.array(y_pred)
    y_true_onehot = np.vstack(y_true_onehot)
    y_pred_onehot = np.vstack(y_pred_onehot)

    # Confusion Matrix
    cm = confusion_matrix(y_true, y_pred)
    class_names = ['Normal', 'Benign', 'Malignant']
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=class_names)
    disp.plot(cmap=plt.cm.Blues)
    plt.title("Confusion Matrix (Patient-Level)")
    plt.tight_layout()
    plt.savefig("confusion_matrix_patient_level.png")
    plt.close()

    # Metrics (only class labels used)
    accuracy = accuracy_score(y_true, y_pred)
    f1 = f1_score(y_true, y_pred, average="macro")
    auc = roc_auc_score(y_true_onehot, y_pred_onehot, average="micro")

    # ROC from binary class outputs
    fpr, tpr, _ = roc_curve(y_true_onehot.ravel(), y_pred_onehot.ravel())
    sensitivity_at_90_spec = np.interp(0.10, fpr, tpr)
    specificity_at_90_sens = 1 - np.interp(0.90, tpr, fpr)
    avg_metric = np.mean([auc, sensitivity_at_90_spec, specificity_at_90_sens])

    # Results
    print("\n--- Evaluation Results (Per-Patient, Rule-Based Only) ---")
    print(f"Accuracy                : {accuracy:.4f}")
    print(f"F1 Score (Macro)       : {f1:.4f}")
    print(f"AUC (Micro)            : {auc:.4f}")
    print(f"Sens @ 90% Specificity : {sensitivity_at_90_spec:.4f}")
    print(f"Spec @ 90% Sensitivity : {specificity_at_90_sens:.4f}")
    print(f"Average Metric         : {avg_metric:.4f}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--csv_path', type=str, required=True)
    parser.add_argument('--model_path', type=str, required=True)
    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--model_type', type=str, default='swin_tiny_patch4_window7_224.ms_in22k')
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    dataset = BreastMRI(csv_file=args.csv_path, split="test")
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=4)

    model = MultiViewSwinModel(model_name=args.model_type)
    model.load_state_dict(torch.load(args.model_path, map_location=device))
    model.to(device)

    evaluate(model, dataloader, device)