lwm-spectro / example_inference.py

Upload example_inference.py with huggingface_hub

652afc6 verified 12 months ago

11.3 kB

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

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