#!/usr/bin/env python # -*- coding: utf-8 -*- """ Comprehensive X-ray Diffusion Model Evaluation Script Evaluates checkpoint_epoch_480.pt and extracts all possible metrics Usage: python evaluate_model.py """ import os import json import torch import numpy as np import matplotlib.pyplot as plt import seaborn as sns from pathlib import Path from tqdm import tqdm from PIL import Image from sklearn.manifold import TSNE import cv2 import logging import pandas as pd from matplotlib.colors import LinearSegmentedColormap import warnings from transformers import AutoTokenizer # Configure paths BASE_DIR = Path(__file__).parent CHECKPOINTS_DIR = BASE_DIR / "outputs" / "diffusion_checkpoints" VAE_CHECKPOINTS_DIR = BASE_DIR / "outputs" / "vae_checkpoints" DEFAULT_MODEL_PATH = str(CHECKPOINTS_DIR / "best_model.pt") TOKENIZER_NAME = os.environ.get("TOKENIZER_NAME", "dmis-lab/biobert-base-cased-v1.1") OUTPUT_DIR = os.environ.get("OUTPUT_DIR", str(BASE_DIR / "outputs" / "generated")) METRICS_DIR = BASE_DIR / "outputs" / "metrics" DATASET_PATH = os.environ.get("DATASET_PATH", str(BASE_DIR / "dataset")) IMAGES_PATH = os.environ.get("IMAGES_PATH", str(Path(DATASET_PATH) / "images" / "images_normalized")) # Import project modules from xray_generator.models.diffusion import DiffusionModel from xray_generator.models.vae import MedicalVAE from xray_generator.models.text_encoder import MedicalTextEncoder from xray_generator.models.unet import DiffusionUNet from xray_generator.utils.processing import get_device, apply_clahe, create_transforms from xray_generator.utils.dataset import ChestXrayDataset # Set up logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Suppress specific warnings warnings.filterwarnings('ignore', category=UserWarning) # Create directories if they don't exist os.makedirs(OUTPUT_DIR, exist_ok=True) os.makedirs(METRICS_DIR, exist_ok=True) os.makedirs(os.path.join(OUTPUT_DIR, "visualizations"), exist_ok=True) # Configure device device = get_device() logger.info(f"Using device: {device}") def load_diffusion_model(checkpoint_path): """Load a diffusion model from checkpoint""" logger.info(f"Loading diffusion model from {checkpoint_path}") try: checkpoint = torch.load(checkpoint_path, map_location=device) # Get model configuration config = checkpoint.get('config', {}) latent_channels = config.get('latent_channels', 8) model_channels = config.get('model_channels', 48) # Initialize model components vae = MedicalVAE( in_channels=1, out_channels=1, latent_channels=latent_channels, hidden_dims=[model_channels, model_channels*2, model_channels*4, model_channels*8] ).to(device) text_encoder = MedicalTextEncoder( model_name=config.get('text_model', "dmis-lab/biobert-base-cased-v1.1"), projection_dim=768, freeze_base=True ).to(device) unet = DiffusionUNet( in_channels=latent_channels, model_channels=model_channels, out_channels=latent_channels, num_res_blocks=2, attention_resolutions=(8, 16, 32), dropout=0.1, channel_mult=(1, 2, 4, 8), context_dim=768 ).to(device) # Load state dictionaries if 'vae_state_dict' in checkpoint: vae.load_state_dict(checkpoint['vae_state_dict']) logger.info("Loaded VAE weights") if 'text_encoder_state_dict' in checkpoint: text_encoder.load_state_dict(checkpoint['text_encoder_state_dict']) logger.info("Loaded text encoder weights") if 'unet_state_dict' in checkpoint: unet.load_state_dict(checkpoint['unet_state_dict']) logger.info("Loaded UNet weights") # Create diffusion model model = DiffusionModel( vae=vae, unet=unet, text_encoder=text_encoder, scheduler_type=config.get('scheduler_type', "ddim"), num_train_timesteps=config.get('num_train_timesteps', 1000), beta_schedule=config.get('beta_schedule', "linear"), prediction_type=config.get('prediction_type', "epsilon"), guidance_scale=config.get('guidance_scale', 7.5), device=device ) return model, checkpoint except Exception as e: logger.error(f"Error loading model: {e}") import traceback logger.error(traceback.format_exc()) raise RuntimeError(f"Failed to load model: {e}") def load_tokenizer(): """Load tokenizer for text conditioning""" try: tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME) logger.info(f"Loaded tokenizer: {TOKENIZER_NAME}") return tokenizer except Exception as e: logger.error(f"Error loading tokenizer: {e}") return None def load_dataset(split_ratio=0.1): """Load a small subset of the dataset for evaluation""" # Check if dataset path exists if not os.path.exists(DATASET_PATH): logger.error(f"Dataset path {DATASET_PATH} does not exist.") return None # Try to find the reports and projections CSV files reports_csv = None projections_csv = None for root, dirs, files in os.walk(BASE_DIR): for file in files: if file.endswith('.csv'): if 'report' in file.lower(): reports_csv = os.path.join(root, file) elif 'projection' in file.lower(): projections_csv = os.path.join(root, file) if not reports_csv or not projections_csv: logger.error(f"Could not find reports or projections CSV files.") logger.info("Creating dummy dataset for evaluation...") # Create a dummy dataset with random noise class DummyDataset: def __init__(self, size=50): self.size = size def __len__(self): return self.size def __getitem__(self, idx): # Create random image img = torch.randn(1, 256, 256) # Normalize to [-1, 1] img = torch.clamp(img, -1, 1) # Create dummy text report = "Normal chest X-ray with no significant findings." # Create dummy encoding input_ids = torch.ones(256, dtype=torch.long) attention_mask = torch.ones(256, dtype=torch.long) return { 'image': img, 'report': report, 'input_ids': input_ids, 'attention_mask': attention_mask, 'uid': f'dummy_{idx}', 'filename': f'dummy_{idx}.png' } dataset = DummyDataset() logger.info(f"Created dummy dataset with {len(dataset)} samples") # Create dataloader from torch.utils.data import DataLoader from xray_generator.utils.processing import custom_collate_fn dataloader = DataLoader( dataset, batch_size=8, shuffle=False, collate_fn=custom_collate_fn ) return dataloader # Load the actual dataset logger.info(f"Loading dataset from {DATASET_PATH}") logger.info(f"Reports CSV: {reports_csv}") logger.info(f"Projections CSV: {projections_csv}") try: # Create transforms _, val_transform = create_transforms(256) # Create dataset dataset = ChestXrayDataset( reports_csv=reports_csv, projections_csv=projections_csv, image_folder=IMAGES_PATH, # Use the images subfolder path transform=val_transform, target_size=(256, 256), filter_frontal=True, tokenizer_name=TOKENIZER_NAME, max_length=256, use_clahe=True ) # Take a small subset for evaluation from torch.utils.data import Subset import random # Set seed for reproducibility random.seed(42) # Select random subset of indices indices = random.sample(range(len(dataset)), max(1, int(len(dataset) * split_ratio))) subset = Subset(dataset, indices) # Create dataloader from torch.utils.data import DataLoader from xray_generator.utils.processing import custom_collate_fn dataloader = DataLoader( subset, batch_size=8, shuffle=False, collate_fn=custom_collate_fn ) logger.info(f"Created dataloader with {len(subset)} samples") return dataloader except Exception as e: logger.error(f"Error loading dataset: {e}") return None class ModelMetrics: """Class to extract and calculate metrics from the model""" def __init__(self, model, checkpoint): self.model = model self.checkpoint = checkpoint self.metrics = {} def extract_checkpoint_metadata(self): """Extract metadata from the checkpoint""" metadata = {} # Extract epoch number if available if 'epoch' in self.checkpoint: metadata['epoch'] = self.checkpoint['epoch'] # Extract loss values if available if 'best_metrics' in self.checkpoint: metadata['best_metrics'] = self.checkpoint['best_metrics'] # Extract optimizer state if available if 'optimizer_state_dict' in self.checkpoint: optimizer = self.checkpoint['optimizer_state_dict'] if 'param_groups' in optimizer: metadata['optimizer_param_groups'] = len(optimizer['param_groups']) if len(optimizer['param_groups']) > 0: metadata['learning_rate'] = optimizer['param_groups'][0].get('lr', None) # Extract model config if available if 'config' in self.checkpoint: metadata['config'] = self.checkpoint['config'] # Extract scheduler state if available if 'scheduler_state_dict' in self.checkpoint: metadata['scheduler_state_present'] = True # Extract global step if available if 'global_step' in self.checkpoint: metadata['global_step'] = self.checkpoint['global_step'] self.metrics['checkpoint_metadata'] = metadata return metadata def extract_model_architecture(self): """Extract model architecture information""" architecture = {} # VAE architecture vae_info = { 'in_channels': self.model.vae.encoder.conv_in.in_channels, 'out_channels': self.model.vae.decoder.final[-1].out_channels, 'latent_channels': self.model.vae.latent_channels, 'encoder_blocks': len(self.model.vae.encoder.down_blocks), 'decoder_blocks': len(self.model.vae.decoder.up_blocks), } # UNet architecture unet_info = { 'in_channels': self.model.unet.in_channels, 'out_channels': self.model.unet.out_channels, 'model_channels': self.model.unet.model_channels, 'attention_resolutions': self.model.unet.attention_resolutions, 'channel_mult': self.model.unet.channel_mult, 'context_dim': self.model.unet.context_dim, 'input_blocks': len(self.model.unet.input_blocks), 'output_blocks': len(self.model.unet.output_blocks), } # Text encoder architecture text_encoder_info = { 'model_name': self.model.text_encoder.model_name, 'hidden_dim': self.model.text_encoder.hidden_dim, 'projection_dim': self.model.text_encoder.projection_dim, } # Diffusion process parameters diffusion_info = { 'scheduler_type': self.model.scheduler_type, 'num_train_timesteps': self.model.num_train_timesteps, 'beta_schedule': self.model.beta_schedule, 'prediction_type': self.model.prediction_type, 'guidance_scale': self.model.guidance_scale, } architecture['vae'] = vae_info architecture['unet'] = unet_info architecture['text_encoder'] = text_encoder_info architecture['diffusion'] = diffusion_info self.metrics['architecture'] = architecture return architecture def count_parameters(self): """Count model parameters""" param_counts = {} def count_params(model): return sum(p.numel() for p in model.parameters()) def count_trainable_params(model): return sum(p.numel() for p in model.parameters() if p.requires_grad) # VAE parameters param_counts['vae_total'] = count_params(self.model.vae) param_counts['vae_trainable'] = count_trainable_params(self.model.vae) # UNet parameters param_counts['unet_total'] = count_params(self.model.unet) param_counts['unet_trainable'] = count_trainable_params(self.model.unet) # Text encoder parameters param_counts['text_encoder_total'] = count_params(self.model.text_encoder) param_counts['text_encoder_trainable'] = count_trainable_params(self.model.text_encoder) # Total parameters param_counts['total'] = param_counts['vae_total'] + param_counts['unet_total'] + param_counts['text_encoder_total'] param_counts['trainable'] = param_counts['vae_trainable'] + param_counts['unet_trainable'] + param_counts['text_encoder_trainable'] # Memory footprint (in MB) param_memory = 0 buffer_memory = 0 for module in [self.model.vae, self.model.unet, self.model.text_encoder]: param_memory += sum(p.nelement() * p.element_size() for p in module.parameters()) buffer_memory += sum(b.nelement() * b.element_size() for b in module.buffers()) param_counts['memory_footprint_mb'] = (param_memory + buffer_memory) / (1024 * 1024) self.metrics['parameters'] = param_counts return param_counts def analyze_beta_schedule(self): """Analyze the beta schedule used in the diffusion model""" beta_info = {} # Get beta schedule info betas = self.model.betas.cpu().numpy() beta_info['min'] = float(betas.min()) beta_info['max'] = float(betas.max()) beta_info['mean'] = float(betas.mean()) beta_info['std'] = float(betas.std()) # Get alphas info alphas_cumprod = self.model.alphas_cumprod.cpu().numpy() beta_info['alphas_cumprod_min'] = float(alphas_cumprod.min()) beta_info['alphas_cumprod_max'] = float(alphas_cumprod.max()) # Plot beta schedule plt.figure(figsize=(10, 6)) plt.plot(betas, label='Beta Schedule') plt.xlabel('Timestep') plt.ylabel('Beta Value') plt.title(f'Beta Schedule ({self.model.beta_schedule})') plt.legend() plt.grid(True, alpha=0.3) plt.savefig(os.path.join(OUTPUT_DIR, 'visualizations', 'beta_schedule.png')) plt.close() # Plot alphas_cumprod plt.figure(figsize=(10, 6)) plt.plot(alphas_cumprod, label='Cumulative Product of Alphas') plt.xlabel('Timestep') plt.ylabel('Alpha Cumprod Value') plt.title('Alphas Cumulative Product') plt.legend() plt.grid(True, alpha=0.3) plt.savefig(os.path.join(OUTPUT_DIR, 'visualizations', 'alphas_cumprod.png')) plt.close() self.metrics['beta_schedule'] = beta_info return beta_info def analyze_vae_latent_space(self, dataloader): """Analyze the VAE latent space""" logger.info("Analyzing VAE latent space...") latent_info = {} latent_vectors = [] orig_images = [] recon_images = [] # Set model to eval mode self.model.vae.eval() with torch.no_grad(): # Process a few batches for i, batch in enumerate(tqdm(dataloader, desc="Processing batches")): if i >= 5: # Limit to 5 batches for efficiency break # Get images images = batch['image'].to(device) # Get latent vectors mu, logvar = self.model.vae.encode(images) # Store latent vectors latent_vectors.append(mu.cpu().numpy()) # Store original images (first batch only) if i == 0: orig_images = images[:8].cpu() # Store up to 8 images # Generate reconstructions recon, _, _ = self.model.vae(images[:8]) recon_images = recon.cpu() # Concatenate latent vectors latent_vectors = np.concatenate(latent_vectors, axis=0) # Calculate latent space statistics latent_info['mean'] = float(np.mean(latent_vectors)) latent_info['std'] = float(np.std(latent_vectors)) latent_info['min'] = float(np.min(latent_vectors)) latent_info['max'] = float(np.max(latent_vectors)) latent_info['dimensions'] = latent_vectors.shape[1] # Calculate active dimensions (standard deviation > 0.1) active_dims = np.sum(np.std(latent_vectors, axis=0) > 0.1) latent_info['active_dimensions'] = int(active_dims) latent_info['active_dimensions_ratio'] = float(active_dims / latent_vectors.shape[1]) # Save visualization of latent space (t-SNE) if len(latent_vectors) > 10: try: # Subsample for efficiency sample_indices = np.random.choice(len(latent_vectors), min(500, len(latent_vectors)), replace=False) sampled_vectors = latent_vectors[sample_indices] # Apply t-SNE tsne = TSNE(n_components=2, random_state=42) latent_2d = tsne.fit_transform(sampled_vectors.reshape(sampled_vectors.shape[0], -1)) # Plot t-SNE plt.figure(figsize=(10, 10)) plt.scatter(latent_2d[:, 0], latent_2d[:, 1], alpha=0.5) plt.title("t-SNE Visualization of VAE Latent Space") plt.colorbar() plt.savefig(os.path.join(OUTPUT_DIR, 'visualizations', 'vae_latent_tsne.png')) plt.close() except Exception as e: logger.error(f"Error creating t-SNE visualization: {e}") # Save original and reconstructed images if len(orig_images) > 0 and len(recon_images) > 0: # Combine into grid from torchvision.utils import make_grid # Denormalize from [-1, 1] to [0, 1] orig_images = (orig_images + 1) / 2 recon_images = (recon_images + 1) / 2 # Create comparison grid comparison = torch.cat([make_grid(orig_images, nrow=4, padding=2), make_grid(recon_images, nrow=4, padding=2)], dim=2) # Save grid from torchvision.utils import save_image save_image(comparison, os.path.join(OUTPUT_DIR, 'visualizations', 'vae_reconstruction.png')) # Calculate reconstruction error mse = torch.mean((orig_images - recon_images) ** 2).item() latent_info['reconstruction_mse'] = mse self.metrics['vae_latent'] = latent_info return latent_info def generate_samples(self, tokenizer, num_samples=4): """Generate samples from the diffusion model""" logger.info("Generating samples from diffusion model...") # Set model to eval mode self.model.vae.eval() self.model.unet.eval() self.model.text_encoder.eval() # Sample prompts prompts = [ "Normal chest X-ray with clear lungs and no abnormalities.", "Right lower lobe pneumonia with focal consolidation.", "Mild cardiomegaly with pulmonary edema.", "Left pleural effusion with adjacent atelectasis." ] # Create folder for samples samples_dir = os.path.join(OUTPUT_DIR, 'samples') os.makedirs(samples_dir, exist_ok=True) generated_samples = [] with torch.no_grad(): for i, prompt in enumerate(tqdm(prompts[:num_samples], desc="Generating samples")): try: # Generate sample results = self.model.sample( prompt, height=256, width=256, num_inference_steps=50, tokenizer=tokenizer ) # Get image img = results['images'][0] # Convert to numpy and save img_np = img.cpu().numpy().transpose(1, 2, 0) img_np = (img_np * 255).astype(np.uint8) # Remove channel dimension for grayscale if img_np.shape[-1] == 1: img_np = img_np.squeeze(-1) # Save image img_path = os.path.join(samples_dir, f"sample_{i+1}.png") Image.fromarray(img_np).save(img_path) # Save prompt prompt_path = os.path.join(samples_dir, f"prompt_{i+1}.txt") with open(prompt_path, "w") as f: f.write(prompt) # Store generated sample generated_samples.append({ 'prompt': prompt, 'image_path': img_path }) except Exception as e: logger.error(f"Error generating sample {i+1}: {e}") continue # Create a grid of all samples try: # Read all samples sample_images = [] for i in range(num_samples): img_path = os.path.join(samples_dir, f"sample_{i+1}.png") if os.path.exists(img_path): img = Image.open(img_path) img_tensor = torch.tensor(np.array(img) / 255.0).unsqueeze(0) if len(img_tensor.shape) == 3: # Add channel dimension if needed img_tensor = img_tensor.unsqueeze(0) else: img_tensor = img_tensor.permute(0, 3, 1, 2) sample_images.append(img_tensor) if sample_images: # Create grid from torchvision.utils import make_grid grid = make_grid(torch.cat(sample_images, dim=0), nrow=2, padding=2) # Save grid from torchvision.utils import save_image save_image(grid, os.path.join(OUTPUT_DIR, 'visualizations', 'generated_samples_grid.png')) except Exception as e: logger.error(f"Error creating sample grid: {e}") self.metrics['generated_samples'] = generated_samples return generated_samples def measure_inference_speed(self, tokenizer, num_runs=10): """Measure inference speed""" logger.info("Measuring inference speed...") # Set model to eval mode self.model.vae.eval() self.model.unet.eval() self.model.text_encoder.eval() # Sample prompt prompt = "Normal chest X-ray with clear lungs and no abnormalities." # Warm-up run logger.info("Performing warm-up run...") with torch.no_grad(): _ = self.model.sample( prompt, height=256, width=256, num_inference_steps=20, # Use fewer steps for speed tokenizer=tokenizer ) # Measure inference time logger.info(f"Measuring inference time over {num_runs} runs...") inference_times = [] for i in range(num_runs): start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) # Synchronize CUDA operations torch.cuda.synchronize() start.record() with torch.no_grad(): _ = self.model.sample( prompt, height=256, width=256, num_inference_steps=20, # Use fewer steps for speed tokenizer=tokenizer ) end.record() torch.cuda.synchronize() # Calculate elapsed time in milliseconds inference_time = start.elapsed_time(end) inference_times.append(inference_time) logger.info(f"Run {i+1}/{num_runs}: {inference_time:.2f} ms") # Calculate statistics avg_time = np.mean(inference_times) std_time = np.std(inference_times) inference_speed = { 'avg_inference_time_ms': float(avg_time), 'std_inference_time_ms': float(std_time), 'min_inference_time_ms': float(np.min(inference_times)), 'max_inference_time_ms': float(np.max(inference_times)), 'num_runs': num_runs, 'num_inference_steps': 20 } # Plot inference times plt.figure(figsize=(10, 6)) plt.bar(range(1, num_runs + 1), inference_times) plt.axhline(avg_time, color='r', linestyle='--', label=f'Avg: {avg_time:.2f} ms') plt.xlabel('Run #') plt.ylabel('Inference Time (ms)') plt.title('Diffusion Model Inference Time') plt.legend() plt.savefig(os.path.join(OUTPUT_DIR, 'visualizations', 'inference_time.png')) plt.close() self.metrics['inference_speed'] = inference_speed return inference_speed def visualize_unet_attention(self, tokenizer): """Visualize UNet attention maps""" logger.info("Visualizing UNet attention maps...") # This is a complex task and might need model code modification # Here we'll just create a placeholder for this analysis self.metrics['unet_attention'] = { 'note': 'UNet attention visualization requires model modifications to extract attention maps' } return self.metrics['unet_attention'] def visualize_noise_levels(self): """Visualize noise levels at different timesteps""" logger.info("Visualizing noise levels...") # Create a random image x_0 = torch.randn(1, 1, 256, 256).to(device) # Sample timesteps timesteps = torch.linspace(0, self.model.num_train_timesteps - 1, 10).long().to(device) # Create folder for noise visualizations noise_dir = os.path.join(OUTPUT_DIR, 'visualizations', 'noise_levels') os.makedirs(noise_dir, exist_ok=True) # Generate noisy samples at different timesteps with torch.no_grad(): for i, t in enumerate(timesteps): # Add noise noisy_x = self.model.q_sample(x_0, t.unsqueeze(0)) # Convert to image img = noisy_x[0].cpu() # Normalize to [0, 1] img = (img - img.min()) / (img.max() - img.min()) # Save image from torchvision.utils import save_image save_image(img, os.path.join(noise_dir, f"noise_t{t.item()}.png")) # Create a grid of noise levels try: # Read all noise images noise_images = [] for i, t in enumerate(timesteps): img_path = os.path.join(noise_dir, f"noise_t{t.item()}.png") if os.path.exists(img_path): img = Image.open(img_path) img_tensor = torch.tensor(np.array(img) / 255.0) if len(img_tensor.shape) == 2: # Add channel dimension if needed img_tensor = img_tensor.unsqueeze(0) else: img_tensor = img_tensor.permute(2, 0, 1) noise_images.append(img_tensor) if noise_images: # Create grid from torchvision.utils import make_grid grid = make_grid(torch.stack(noise_images), nrow=5, padding=2) # Save grid from torchvision.utils import save_image save_image(grid, os.path.join(OUTPUT_DIR, 'visualizations', 'noise_levels_grid.png')) except Exception as e: logger.error(f"Error creating noise levels grid: {e}") self.metrics['noise_levels'] = { 'timesteps': timesteps.cpu().numpy().tolist(), 'visualization_path': noise_dir } return self.metrics['noise_levels'] def plot_learning_curves(self): """Plot learning curves if available in checkpoint""" logger.info("Plotting learning curves...") # Check if loss values are available if 'best_metrics' not in self.checkpoint: logger.info("No loss values found in checkpoint") return None # Extract metrics metrics = self.checkpoint['best_metrics'] if 'train_loss' in metrics and 'val_loss' in metrics: # Plot training and validation loss plt.figure(figsize=(10, 6)) plt.bar(['Training Loss', 'Validation Loss'], [metrics['train_loss'], metrics['val_loss']]) plt.ylabel('Loss') plt.title('Training and Validation Loss') plt.savefig(os.path.join(OUTPUT_DIR, 'visualizations', 'loss_comparison.png')) plt.close() if 'train_diffusion_loss' in metrics and 'val_diffusion_loss' in metrics: # Plot diffusion loss plt.figure(figsize=(10, 6)) plt.bar(['Training Diffusion Loss', 'Validation Diffusion Loss'], [metrics['train_diffusion_loss'], metrics['val_diffusion_loss']]) plt.ylabel('Diffusion Loss') plt.title('Diffusion Loss') plt.savefig(os.path.join(OUTPUT_DIR, 'visualizations', 'diffusion_loss.png')) plt.close() return metrics def create_parameter_distribution_plots(self): """Plot parameter distributions""" logger.info("Creating parameter distribution plots...") # Collect parameters from different components vae_params = torch.cat([p.detach().cpu().flatten() for p in self.model.vae.parameters()]) unet_params = torch.cat([p.detach().cpu().flatten() for p in self.model.unet.parameters()]) text_encoder_params = torch.cat([p.detach().cpu().flatten() for p in self.model.text_encoder.parameters()]) # Plot parameter distributions plt.figure(figsize=(15, 5)) plt.subplot(1, 3, 1) plt.hist(vae_params.numpy(), bins=50, alpha=0.7) plt.title('VAE Parameters') plt.xlabel('Value') plt.ylabel('Count') plt.subplot(1, 3, 2) plt.hist(unet_params.numpy(), bins=50, alpha=0.7) plt.title('UNet Parameters') plt.xlabel('Value') plt.ylabel('Count') plt.subplot(1, 3, 3) plt.hist(text_encoder_params.numpy(), bins=50, alpha=0.7) plt.title('Text Encoder Parameters') plt.xlabel('Value') plt.ylabel('Count') plt.tight_layout() plt.savefig(os.path.join(OUTPUT_DIR, 'visualizations', 'parameter_distributions.png')) plt.close() # Calculate statistics param_stats = { 'vae': { 'mean': float(vae_params.mean()), 'std': float(vae_params.std()), 'min': float(vae_params.min()), 'max': float(vae_params.max()) }, 'unet': { 'mean': float(unet_params.mean()), 'std': float(unet_params.std()), 'min': float(unet_params.min()), 'max': float(unet_params.max()) }, 'text_encoder': { 'mean': float(text_encoder_params.mean()), 'std': float(text_encoder_params.std()), 'min': float(text_encoder_params.min()), 'max': float(text_encoder_params.max()) } } self.metrics['parameter_stats'] = param_stats return param_stats def generate_text_conditioning_analysis(self, tokenizer): """Analyze the effect of text conditioning on generation""" logger.info("Generating text conditioning analysis...") if tokenizer is None: logger.error("Tokenizer is required for text conditioning analysis") return None # Create a test case with multiple prompts test_prompts = [ "Normal chest X-ray with no abnormalities.", "Severe pneumonia with bilateral infiltrates.", "Cardiomegaly with pulmonary edema.", "Pneumothorax with collapsed left lung." ] # Create folder for text conditioning visualizations text_dir = os.path.join(OUTPUT_DIR, 'visualizations', 'text_conditioning') os.makedirs(text_dir, exist_ok=True) # Generate samples for each prompt generated_images = [] with torch.no_grad(): # Generate one sample with fixed seed for each prompt for i, prompt in enumerate(tqdm(test_prompts, desc="Generating conditioned samples")): try: # Set seed for reproducibility torch.manual_seed(42) # Generate sample results = self.model.sample( prompt, height=256, width=256, num_inference_steps=50, tokenizer=tokenizer ) # Get image img = results['images'][0] # Save image img_np = img.cpu().numpy().transpose(1, 2, 0) img_np = (img_np * 255).astype(np.uint8) if img_np.shape[-1] == 1: img_np = img_np.squeeze(-1) img_path = os.path.join(text_dir, f"prompt_{i+1}.png") Image.fromarray(img_np).save(img_path) # Save prompt prompt_path = os.path.join(text_dir, f"prompt_{i+1}.txt") with open(prompt_path, "w") as f: f.write(prompt) # Store generated image generated_images.append(img.cpu()) except Exception as e: logger.error(f"Error generating sample for prompt {i+1}: {e}") continue # Create a grid of all samples if generated_images: try: # Create grid from torchvision.utils import make_grid grid = make_grid(torch.stack(generated_images), nrow=2, padding=2) # Save grid from torchvision.utils import save_image save_image(grid, os.path.join(OUTPUT_DIR, 'visualizations', 'text_conditioning_grid.png')) except Exception as e: logger.error(f"Error creating text conditioning grid: {e}") # Test different guidance scales on a single prompt guidance_scales = [1.0, 3.0, 7.5, 10.0, 15.0] guidance_images = [] with torch.no_grad(): # Generate samples with different guidance scales for i, scale in enumerate(tqdm(guidance_scales, desc="Testing guidance scales")): try: # Set seed for reproducibility torch.manual_seed(42) # Generate sample results = self.model.sample( test_prompts[0], # Use the first prompt height=256, width=256, num_inference_steps=50, guidance_scale=scale, tokenizer=tokenizer, seed=42 # Fixed seed ) # Get image img = results['images'][0] # Save image img_np = img.cpu().numpy().transpose(1, 2, 0) img_np = (img_np * 255).astype(np.uint8) if img_np.shape[-1] == 1: img_np = img_np.squeeze(-1) img_path = os.path.join(text_dir, f"guidance_{scale}.png") Image.fromarray(img_np).save(img_path) # Store generated image guidance_images.append(img.cpu()) except Exception as e: logger.error(f"Error generating sample for guidance scale {scale}: {e}") continue # Create a grid of guidance scale samples if guidance_images: try: # Create grid from torchvision.utils import make_grid grid = make_grid(torch.stack(guidance_images), nrow=len(guidance_scales), padding=2) # Save grid from torchvision.utils import save_image save_image(grid, os.path.join(OUTPUT_DIR, 'visualizations', 'guidance_scale_grid.png')) except Exception as e: logger.error(f"Error creating guidance scale grid: {e}") self.metrics['text_conditioning'] = { 'test_prompts': test_prompts, 'guidance_scales': guidance_scales, 'visualization_path': text_dir } return self.metrics['text_conditioning'] def analyze_all(self, dataloader, tokenizer): """Run all analysis methods and collect metrics""" # Extract checkpoint metadata self.extract_checkpoint_metadata() # Extract model architecture information self.extract_model_architecture() # Count parameters self.count_parameters() # Analyze beta schedule self.analyze_beta_schedule() # Analyze VAE latent space if dataloader is not None: self.analyze_vae_latent_space(dataloader) # Generate samples if tokenizer is not None: self.generate_samples(tokenizer) # Measure inference speed if tokenizer is not None: self.measure_inference_speed(tokenizer, num_runs=5) # Visualize UNet attention if tokenizer is not None: self.visualize_unet_attention(tokenizer) # Visualize noise levels self.visualize_noise_levels() # Plot learning curves self.plot_learning_curves() # Create parameter distribution plots self.create_parameter_distribution_plots() # Generate text conditioning analysis if tokenizer is not None: self.generate_text_conditioning_analysis(tokenizer) # Save all metrics to file with open(os.path.join(METRICS_DIR, 'diffusion_metrics.json'), 'w') as f: # Convert non-serializable values to strings or lists serializable_metrics = json.loads( json.dumps(self.metrics, default=lambda o: str(o) if not isinstance(o, (int, float, str, bool, list, dict, type(None))) else o) ) json.dump(serializable_metrics, f, indent=2) return self.metrics def create_model_summary(metrics): """Create a human-readable summary of model metrics""" logger.info("Creating model summary...") summary = [] # Add header summary.append("# X-ray Diffusion Model Evaluation Summary") summary.append("\n## Model Information") # Add model architecture if 'architecture' in metrics: arch = metrics['architecture'] summary.append("\n### Diffusion Model") summary.append(f"- Scheduler Type: {arch['diffusion']['scheduler_type']}") summary.append(f"- Timesteps: {arch['diffusion']['num_train_timesteps']}") summary.append(f"- Beta Schedule: {arch['diffusion']['beta_schedule']}") summary.append(f"- Prediction Type: {arch['diffusion']['prediction_type']}") summary.append(f"- Guidance Scale: {arch['diffusion']['guidance_scale']}") summary.append("\n### VAE") summary.append(f"- Latent Channels: {arch['vae']['latent_channels']}") summary.append(f"- Encoder Blocks: {arch['vae']['encoder_blocks']}") summary.append(f"- Decoder Blocks: {arch['vae']['decoder_blocks']}") summary.append("\n### UNet") summary.append(f"- Model Channels: {arch['unet']['model_channels']}") summary.append(f"- Attention Resolutions: {arch['unet']['attention_resolutions']}") summary.append(f"- Channel Multipliers: {arch['unet']['channel_mult']}") summary.append("\n### Text Encoder") summary.append(f"- Model: {arch['text_encoder']['model_name']}") summary.append(f"- Hidden Dimension: {arch['text_encoder']['hidden_dim']}") summary.append(f"- Projection Dimension: {arch['text_encoder']['projection_dim']}") # Add parameter counts if 'parameters' in metrics: params = metrics['parameters'] summary.append("\n## Parameter Counts") summary.append(f"- Total Parameters: {params['total']:,}") summary.append(f"- Trainable Parameters: {params['trainable']:,}") summary.append(f"- Memory Footprint: {params['memory_footprint_mb']:.2f} MB") summary.append("\n### Component Breakdown") summary.append(f"- VAE: {params['vae_total']:,} parameters ({params['vae_trainable']:,} trainable)") summary.append(f"- UNet: {params['unet_total']:,} parameters ({params['unet_trainable']:,} trainable)") summary.append(f"- Text Encoder: {params['text_encoder_total']:,} parameters ({params['text_encoder_trainable']:,} trainable)") # Add training information if 'checkpoint_metadata' in metrics: meta = metrics['checkpoint_metadata'] summary.append("\n## Training Information") if 'epoch' in meta: summary.append(f"- Trained for {meta['epoch']} epochs") if 'global_step' in meta: summary.append(f"- Global steps: {meta['global_step']}") if 'best_metrics' in meta: summary.append("\n### Best Metrics") best = meta['best_metrics'] for key, value in best.items(): summary.append(f"- {key}: {value}") # Add VAE latent information if 'vae_latent' in metrics: latent = metrics['vae_latent'] summary.append("\n## VAE Latent Space Analysis") summary.append(f"- Latent Dimensions: {latent.get('dimensions', 'N/A')}") summary.append(f"- Active Dimensions: {latent.get('active_dimensions', 'N/A')} ({latent.get('active_dimensions_ratio', 'N/A'):.2%})") if 'reconstruction_mse' in latent: summary.append(f"- Reconstruction MSE: {latent['reconstruction_mse']:.6f}") # Add inference speed if 'inference_speed' in metrics: speed = metrics['inference_speed'] summary.append("\n## Inference Performance") summary.append(f"- Average Inference Time: {speed['avg_inference_time_ms']:.2f} ms") summary.append(f"- Standard Deviation: {speed['std_inference_time_ms']:.2f} ms") summary.append(f"- Range: {speed['min_inference_time_ms']:.2f} - {speed['max_inference_time_ms']:.2f} ms") # Add visualization paths summary.append("\n## Visualizations") summary.append(f"- All visualizations saved to: {os.path.join(OUTPUT_DIR, 'visualizations')}") if 'generated_samples' in metrics: summary.append(f"- Generated samples saved to: {os.path.join(OUTPUT_DIR, 'samples')}") # Save summary to file summary_text = "\n".join(summary) with open(os.path.join(METRICS_DIR, 'model_summary.md'), 'w') as f: f.write(summary_text) logger.info(f"Model summary saved to {os.path.join(METRICS_DIR, 'model_summary.md')}") return summary_text def main(): """Main function to run all analyses""" logger.info("Starting model evaluation script") # Load diffusion model from checkpoint diffusion_model, checkpoint = load_diffusion_model( os.path.join(CHECKPOINTS_DIR, "checkpoint_epoch_480.pt") ) # Load tokenizer tokenizer = load_tokenizer() # Load dataset dataloader = load_dataset() # Create metrics calculator metrics_calculator = ModelMetrics(diffusion_model, checkpoint) # Run all analyses metrics = metrics_calculator.analyze_all(dataloader, tokenizer) # Create human-readable summary summary = create_model_summary(metrics) logger.info("Model evaluation complete") logger.info(f"Results saved to {METRICS_DIR}") logger.info(f"Visualizations saved to {os.path.join(OUTPUT_DIR, 'visualizations')}") if __name__ == "__main__": main()