demonstrate_augmentations.py

import os
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import random

# Create directory for visualization outputs
os.makedirs('/home/ubuntu/augmentation_demo', exist_ok=True)

# Function to load and crop a spectrogram image
def load_spectrogram(file_path):
    """Load a spectrogram image and crop to the spectrogram area."""
    image = Image.open(file_path).convert('L')
    # Crop to spectrogram area (removing axes, colorbar, etc.)
    image = image.crop((80, 80, 800, 350))
    return image

# Augmentation functions from our data preprocessing module
def time_shift(spectrogram, max_shift_percent=0.2):
    """Randomly shift the spectrogram in time."""
    width = spectrogram.shape[1]
    shift_amount = int(width * np.random.uniform(-max_shift_percent, max_shift_percent))
    
    shifted = np.zeros_like(spectrogram)
    if shift_amount > 0:
        shifted[:, shift_amount:] = spectrogram[:, :width-shift_amount]
    elif shift_amount < 0:
        shifted[:, :width+shift_amount] = spectrogram[:, -shift_amount:]
    else:
        shifted = spectrogram
        
    return shifted

def time_mask(spectrogram, max_mask_percent=0.2, num_masks=2):
    """Apply random masking in time dimension."""
    width = spectrogram.shape[1]
    masked = spectrogram.copy()
    
    for _ in range(num_masks):
        mask_width = int(width * np.random.uniform(0, max_mask_percent))
        mask_start = np.random.randint(0, width - mask_width)
        masked[:, mask_start:mask_start + mask_width] = 0
        
    return masked

def freq_mask(spectrogram, max_mask_percent=0.2, num_masks=2):
    """Apply random masking in frequency dimension."""
    height = spectrogram.shape[0]
    masked = spectrogram.copy()
    
    for _ in range(num_masks):
        mask_height = int(height * np.random.uniform(0, max_mask_percent))
        mask_start = np.random.randint(0, height - mask_height)
        masked[mask_start:mask_start + mask_height, :] = 0
        
    return masked

def freq_shift(spectrogram, max_shift_percent=0.2):
    """Randomly shift the spectrogram in frequency."""
    height = spectrogram.shape[0]
    shift_amount = int(height * np.random.uniform(-max_shift_percent, max_shift_percent))
    
    shifted = np.zeros_like(spectrogram)
    if shift_amount > 0:
        shifted[shift_amount:, :] = spectrogram[:height-shift_amount, :]
    elif shift_amount < 0:
        shifted[:height+shift_amount, :] = spectrogram[-shift_amount:, :]
    else:
        shifted = spectrogram
        
    return shifted

def amplitude_scale(spectrogram, min_factor=0.5, max_factor=1.5):
    """Randomly scale the amplitude of the spectrogram."""
    scale_factor = np.random.uniform(min_factor, max_factor)
    return spectrogram * scale_factor

def add_gaussian_noise(spectrogram, max_noise_percent=0.1):
    """Add random Gaussian noise to the spectrogram."""
    noise_level = np.random.uniform(0, max_noise_percent)
    noise = np.random.normal(0, noise_level * np.mean(spectrogram), spectrogram.shape)
    return spectrogram + noise

def apply_augmentations(spectrogram, p=0.5):
    """Apply a random combination of augmentations to the spectrogram."""
    augmented = spectrogram.copy()
    
    # Time-domain augmentations
    if np.random.random() < p:
        augmented = time_shift(augmented)
    if np.random.random() < p:
        augmented = time_mask(augmented)
        
    # Frequency-domain augmentations
    if np.random.random() < p:
        augmented = freq_shift(augmented)
    if np.random.random() < p:
        augmented = freq_mask(augmented)
        
    # Intensity augmentations
    if np.random.random() < p:
        augmented = amplitude_scale(augmented)
    if np.random.random() < p:
        augmented = add_gaussian_noise(augmented)
        
    return augmented

