evaluate_brats.py

import os
import logging
import argparse
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Subset, random_split
from torchvision import transforms
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from utils import (
    generalized_energy_distance_iou, hm_iou_cal, dice_max_cal2, dice_avg_cal
)
from brats_dataset import BratsDataset
from asam import ASAM
from datetime import datetime
import random
import numpy as np

def set_seed(seed):
    """Set all random seeds for reproducibility"""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    os.environ['PYTHONHASHSEED'] = str(seed)
    print(f"Random seed set to {seed} for reproducibility")

def visualize_test_results(image, targets, predictions_16, save_path='test_visualization.png', threshold=0.5):
    """
    Visualize test results: original image, 3 GT labels, and dynamic number of predictions
    Args:
        image: torch.Size([1, 128, 128]) - original image
        targets: torch.Size([3, 128, 128]) - 3 GT labels
        predictions_16: torch.Size([N, 128, 128]) - N prediction results
        save_path: path to save the image
        threshold: binarization threshold
    """
    try:
        import matplotlib.pyplot as plt
        
        image = image.cpu().detach().numpy().squeeze()
        targets = targets.cpu().detach().numpy()
        predictions_16 = predictions_16.cpu().detach().numpy()
        
        num_predictions = predictions_16.shape[0]
        
        if len(image.shape) == 3 and image.shape[0] == 3:
            image = np.mean(image, axis=0)
        elif len(image.shape) == 3 and image.shape[0] == 1:
            image = image.squeeze(0)
        
        num_rows = 1 + (num_predictions + 3) // 4
        num_cols = 5
        
        fig, axes = plt.subplots(num_rows, num_cols, figsize=(20, 4*num_rows))
        fig.suptitle(f'Test Results: Original + 3 GT + {num_predictions} Predictions (threshold={threshold})', fontsize=16)
        
        if num_rows == 1:
            axes = axes.reshape(1, -1)
        
        axes[0, 0].imshow(image, cmap='gray')
        axes[0, 0].set_title('Original Image')
        axes[0, 0].axis('off')
        
        for i in range(3):
            gt_data = targets[i]
            if len(gt_data.shape) == 3 and gt_data.shape[0] == 3:
                gt_data = np.mean(gt_data, axis=0)
            elif len(gt_data.shape) == 3 and gt_data.shape[0] == 1:
                gt_data = gt_data.squeeze(0)
                
            axes[0, i+1].imshow(gt_data, cmap='gray', vmin=0, vmax=1)
            axes[0, i+1].set_title(f'GT {i+1}')
            axes[0, i+1].axis('off')
            
        axes[0, 4].axis('off')
            
        for i in range(num_predictions):
            row = (i // 4) + 1
            col = i % 4
            
            if row >= num_rows:
                break
                
            pred_data = predictions_16[i]
            if len(pred_data.shape) == 3 and pred_data.shape[0] == 3:
                pred_data = np.mean(pred_data, axis=0)
            elif len(pred_data.shape) == 3 and pred_data.shape[0] == 1:
                pred_data = pred_data.squeeze(0)
            
            pred_prob = 1 / (1 + np.exp(-pred_data))
            pred_binary = (pred_prob > 0.5).astype(np.float32)
            
            axes[row, col].imshow(pred_binary, cmap='gray', vmin=0, vmax=1)
            axes[row, col].set_title(f'Pred {i+1}')
            axes[row, col].axis('off')
        
        for row in range(num_rows):
            for col in range(num_cols):
                if row == 0 and col < 4:
                    continue
                if row > 0 and col < 4 and (row-1)*4 + col < num_predictions:
                    continue
                axes[row, col].axis('off')
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"Test visualization results saved to: {save_path}")
    except Exception as e:
        print(f"Test visualization error: {e}")

