evaluate_lidc.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
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 lidc_dataset import LIDC_IDRI, RandomGenerator
from asam import ASAM
from datetime import datetime
import random
import numpy as np

# Set random seed for reproducibility
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, 4 GT labels, 16 predictions
    Args:
        image: torch.Size([1, 128, 128]) - original image
        targets: torch.Size([4, 128, 128]) - 4 GT labels  
        predictions_16: torch.Size([16, 128, 128]) - 16 prediction results
        save_path: path to save the image
        threshold: binarization threshold
    """
    try:
        import matplotlib.pyplot as plt
        
        # Convert tensors to numpy and move to CPU
        image = image.cpu().detach().numpy().squeeze()
        targets = targets.cpu().detach().numpy()
        predictions_16 = predictions_16.cpu().detach().numpy()
        
        # Create 5x5 subplot layout
        fig, axes = plt.subplots(5, 5, figsize=(20, 20))
        fig.suptitle(f'Test Results: Original + 4 GT + 16 Predictions (threshold={threshold})', fontsize=16)
        
        # Row 1: original image + 4 GT labels
        # Show original image
        axes[0, 0].imshow(image, cmap='gray')
        axes[0, 0].set_title('Original Image')
        axes[0, 0].axis('off')
        
        # Show 4 GT labels
        for i in range(4):
            axes[0, i+1].imshow(targets[i], cmap='gray', vmin=0, vmax=1)
            axes[0, i+1].set_title(f'GT {i+1}')
            axes[0, i+1].axis('off')
            
        # Rows 2-5: 16 prediction results
        for i in range(16):
            row = (i // 4) + 1
            col = i % 4
            
            # Apply sigmoid activation and binarization (threshold 0.5)
            pred_prob = 1 / (1 + np.exp(-predictions_16[i]))
            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')
            
        # Hide unused subplot
        axes[row, 4].axis('off')
        
        # Adjust subplot spacing
        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}")

# Apply random seed
set_seed(2025)
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=4):
        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"
        
        # Input projection: map 16 channels to hidden_dim
        self.input_projection = nn.Conv2d(in_channels, hidden_dim, 1)
        
        # Position encoding: provide positional information for each spatial location
        self.position_encoding = nn.Parameter(torch.randn(1, hidden_dim, 128, 128))
        
        # Multi-Head Self-Attention
        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)
        
        # Attention output projection
        self.attention_output = nn.Conv2d(hidden_dim, hidden_dim, 1)
        
        # LayerNorm for transformer
        self.layer_norm1 = nn.LayerNorm(hidden_dim)
        self.layer_norm2 = nn.LayerNorm(hidden_dim)
        
        # Feed Forward Network
        self.ffn = nn.Sequential(
            nn.Conv2d(hidden_dim, hidden_dim * 4, 1),
            nn.GELU(),
            nn.Conv2d(hidden_dim * 4, hidden_dim, 1),
        )
        
        # Output projection: generate 4 output channels
        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)
        )
        
        # Dropout
        self.dropout = nn.Dropout(0.1)
        
        # Initialize weights
        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
        
        # Use chunked computation for high resolution to avoid OOM
        if H * W > 8192:
            return self._chunked_attention(x)
        
        # Generate Q, K, V
        q = self.query_projection(x)
        k = self.key_projection(x)
        v = self.value_projection(x)
        
        # Reshape to multi-head format: (B, num_heads, head_dim, H*W)
        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)
        
        # Compute attention scores: (B, num_heads, H*W, 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)
        
        # Apply attention weights: (B, num_heads, head_dim, H*W)
        attended_values = torch.matmul(v, attention_weights.transpose(-2, -1))
        
        # Reshape back: (B, C, H, W)
        attended_values = attended_values.view(B, C, H, W)
        
        # Output projection
        output = self.attention_output(attended_values)
        
        return output
    
    def _chunked_attention(self, x, chunk_size=64):
        """
        Chunked attention 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
        
        # Generate Q, K, V
        q = self.query_projection(x)
        k = self.key_projection(x)
        v = self.value_projection(x)
        
        # Reshape to sequence format: (B, num_heads, head_dim, H*W)
        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)
        
        # Chunked attention computation
        for i in range(0, seq_len, chunk_size):
            end_i = min(i + chunk_size, seq_len)
            q_chunk = q[:, :, :, i:end_i]
            
            # Compute attention scores between current chunk and all positions
            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)
            
            # Apply attention weights
            attended_chunk = torch.matmul(v, attention_weights.transpose(-2, -1))
            attended_values[:, :, :, i:end_i] = attended_chunk
        
        # Reshape back to original shape: (B, C, H, W)
        attended_values = attended_values.view(B, C, H, W)
        
        # Output projection
        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)
        """
        # Stack 16 predictions
        logits_stacked = torch.stack(logits_list, dim=1)  # (B, 16, H, W)
        B, N, H, W = logits_stacked.shape
        
        # Input projection
        x = self.input_projection(logits_stacked)  # (B, hidden_dim, H, W)
        
        # Add positional encoding
        pos_encoding = self.position_encoding[:, :, :H, :W]
        x = x + pos_encoding
        
        # Transformer Block 1: Multi-Head Self-Attention + Residual
        residual = x
        x_flat = x.permute(0, 2, 3, 1).contiguous().view(B * H * W, -1)  # (B*H*W, C)
        x_flat = self.layer_norm1(x_flat)
        x = x_flat.view(B, H, W, -1).permute(0, 3, 1, 2).contiguous()  # (B, C, H, W)
        
        attention_output = self.multi_head_attention(x)
        x = residual + self.dropout(attention_output)
        
        # Transformer Block 2: Feed Forward Network + Residual
        residual = x
        x_flat = x.permute(0, 2, 3, 1).contiguous().view(B * H * W, -1)  # (B*H*W, C)
        x_flat = self.layer_norm2(x_flat)
        x = x_flat.view(B, H, W, -1).permute(0, 3, 1, 2).contiguous()  # (B, C, H, W)
        
        ffn_output = self.ffn(x)
        x = residual + self.dropout(ffn_output)
        
        # Output projection: generate 4 output channels
        output = self.output_projection(x)  # (B, 4, H, W)
        
        return output


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

# Parse command-line arguments
parser = argparse.ArgumentParser(description='Evaluate the model with specified epochs and weights.')
parser.add_argument('--sam_combined_weights_path', type=str, default='/path/to/final_weights.pth', help='Path to the combined weights file.')
parser.add_argument('--gpuid', type=int, default=3, 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()

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

# Load combined weights
sam_combined_weights = torch.load(args.sam_combined_weights_path, map_location=device)
# cross_attention_combined_weights = torch.load(args.cross_attention_combined_weights_path, map_location=device)

# Initialize SAM networks
def initialize_networks():
    networks = []
    for epoch in [10, 20, 30, 50]:
        epoch_key = f'epoch_{epoch}'
        if epoch_key in sam_combined_weights:
            net = ASAM().to(device)
            net.load_state_dict(sam_combined_weights[epoch_key]['model_state_dict'])
            networks.append((net, sam_combined_weights[epoch_key]['mask_weights'].to(device)))
        else:
            print(f"Warning: Weights for epoch {epoch} not found.")
    return networks

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):
    # Use 4 different networks to generate 16 logits_high_res (4 per network)
    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

# Create test visualization directory
vis_dir = f'test_vis/test_vis_{timestamp}'
os.makedirs(vis_dir, exist_ok=True)

# Initialize scores
ged_score = dice_max2_score = hm_iou_score = dmean_score = 0

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


# Prepare dataset
db = LIDC_IDRI(dataset_location='/path/to/data/', transform=transforms.Compose([
    RandomGenerator(output_size=[128, 128])
]))
dataset_size = len(db)
indices = list(range(dataset_size))
train_split = int(np.floor(0.6 * dataset_size))
validation_split = int(np.floor(0.8 * dataset_size))
test_indices = indices[validation_split:]
test_dataset = Subset(db, test_indices)

test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)

print(f"Total dataset size: {dataset_size}")
print(f"Test set size: {len(test_indices)}")

# Hyperparameter: number of samples per network
samples_per_net = args.total_samples // len(cross_attention_modules) // 4

# Evaluate the model
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 = sampled_batch['image_oc'].to(device)
        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):
                # Use 4 SAM networks to generate 16 logits_high_res (4 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) # (B, 4, H, W)
                # print(enhanced_logits_high_res.shape) # (B, 4, H, W)
                for i in range(4):
                    enhanced_logits_high_res_i = enhanced_logits_high_res[:, i:i+1]
                    # print(enhanced_logits_high_res_i.shape) # (B, 1, H, W)
                    for j in range(image_batch.shape[0]):
                        # print(enhanced_logits_high_res_i[j].shape) # (1, H, W)
                        pred_list[j].append(enhanced_logits_high_res_i[j])
                        # print(pred_list[j].shape)
        
        # print(len(pred_list))
        for index in range(len(pred_list)):  # batch dimension
            # print(len(pred_list[index]))
            pred_eval = torch.cat(pred_list[index], 0)
            # print(pred_eval.shape) # (16, H, W)
            pred_eval = (pred_eval > 0).cpu().detach().int()
            # print(pred_eval.shape)

            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
            # Pass original pred_list[index] instead of processed pred_eval
            dmean_score += dice_avg_cal(pred_list[index], label_four_batch[index])

            total_test_samples += 1
                        
            # Visualize first 10 samples
            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 / len(test_indices)
dice_max2 = dice_max2_score / len(test_indices)
hm_iou = hm_iou_score / len(test_indices)
dmean = dmean_score / len(test_indices)

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}")