# Function to demonstrate SimCLR augmentation pairs
def demonstrate_simclr_pairs(file_path, num_pairs=5):
    """Generate multiple SimCLR augmentation pairs from a single spectrogram."""
    # Load and convert image to numpy array
    image = load_spectrogram(file_path)
    spectrogram = np.array(image)
    
    # Create figure for visualization
    plt.figure(figsize=(15, num_pairs * 4))
    
    # Original spectrogram
    plt.subplot(num_pairs + 1, 2, 1)
    plt.imshow(spectrogram, aspect='auto', cmap='viridis')
    plt.title('Original Spectrogram')
    plt.colorbar()
    
    plt.subplot(num_pairs + 1, 2, 2)
    plt.imshow(spectrogram, aspect='auto', cmap='viridis')
    plt.title('Original Spectrogram')
    plt.colorbar()
    
    # Generate augmentation pairs
    for i in range(num_pairs):
        # Create two different augmentations of the same spectrogram
        augmented1 = apply_augmentations(spectrogram, p=0.8)
        augmented2 = apply_augmentations(spectrogram, p=0.8)
        
        # Plot the pair
        plt.subplot(num_pairs + 1, 2, 2*i + 3)
        plt.imshow(augmented1, aspect='auto', cmap='viridis')
        plt.title(f'Augmented View 1 - Pair {i+1}')
        plt.colorbar()
        
        plt.subplot(num_pairs + 1, 2, 2*i + 4)
        plt.imshow(augmented2, aspect='auto', cmap='viridis')
        plt.title(f'Augmented View 2 - Pair {i+1}')
        plt.colorbar()
    
    plt.tight_layout()
    
    # Extract category from filename for saving
    filename = os.path.basename(file_path)
    category = "unknown"
    if filename.startswith('env_noise'):
        category = 'env_noise'
    elif 'bio_signal_whale' in filename:
        category = 'bio_whale'
    elif 'bio_signal_fish' in filename:
        category = 'bio_fish'
    elif 'bio_signal_coral' in filename:
        category = 'bio_coral'
    elif 'manmade_boat' in filename:
        category = 'manmade_boat'
    elif 'manmade_ship' in filename:
        category = 'manmade_ship'
    elif 'manmade_submarine' in filename:
        category = 'manmade_submarine'
    elif 'manmade_speedboat' in filename:
        category = 'manmade_speedboat'
    elif 'transient' in filename:
        category = 'transient'
    
    # Save the figure
    plt.savefig(f'/home/ubuntu/augmentation_demo/simclr_pairs_{category}.png')
    plt.close()
    
    return f'/home/ubuntu/augmentation_demo/simclr_pairs_{category}.png'

# Function to demonstrate individual augmentations
def demonstrate_individual_augmentations(file_path):
    """Show the effect of each individual augmentation type."""
    # Load and convert image to numpy array
    image = load_spectrogram(file_path)
    spectrogram = np.array(image)
    
    # Create figure for visualization
    plt.figure(figsize=(15, 15))
    
    # Original spectrogram
    plt.subplot(4, 2, 1)
    plt.imshow(spectrogram, aspect='auto', cmap='viridis')
    plt.title('Original Spectrogram')
    plt.colorbar()
    
    # Time shift
    plt.subplot(4, 2, 2)
    shifted = time_shift(spectrogram, max_shift_percent=0.3)
    plt.imshow(shifted, aspect='auto', cmap='viridis')
    plt.title('Time Shift')
    plt.colorbar()
    
    # Time mask
    plt.subplot(4, 2, 3)
    time_masked = time_mask(spectrogram, max_mask_percent=0.2, num_masks=3)
    plt.imshow(time_masked, aspect='auto', cmap='viridis')
    plt.title('Time Masking')
    plt.colorbar()
    
    # Frequency mask
    plt.subplot(4, 2, 4)
    freq_masked = freq_mask(spectrogram, max_mask_percent=0.2, num_masks=3)
    plt.imshow(freq_masked, aspect='auto', cmap='viridis')
    plt.title('Frequency Masking')
    plt.colorbar()
    
    # Frequency shift
    plt.subplot(4, 2, 5)
    freq_shifted = freq_shift(spectrogram, max_shift_percent=0.3)
    plt.imshow(freq_shifted, aspect='auto', cmap='viridis')
    plt.title('Frequency Shift')
    plt.colorbar()
    
    # Amplitude scaling
    plt.subplot(4, 2, 6)
    scaled = amplitude_scale(spectrogram, min_factor=0.3, max_factor=1.8)
    plt.imshow(scaled, aspect='auto', cmap='viridis')
    plt.title('Amplitude Scaling')
    plt.colorbar()
    
    # Gaussian noise
    plt.subplot(4, 2, 7)
    noisy = add_gaussian_noise(spectrogram, max_noise_percent=0.2)
    plt.imshow(noisy, aspect='auto', cmap='viridis')
    plt.title('Gaussian Noise')
    plt.colorbar()
    
    # Combined augmentations
    plt.subplot(4, 2, 8)
    combined = apply_augmentations(spectrogram, p=1.0)  # Apply all augmentations
    plt.imshow(combined, aspect='auto', cmap='viridis')
    plt.title('Combined Augmentations')
    plt.colorbar()
    
    plt.tight_layout()
    
    # Extract category from filename for saving
    filename = os.path.basename(file_path)
    category = "unknown"
    if filename.startswith('env_noise'):
        category = 'env_noise'
    elif 'bio_signal_whale' in filename:
        category = 'bio_whale'
    elif 'bio_signal_fish' in filename:
        category = 'bio_fish'
    elif 'bio_signal_coral' in filename:
        category = 'bio_coral'
    elif 'manmade_boat' in filename:
        category = 'manmade_boat'
    elif 'manmade_ship' in filename:
        category = 'manmade_ship'
    elif 'manmade_submarine' in filename:
        category = 'manmade_submarine'
    elif 'manmade_speedboat' in filename:
        category = 'manmade_speedboat'
    elif 'transient' in filename:
        category = 'transient'
    
    # Save the figure
    plt.savefig(f'/home/ubuntu/augmentation_demo/individual_augmentations_{category}.png')
    plt.close()
    
    return f'/home/ubuntu/augmentation_demo/individual_augmentations_{category}.png'