def initialize_cross_attention_modules(checkpoint_dir):
    cross_attention_modules = []
    for epoch in [1,2,3,4]:  
        epoch_path = os.path.join(checkpoint_dir, f'epoch_{epoch}_weights.pth')
        if os.path.exists(epoch_path):
            try:
                epoch_weights = torch.load(epoch_path, map_location=device)
                ca = SpatialAttention(in_channels=16, hidden_dim=128).to(device)
                ca.load_state_dict(epoch_weights['attention_state'])
                cross_attention_modules.append(ca)
                
                print(f"Successfully loaded weights for epoch {epoch} from {checkpoint_dir}")
            except Exception as e:
                print(f"Error loading weights for epoch {epoch}: {str(e)}")
        else:
            print(f"Warning: Weights file for epoch {epoch} not found at {epoch_path}")

    return cross_attention_modules

def generate_logits_high_res_list(image_batch, image_batch_oc, box1024_batch, boxshift_batch, label_batch, device, networks):
    logits_high_res_list = []
    for net, weights in networks:
        for j in range(4):
            outputs = net.forward(image_batch, image_batch_oc, box1024_batch, boxshift_batch, label_batch, device, train=False)
            logits_high = outputs['masks'].to(device) * weights.unsqueeze(-1)
            logits_high_res = logits_high.sum(1)
            logits_high_res_list.append(logits_high_res)
    return logits_high_res_list

def load_specific_epoch_weights(checkpoint_dir, epoch_num, device):
    """Load weights from specific epoch"""
    weight_path = os.path.join(checkpoint_dir, f'weights_epoch_{epoch_num}.pth')
    weight = torch.load(weight_path, map_location=device)
    pretrained_state_dict = weight['model_state_dict']
    mask_weights = weight['mask_weights']
    print(f"Loading weights from epoch {epoch_num} from {checkpoint_dir}")
    return pretrained_state_dict, mask_weights

def freeze_all(model):
    """Freeze all parameters"""
    print("Freezing all parameters...")
    
    total_params = 0
    frozen_params = 0
    
    for name, param in model.named_parameters():
        total_params += param.numel()
        param.requires_grad = False
        frozen_params += param.numel()
    
    print(f"Total parameters: {total_params:,}")
    print(f"Frozen parameters: {frozen_params:,}")
    
    return total_params, frozen_params
set_seed(1234)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

