Upload example_inference.py with huggingface_hub

example_inference.py

@@ -0,0 +1,334 @@
+"""
+Example Inference Script for LWM-Spectro Model
+
+This script demonstrates how to:
+1. Load the pre-trained MoE model
+2. Load and preprocess a spectrogram
+3. Perform inference
+4. Interpret results
+"""
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+from pathlib import Path
+import sys
+
+# Add project root to path
+sys.path.append(str(Path(__file__).parent))
+
+from pretraining.pretrained_model import PretrainedLWM
+
+
+class SpectrogramClassifier:
+    """Wrapper class for easy inference with LWM-Spectro model"""
+    
+    def __init__(self, model_path, device='cuda'):
+        """
+        Initialize the classifier
+        
+        Args:
+            model_path: Path to the trained model checkpoint (.pth file)
+            device: 'cuda' or 'cpu'
+        """
+        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
+        print(f"Using device: {self.device}")
+        
+        # Load model
+        self.model = self._load_model(model_path)
+        self.model.eval()
+        
+        # Class mapping
+        self.classes = ['LTE', 'WiFi', '5G']
+        
+    def _load_model(self, model_path):
+        """Load the trained model from checkpoint"""
+        checkpoint = torch.load(model_path, map_location=self.device)
+        
+        # Handle different checkpoint formats
+        if isinstance(checkpoint, dict):
+            if 'model_state_dict' in checkpoint:
+                state_dict = checkpoint['model_state_dict']
+            elif 'state_dict' in checkpoint:
+                state_dict = checkpoint['state_dict']
+            else:
+                state_dict = checkpoint
+        else:
+            state_dict = checkpoint
+        
+        # Initialize model (adjust architecture as needed)
+        model = PretrainedLWM()  # or your specific model class
+        
+        # Load state dict
+        model.load_state_dict(state_dict, strict=False)
+        model.to(self.device)
+        
+        return model
+    
+    def load_spectrogram(self, image_path, target_size=(128, 128)):
+        """
+        Load and preprocess a spectrogram image
+        
+        Args:
+            image_path: Path to spectrogram image file
+            target_size: Target size for resizing (height, width)
+            
+        Returns:
+            Preprocessed tensor ready for inference
+        """
+        # Load image
+        img = Image.open(image_path).convert('L')  # Convert to grayscale
+        
+        # Resize
+        img = img.resize((target_size[1], target_size[0]), Image.BILINEAR)
+        
+        # Convert to numpy array and normalize
+        img_array = np.array(img, dtype=np.float32) / 255.0
+        
+        # Convert to tensor [1, 1, H, W]
+        tensor = torch.from_numpy(img_array).unsqueeze(0).unsqueeze(0)
+        
+        return tensor.to(self.device)
+    
+    def predict(self, spectrogram, return_probs=False):
+        """
+        Perform inference on a spectrogram
+        
+        Args:
+            spectrogram: Preprocessed spectrogram tensor or path to image file
+            return_probs: If True, return class probabilities along with prediction
+            
+        Returns:
+            If return_probs=False: predicted class name
+            If return_probs=True: (predicted class name, probability dict)
+        """
+        # Load spectrogram if path is provided
+        if isinstance(spectrogram, (str, Path)):
+            spectrogram = self.load_spectrogram(spectrogram)
+        
+        # Inference
+        with torch.no_grad():
+            output = self.model(spectrogram)
+            probabilities = F.softmax(output, dim=1)
+            predicted_idx = torch.argmax(probabilities, dim=1).item()
+            
+        predicted_class = self.classes[predicted_idx]
+        
+        if return_probs:
+            prob_dict = {
+                cls: probabilities[0, i].item() 
+                for i, cls in enumerate(self.classes)
+            }
+            return predicted_class, prob_dict
+        
+        return predicted_class
+    
+    def predict_batch(self, spectrogram_paths):
+        """
+        Perform batch inference on multiple spectrograms
+        
+        Args:
+            spectrogram_paths: List of paths to spectrogram images
+            
+        Returns:
+            List of predictions
+        """
+        predictions = []
+        for path in spectrogram_paths:
+            pred = self.predict(path)
+            predictions.append(pred)
+        
+        return predictions
+    
+    def visualize_prediction(self, image_path, save_path=None):
+        """
+        Visualize spectrogram with prediction
+        
+        Args:
+            image_path: Path to spectrogram image
+            save_path: Optional path to save the visualization
+        """
+        # Load original image for display
+        img = Image.open(image_path)
+        
+        # Get prediction with probabilities
+        pred_class, probs = self.predict(image_path, return_probs=True)
+        
+        # Create visualization
+        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
+        
+        # Display spectrogram
+        ax1.imshow(img, cmap='viridis')
+        ax1.set_title(f'Input Spectrogram\nPredicted: {pred_class}', fontsize=14, fontweight='bold')
+        ax1.axis('off')
+        
+        # Display probability distribution
+        classes = list(probs.keys())
+        probabilities = list(probs.values())
+        colors = ['#1f77b4', '#ff7f0e', '#2ca02c']
+        
+        bars = ax2.barh(classes, probabilities, color=colors)
+        ax2.set_xlabel('Probability', fontsize=12)
+        ax2.set_title('Class Probabilities', fontsize=14, fontweight='bold')
+        ax2.set_xlim(0, 1)
+        
+        # Add probability values on bars
+        for bar, prob in zip(bars, probabilities):
+            width = bar.get_width()
+            ax2.text(width, bar.get_y() + bar.get_height()/2, 
+                    f'{prob:.3f}', ha='left', va='center', fontsize=11)
+        
+        plt.tight_layout()
+        
+        if save_path:
+            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+            print(f"Visualization saved to: {save_path}")
+        
+        plt.show()
+
+
+# ============================================================================
+# Example Usage
+# ============================================================================
+
+def example_single_inference():
+    """Example: Single spectrogram inference"""
+    print("=" * 60)
+    print("Example 1: Single Spectrogram Inference")
+    print("=" * 60)
+    
+    # Initialize classifier
+    model_path = "mixture/runs/embedding_router/moe_checkpoint.pth"
+    classifier = SpectrogramClassifier(model_path, device='cuda')
+    
+    # Single inference
+    image_path = "spectrograms/5G/QPSK/rate1-2/SNR10dB/sample_0001.png"
+    prediction = classifier.predict(image_path)
+    print(f"\nPrediction: {prediction}")
+    
+    # With probabilities
+    pred_class, probs = classifier.predict(image_path, return_probs=True)
+    print(f"\nPredicted Class: {pred_class}")
+    print("\nClass Probabilities:")
+    for cls, prob in probs.items():
+        print(f"  {cls}: {prob:.4f}")
+
+
+def example_batch_inference():
+    """Example: Batch inference on multiple spectrograms"""
+    print("\n" + "=" * 60)
+    print("Example 2: Batch Inference")
+    print("=" * 60)
+    
+    # Initialize classifier
+    model_path = "mixture/runs/embedding_router/moe_checkpoint.pth"
+    classifier = SpectrogramClassifier(model_path, device='cuda')
+    
+    # Multiple images
+    image_paths = [
+        "spectrograms/5G/QPSK/rate1-2/SNR10dB/sample_0001.png",
+        "spectrograms/LTE/QAM16/rate1-2/SNR10dB/sample_0001.png",
+        "spectrograms/WiFi/QAM64/rate3-4/sample_0001.png",
+    ]
+    
+    # Batch prediction
+    predictions = classifier.predict_batch(image_paths)
+    
+    print("\nBatch Predictions:")
+    for path, pred in zip(image_paths, predictions):
+        print(f"  {Path(path).name}: {pred}")
+
+
+def example_visualization():
+    """Example: Visualize prediction with probabilities"""
+    print("\n" + "=" * 60)
+    print("Example 3: Prediction Visualization")
+    print("=" * 60)
+    
+    # Initialize classifier
+    model_path = "mixture/runs/embedding_router/moe_checkpoint.pth"
+    classifier = SpectrogramClassifier(model_path, device='cuda')
+    
+    # Visualize prediction
+    image_path = "spectrograms/5G/QPSK/rate1-2/SNR10dB/sample_0001.png"
+    classifier.visualize_prediction(image_path, save_path="prediction_result.png")
+
+
+def example_custom_preprocessing():
+    """Example: Custom preprocessing and inference"""
+    print("\n" + "=" * 60)
+    print("Example 4: Custom Preprocessing")
+    print("=" * 60)
+    
+    # Initialize classifier
+    model_path = "mixture/runs/embedding_router/moe_checkpoint.pth"
+    classifier = SpectrogramClassifier(model_path, device='cuda')
+    
+    # Load and custom preprocess
+    img = Image.open("spectrograms/5G/QPSK/rate1-2/SNR10dB/sample_0001.png")
+    img_array = np.array(img.convert('L'), dtype=np.float32) / 255.0
+    
+    # Apply custom transformations (example: add noise)
+    noise = np.random.normal(0, 0.01, img_array.shape)
+    img_array_noisy = np.clip(img_array + noise, 0, 1)
+    
+    # Convert to tensor
+    tensor = torch.from_numpy(img_array_noisy).unsqueeze(0).unsqueeze(0)
+    tensor = tensor.to(classifier.device)
+    
+    # Predict
+    prediction = classifier.predict(tensor)
+    print(f"\nPrediction on noisy image: {prediction}")
+
+
+def example_error_analysis():
+    """Example: Analyze predictions across different SNR levels"""
+    print("\n" + "=" * 60)
+    print("Example 5: SNR-based Error Analysis")
+    print("=" * 60)
+    
+    # Initialize classifier
+    model_path = "mixture/runs/embedding_router/moe_checkpoint.pth"
+    classifier = SpectrogramClassifier(model_path, device='cuda')
+    
+    # Test across different SNR levels
+    snr_levels = ['SNR-5dB', 'SNR0dB', 'SNR5dB', 'SNR10dB', 'SNR15dB', 'SNR20dB', 'SNR25dB']
+    base_path = Path("spectrograms/5G/QPSK/rate1-2")
+    
+    print("\nPredictions across SNR levels:")
+    for snr in snr_levels:
+        snr_path = base_path / snr / "sample_0001.png"
+        if snr_path.exists():
+            pred_class, probs = classifier.predict(str(snr_path), return_probs=True)
+            confidence = max(probs.values())
+            print(f"  {snr}: {pred_class} (confidence: {confidence:.3f})")
+
+
+if __name__ == "__main__":
+    print("\n" + "=" * 60)
+    print("LWM-Spectro Inference Examples")
+    print("=" * 60)
+    
+    try:
+        # Run examples
+        example_single_inference()
+        example_batch_inference()
+        example_visualization()
+        example_custom_preprocessing()
+        example_error_analysis()
+        
+        print("\n" + "=" * 60)
+        print("All examples completed successfully!")
+        print("=" * 60)
+        
+    except FileNotFoundError as e:
+        print(f"\nError: {e}")
+        print("\nNote: Update the file paths in the examples to match your directory structure.")
+    except Exception as e:
+        print(f"\nError: {e}")
+        print("\nPlease ensure:")
+        print("  1. Model checkpoint exists at specified path")
+        print("  2. Spectrogram images are available")
+        print("  3. All dependencies are installed")