# Function to demonstrate augmentations across different signal types
def demonstrate_augmentations_across_categories(data_dir):
    """Show how augmentations affect different types of underwater acoustic signals."""
    # Get all spectrogram files
    all_files = [f for f in os.listdir(data_dir) if f.endswith('.png')]
    
    # Categorize files
    categories = {
        'env_noise': [],
        'bio_whale': [],
        'bio_fish': [],
        'bio_coral': [],
        'manmade_boat': [],
        'manmade_ship': [],
        'manmade_submarine': [],
        'manmade_speedboat': [],
        'transient': []
    }
    
    for file in all_files:
        if file.startswith('env_noise'):
            categories['env_noise'].append(file)
        elif 'bio_signal_whale' in file:
            categories['bio_whale'].append(file)
        elif 'bio_signal_fish' in file:
            categories['bio_fish'].append(file)
        elif 'bio_signal_coral' in file:
            categories['bio_coral'].append(file)
        elif 'manmade_boat' in file:
            categories['manmade_boat'].append(file)
        elif 'manmade_ship' in file:
            categories['manmade_ship'].append(file)
        elif 'manmade_submarine' in file:
            categories['manmade_submarine'].append(file)
        elif 'manmade_speedboat' in file:
            categories['manmade_speedboat'].append(file)
        elif 'transient' in file:
            categories['transient'].append(file)
    
    # Sample one file from each non-empty category
    sample_files = []
    for category, files in categories.items():
        if files:  # If category is not empty
            sample_files.append(os.path.join(data_dir, random.choice(files)))
    
    # Demonstrate augmentations for each sample
    simclr_demo_files = []
    individual_demo_files = []
    
    for file_path in sample_files:
        simclr_demo_file = demonstrate_simclr_pairs(file_path, num_pairs=3)
        individual_demo_file = demonstrate_individual_augmentations(file_path)
        
        simclr_demo_files.append(simclr_demo_file)
        individual_demo_files.append(individual_demo_file)
    
    return simclr_demo_files, individual_demo_files

# Main function to run the demonstration
def main():
    data_dir = '/home/ubuntu/data'
    
    print("Demonstrating SimCLR augmentations for underwater acoustic spectrograms...")
    simclr_demo_files, individual_demo_files = demonstrate_augmentations_across_categories(data_dir)
    
    print("SimCLR augmentation pair demonstrations saved to:")
    for file in simclr_demo_files:
        print(f"  - {file}")
    
    print("\nIndividual augmentation demonstrations saved to:")
    for file in individual_demo_files:
        print(f"  - {file}")
    
    print("\nDemonstration complete!")

if __name__ == "__main__":
    main()