class SpatialAttention(nn.Module):
    """Transformer-based spatial attention module using Multi-Head Self-Attention"""
    def __init__(self, in_channels=16, hidden_dim=128, num_heads=8, num_output=3):
        super().__init__()
        self.in_channels = in_channels
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.num_output = num_output
        self.head_dim = hidden_dim // num_heads
        
        assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads"
        
        self.input_projection = nn.Conv2d(in_channels, hidden_dim, 1)
        self.position_encoding = nn.Parameter(torch.randn(1, hidden_dim, 128, 128))
        
        self.query_projection = nn.Conv2d(hidden_dim, hidden_dim, 1)
        self.key_projection = nn.Conv2d(hidden_dim, hidden_dim, 1)
        self.value_projection = nn.Conv2d(hidden_dim, hidden_dim, 1)
        
        self.attention_output = nn.Conv2d(hidden_dim, hidden_dim, 1)
        
        self.layer_norm1 = nn.LayerNorm(hidden_dim)
        self.layer_norm2 = nn.LayerNorm(hidden_dim)
        
        self.ffn = nn.Sequential(
            nn.Conv2d(hidden_dim, hidden_dim * 4, 1),
            nn.GELU(),
            nn.Conv2d(hidden_dim * 4, hidden_dim, 1),
        )
        
        self.output_projection = nn.Sequential(
            nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
            nn.GELU(),
            nn.Conv2d(hidden_dim, hidden_dim // 2, 3, padding=1),
            nn.GELU(),
            nn.Conv2d(hidden_dim // 2, num_output, 1)
        )
        
        self.dropout = nn.Dropout(0.1)
        self._init_weights()
        
    def _init_weights(self):
        """Initialize weights"""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
    
    def multi_head_attention(self, x):
        """
        Multi-Head Self-Attention with memory-efficient implementation
        Args:
            x: (B, C, H, W)
        Returns:
            output: (B, C, H, W)
        """
        B, C, H, W = x.shape
        
        if H * W > 8192:
            return self._chunked_attention(x)
        
        q = self.query_projection(x)
        k = self.key_projection(x)
        v = self.value_projection(x)
        
        q = q.view(B, self.num_heads, self.head_dim, H * W)
        k = k.view(B, self.num_heads, self.head_dim, H * W)
        v = v.view(B, self.num_heads, self.head_dim, H * W)
        
        attention_scores = torch.matmul(q.transpose(-2, -1), k) / (self.head_dim ** 0.5)
        attention_weights = torch.softmax(attention_scores, dim=-1)
        attention_weights = self.dropout(attention_weights)
        
        attended_values = torch.matmul(v, attention_weights.transpose(-2, -1))
        attended_values = attended_values.view(B, C, H, W)
        
        output = self.attention_output(attended_values)
        
        return output
    
    def _chunked_attention(self, x, chunk_size=64):
        """
        Chunked attention computation to save memory
        Args:
            x: (B, C, H, W)
            chunk_size: Size of each chunk
        Returns:
            output: (B, C, H, W)
        """
        B, C, H, W = x.shape
        
        q = self.query_projection(x)
        k = self.key_projection(x)
        v = self.value_projection(x)
        
        q = q.view(B, self.num_heads, self.head_dim, H * W)
        k = k.view(B, self.num_heads, self.head_dim, H * W)
        v = v.view(B, self.num_heads, self.head_dim, H * W)
        
        seq_len = H * W
        attended_values = torch.zeros_like(v)
        
        for i in range(0, seq_len, chunk_size):
            end_i = min(i + chunk_size, seq_len)
            q_chunk = q[:, :, :, i:end_i]
            
            attention_scores = torch.matmul(q_chunk.transpose(-2, -1), k) / (self.head_dim ** 0.5)
            attention_weights = torch.softmax(attention_scores, dim=-1)
            attention_weights = self.dropout(attention_weights)
            
            attended_chunk = torch.matmul(v, attention_weights.transpose(-2, -1))
            attended_values[:, :, :, i:end_i] = attended_chunk
        
        attended_values = attended_values.view(B, C, H, W)
        output = self.attention_output(attended_values)
        
        return output
    
    def forward(self, logits_list):
        """
        Args:
            logits_list: list of (B, H, W) tensors, length=16
        Returns:
            output: (B, 4, H, W)
        """
        logits_stacked = torch.stack(logits_list, dim=1)
        B, N, H, W = logits_stacked.shape
        
        x = self.input_projection(logits_stacked)
        
        pos_encoding = self.position_encoding[:, :, :H, :W]
        x = x + pos_encoding
        
        residual = x
        x_flat = x.permute(0, 2, 3, 1).contiguous().view(B * H * W, -1)
        x_flat = self.layer_norm1(x_flat)
        x = x_flat.view(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
        
        attention_output = self.multi_head_attention(x)
        x = residual + self.dropout(attention_output)
        
        residual = x
        x_flat = x.permute(0, 2, 3, 1).contiguous().view(B * H * W, -1)
        x_flat = self.layer_norm2(x_flat)
        x = x_flat.view(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
        
        ffn_output = self.ffn(x)
        x = residual + self.dropout(ffn_output)
        
        output = self.output_projection(x)
        
        return output

logging.basicConfig(filename=f'test_log_{timestamp}.txt', level=logging.INFO, format='%(asctime)s - %(message)s')

parser = argparse.ArgumentParser(description='Evaluate the model with specified epochs and weights.')
parser.add_argument('--gpuid', type=int, default=0, help='ID of the GPU to use.(2,3,4,5,6)')
parser.add_argument('--batch_size', type=int, default=40, help='Batch size for data loading.')
parser.add_argument('--total_samples', type=int, default=16, help='Total number of samples to generate.')
args = parser.parse_args()

device = torch.device(f'cuda:{args.gpuid}' if torch.cuda.is_available() else 'cpu')

target_epochs = [450, 460, 470, 480]
checkpoint_dir = "/path/to/Brats/asam_ckpoints"
networks = []
for i, epoch_num in enumerate(target_epochs):
    pretrained_state_dict_i, mask_weights_i = load_specific_epoch_weights(checkpoint_dir, epoch_num, device)
    
    net = ASAM(dataset='brats').to(device)
    net.load_state_dict(pretrained_state_dict_i)  
    _, _ = freeze_all(net)
    networks.append((net,mask_weights_i.to(device)))
    
    print(f"Initialized network {i+1} with epoch {epoch_num} weights")

print(f"Initialized {len(networks)} networks with different epoch weights")

vis_dir = f'test_vis/test_vis_{timestamp}'
os.makedirs(vis_dir, exist_ok=True)



checkpoint_dir = '/path/to/checkpoints'
cross_attention_modules = initialize_cross_attention_modules(checkpoint_dir)

csv_file = '/path/to/image_mask_paths1.csv'
dataset = BratsDataset(csv_file=csv_file)
dataset_size = len(dataset)
train_size = int(dataset_size * 0.6)
val_size = int(dataset_size * 0.2)
test_size = dataset_size - (train_size + val_size)
train_dataset, val_dataset, test_dataset = random_split(dataset, [train_size, val_size, test_size])
print("Train Dataset Length:", len(train_dataset))
print("Test Dataset Length:", len(test_dataset))
test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)

samples_per_net = args.total_samples // len(cross_attention_modules) // 4

ged_score = dice_max2_score = hm_iou_score = dmean_score = 0
total_test_samples = 0
with torch.no_grad():
    for i_batch, sampled_batch in enumerate(test_loader):
        print(f'Processing batch {i_batch}')
        logging.info(f'Processing batch {i_batch}')
        image_batch, label_batch = sampled_batch['image'].to(device), sampled_batch['label'].to(device)
        label_four_batch = sampled_batch['label_four']
        image_batch_oc = image_batch
        box1024_batch = sampled_batch['box_1024'].to(device) 
        boxshift_batch = sampled_batch['box_shift'].to(device)
        pred_list = [[] for _ in range(image_batch.shape[0])]
        
        for cross_attention_module_i in cross_attention_modules:
            for _ in range(samples_per_net):
                logits_high_res_list = generate_logits_high_res_list(image_batch, image_batch_oc, box1024_batch, boxshift_batch, label_batch, device, networks)
                
                enhanced_logits_high_res = cross_attention_module_i(logits_high_res_list)
                for i in range(3):
                    enhanced_logits_high_res_i = enhanced_logits_high_res[:, i:i+1]
                    for j in range(image_batch.shape[0]):
                        pred_list[j].append(enhanced_logits_high_res_i[j])
        
        for index in range(len(pred_list)):
            pred_eval = torch.cat(pred_list[index], 0)
            pred_eval = (pred_eval > 0).cpu().detach().int()

            iou_score_iter, ged_score_iter = generalized_energy_distance_iou(pred_eval, label_four_batch[index])
            score = hm_iou_cal(pred_eval, label_four_batch[index])
            hm_iou_score += score
            dice_max2_score += dice_max_cal2(pred_eval, label_four_batch[index])
            ged_score += ged_score_iter
            dmean_score += dice_avg_cal(pred_list[index], label_four_batch[index])

            total_test_samples += 1
                        
            if total_test_samples <= 10:
                vis_save_path = f'{vis_dir}/test_batch_{i_batch}_sample_{total_test_samples}.png'
                visualize_test_results(
                    image_batch_oc[index], 
                    label_four_batch[index], 
                    pred_eval, 
                    vis_save_path, 
                    threshold=0.5
                )
        
# Calculate average scores
ged = ged_score / test_size
dice_max2 = dice_max2_score / test_size
hm_iou = hm_iou_score / test_size
dmean = dmean_score / test_size

print(f"ged_score: {ged}, dice_max_score2: {dice_max2}, hm_iou_score: {hm_iou}, dmean_score: {dmean}")
logging.info(f"ged_score: {ged}, dice_max_score2: {dice_max2}, hm_iou_score: {hm_iou}, dmean_score: {dmean}")