Add inference.py

inference.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+"""
+Inference script for the conditional diffusion model.
+This script provides easy-to-use functions for generating medical images.
+"""
+
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from model import UNet, marginal_prob_std, diffusion_coeff, Euler_Maruyama_sampler
+
+
+class ConditionalDiffusionInference:
+    """Wrapper class for easy inference with the conditional diffusion model."""
+    
+    def __init__(self, model_path, device='cuda'):
+        """
+        Initialize the inference model.
+        
+        Args:
+            model_path: Path to the trained model checkpoint
+            device: Device to run inference on ('cuda' or 'cpu')
+        """
+        self.device = device
+        self.Lambda = 25.0
+        
+        # Initialize the model
+        self.marginal_prob_std_fn = lambda t: marginal_prob_std(t, Lambda=self.Lambda, device=self.device)
+        self.diffusion_coeff_fn = lambda t: diffusion_coeff(t, Lambda=self.Lambda, device=self.device)
+        
+        self.score_model = UNet(marginal_prob_std=self.marginal_prob_std_fn)
+        self.score_model.load_state_dict(torch.load(model_path, map_location=self.device))
+        self.score_model.to(self.device)
+        self.score_model.eval()
+        
+        print(f"Model loaded successfully on {self.device}")
+    
+    def generate_image(self, conditioning_mask, num_steps=250, eps=1e-3):
+        """
+        Generate a medical image based on a conditioning mask.
+        
+        Args:
+            conditioning_mask: Conditioning mask tensor of shape (1, 4, 256, 256)
+            num_steps: Number of sampling steps
+            eps: Smallest time step for numerical stability
+            
+        Returns:
+            Generated image tensor of shape (1, 4, 256, 256)
+        """
+        if not isinstance(conditioning_mask, torch.Tensor):
+            conditioning_mask = torch.tensor(conditioning_mask, dtype=torch.float32)
+        
+        if conditioning_mask.dim() == 3:
+            conditioning_mask = conditioning_mask.unsqueeze(0)
+        
+        conditioning_mask = conditioning_mask.to(self.device)
+        
+        with torch.no_grad():
+            samples = Euler_Maruyama_sampler(
+                self.score_model,
+                self.marginal_prob_std_fn,
+                self.diffusion_coeff_fn,
+                batch_size=1,
+                x_shape=(4, 256, 256),
+                num_steps=num_steps,
+                device=self.device,
+                eps=eps,
+                y=conditioning_mask
+            )
+        
+        return samples.clamp(0, 1)
+    
+    def generate_batch(self, conditioning_masks, num_steps=250, eps=1e-3):
+        """
+        Generate multiple images based on conditioning masks.
+        
+        Args:
+            conditioning_masks: Conditioning mask tensor of shape (B, 4, 256, 256)
+            num_steps: Number of sampling steps
+            eps: Smallest time step for numerical stability
+            
+        Returns:
+            Generated images tensor of shape (B, 4, 256, 256)
+        """
+        if not isinstance(conditioning_masks, torch.Tensor):
+            conditioning_masks = torch.tensor(conditioning_masks, dtype=torch.float32)
+        
+        if conditioning_masks.dim() == 3:
+            conditioning_masks = conditioning_masks.unsqueeze(0)
+        
+        conditioning_masks = conditioning_masks.to(self.device)
+        batch_size = conditioning_masks.shape[0]
+        
+        with torch.no_grad():
+            samples = Euler_Maruyama_sampler(
+                self.score_model,
+                self.marginal_prob_std_fn,
+                self.diffusion_coeff_fn,
+                batch_size=batch_size,
+                x_shape=(4, 256, 256),
+                num_steps=num_steps,
+                device=self.device,
+                eps=eps,
+                y=conditioning_masks
+            )
+        
+        return samples.clamp(0, 1)
+    
+    def visualize_generation(self, conditioning_mask, generated_image, save_path=None):
+        """
+        Visualize the conditioning mask and generated image.
+        
+        Args:
+            conditioning_mask: Conditioning mask tensor
+            generated_image: Generated image tensor
+            save_path: Optional path to save the visualization
+        """
+        fig, axes = plt.subplots(2, 4, figsize=(16, 8))
+        
+        # Plot conditioning mask
+        for i in range(4):
+            axes[0, i].imshow(conditioning_mask[0, i].cpu().numpy(), cmap='gray')
+            axes[0, i].set_title(f'Conditioning Mask {i+1}')
+            axes[0, i].axis('off')
+        
+        # Plot generated image
+        for i in range(4):
+            axes[1, i].imshow(generated_image[0, i].cpu().numpy(), cmap='gray')
+            axes[1, i].set_title(f'Generated Image {i+1}')
+            axes[1, i].axis('off')
+        
+        plt.tight_layout()
+        
+        if save_path:
+            plt.savefig(save_path, dpi=150, bbox_inches='tight')
+            print(f"Visualization saved to {save_path}")
+        
+        plt.show()
+
+
+def main():
+    """Example usage of the inference model."""
+    
+    # Initialize the model
+    model_path = "ckpt_3D_v2.pth"  # Update with your model path
+    inference_model = ConditionalDiffusionInference(model_path, device='cuda')
+    
+    # Create a random conditioning mask (replace with your actual mask)
+    conditioning_mask = torch.randn(1, 4, 256, 256)
+    
+    # Generate image
+    print("Generating image...")
+    generated_image = inference_model.generate_image(conditioning_mask)
+    
+    # Visualize results
+    inference_model.visualize_generation(conditioning_mask, generated_image, "generation_result.png")
+    
+    print("Generation complete!")
+
+
+if __name__ == "__main__":
+    main()