evaluate.py

# simplified_evaluate.py
import os
import argparse
import numpy as np
import torch
import matplotlib.pyplot as plt
from tqdm import tqdm
from torch.serialization import safe_globals
from argparse import Namespace

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', type=str, default='./protein_fragment_dataset')
    parser.add_argument('--output_dir', type=str, default='./evaluation')
    parser.add_argument('--checkpoint', type=str, default='./output_small/model_best.pth')
    parser.add_argument('--num_visualizations', type=int, default=20)
    return parser.parse_args()

def visualize_samples(model, dataloader, device, output_dir, num_examples=10):
    """Visualize original, masked, and reconstructed distance maps"""
    os.makedirs(os.path.join(output_dir, 'visualizations'), exist_ok=True)
    
    model.eval()
    
    # Get a batch
    batch = next(iter(dataloader))
    distance_maps = batch['distance_map'].to(device)
    
    # Only use a subset for visualization
    distance_maps = distance_maps[:num_examples]
    
    # Forward pass
    with torch.no_grad():
        reconstructions, masks = model(distance_maps, mask_ratio=0.75)
    
    # Create visualizations
    for i in range(min(num_examples, distance_maps.size(0))):
        # Get maps and convert to numpy
        orig_map = distance_maps[i, 0].cpu().numpy()
        recon_map = reconstructions[i, 0].cpu().numpy()
        mask = masks[i].cpu().numpy()
        
        # Create masked version of original
        h, w = orig_map.shape
        ph, pw = h // 8, w // 8  # Assuming patch size is 8
        
        # Reshape mask (handle if it includes CLS token)
        if len(mask) > ph*pw:
            mask = mask[1:]  # Remove CLS token
        
        mask_2d = mask.reshape(ph, pw)
        mask_full = np.repeat(np.repeat(mask_2d, 8, axis=0), 8, axis=1)
        
        # Apply mask (0 = masked in our convention)
        masked_map = orig_map.copy()
        masked_map[mask_full < 0.5] = 0
        
        # Create figure
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        
        # Original
        im0 = axes[0].imshow(orig_map, cmap='viridis')
        axes[0].set_title("Original")
        plt.colorbar(im0, ax=axes[0])
        
        # Masked (input)
        im1 = axes[1].imshow(masked_map, cmap='viridis')
        axes[1].set_title("Masked (Input)")
        plt.colorbar(im1, ax=axes[1])
        
        # Reconstruction
        im2 = axes[2].imshow(recon_map, cmap='viridis')
        axes[2].set_title("Reconstructed")
        plt.colorbar(im2, ax=axes[2])
        
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'visualizations', f'sample_{i}.png'))
        plt.close()
    
    print(f"Saved {num_examples} visualizations to {os.path.join(output_dir, 'visualizations')}")

def main():
    args = parse_args()
    
    # Set device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    # Create output directory
    os.makedirs(args.output_dir, exist_ok=True)
    
    # Import the model (do this here to avoid circular imports)
    from model import MaskedAutoencoderViT
    from dataset import get_dataloaders
    
    # Create model with the SAME architecture as used in training
    # These parameters MUST match what you used in training
    model = MaskedAutoencoderViT(
        img_size=64,
        patch_size=8,
        in_channels=1,
        # Small model parameters - THIS IS IMPORTANT
        embed_dim=384,
        depth=6,
        num_heads=6,
        decoder_embed_dim=256, 
        decoder_depth=4,
        decoder_num_heads=8,
        mask_ratio=0.75
    ).to(device)
    
    # Load checkpoint
    print(f"Loading checkpoint from {args.checkpoint}")
    
    # Use a context manager to allow Namespace
    with safe_globals([Namespace]):
        checkpoint = torch.load(args.checkpoint, map_location=device)
    
    # If state_dict is nested in the checkpoint
    if 'state_dict' in checkpoint:
        state_dict = checkpoint['state_dict']
    else:
        state_dict = checkpoint
    
    # Load weights
    model.load_state_dict(state_dict)
    
    # Create dataloaders
    print("Loading dataset...")
    dataloaders = get_dataloaders(
        args.data_dir, 
        batch_size=16,
        num_workers=4,
        visualize_samples=False
    )
    
    # Visualize samples
    print("Visualizing samples...")
    visualize_samples(
        model,
        dataloaders['test'],
        device,
        args.output_dir,
        args.num_visualizations
    )
    
    print("Evaluation complete!")

if __name__ == '__main__':
    main()