Create utils/visualization.py

utils/visualization.py

@@ -0,0 +1,356 @@
+"""
+Visualization utilities for Kolam images and training progress.
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import seaborn as sns
+from pathlib import Path
+from typing import List, Optional, Tuple
+import cv2
+
+
+def plot_kolam_grid(images: List[np.ndarray], grid_size: Tuple[int, int] = (4, 4), 
+                   save_path: Optional[str] = None, title: str = "Kolam Designs",
+                   figsize: Tuple[int, int] = (12, 12)) -> None:
+    """Plot a grid of Kolam images."""
+    rows, cols = grid_size
+    fig, axes = plt.subplots(rows, cols, figsize=figsize)
+    
+    # Flatten axes for easier indexing
+    if rows == 1:
+        axes = [axes]
+    elif cols == 1:
+        axes = [[ax] for ax in axes]
+    else:
+        axes = axes.flatten()
+    
+    for i, ax in enumerate(axes):
+        if i < len(images):
+            # Ensure image is in correct format
+            img = images[i]
+            if isinstance(img, torch.Tensor):
+                img = img.detach().cpu().numpy()
+            
+            # Handle different image formats
+            if len(img.shape) == 3 and img.shape[0] == 1:
+                img = img.squeeze(0)
+            elif len(img.shape) == 3 and img.shape[2] == 1:
+                img = img.squeeze(2)
+            
+            # Normalize to [0, 1] if needed
+            if img.max() > 1.0:
+                img = img / 255.0
+            
+            ax.imshow(img, cmap='gray', vmin=0, vmax=1)
+        else:
+            ax.axis('off')
+        
+        ax.set_xticks([])
+        ax.set_yticks([])
+        ax.set_aspect('equal')
+    
+    plt.suptitle(title, fontsize=16, fontweight='bold')
+    plt.tight_layout()
+    
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        print(f"Grid saved to {save_path}")
+    
+    plt.show()
+
+
+def plot_training_curves(train_losses: List[float], val_losses: List[float] = None,
+                        train_accs: List[float] = None, val_accs: List[float] = None,
+                        save_path: Optional[str] = None, title: str = "Training Progress") -> None:
+    """Plot training curves for loss and accuracy."""
+    fig, axes = plt.subplots(1, 2, figsize=(15, 5))
+    
+    # Loss curves
+    axes[0].plot(train_losses, label='Train Loss', color='blue', alpha=0.7)
+    if val_losses:
+        axes[0].plot(val_losses, label='Validation Loss', color='red', alpha=0.7)
+    axes[0].set_xlabel('Epoch')
+    axes[0].set_ylabel('Loss')
+    axes[0].set_title('Training and Validation Loss')
+    axes[0].legend()
+    axes[0].grid(True, alpha=0.3)
+    
+    # Accuracy curves
+    if train_accs is not None:
+        axes[1].plot(train_accs, label='Train Accuracy', color='blue', alpha=0.7)
+        if val_accs:
+            axes[1].plot(val_accs, label='Validation Accuracy', color='red', alpha=0.7)
+        axes[1].set_xlabel('Epoch')
+        axes[1].set_ylabel('Accuracy (%)')
+        axes[1].set_title('Training and Validation Accuracy')
+        axes[1].legend()
+        axes[1].grid(True, alpha=0.3)
+    else:
+        axes[1].axis('off')
+    
+    plt.suptitle(title, fontsize=16, fontweight='bold')
+    plt.tight_layout()
+    
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        print(f"Training curves saved to {save_path}")
+    
+    plt.show()
+
+
+def plot_gan_training_curves(g_losses: List[float], d_losses: List[float],
+                            real_scores: List[float] = None, fake_scores: List[float] = None,
+                            save_path: Optional[str] = None) -> None:
+    """Plot GAN training curves."""
+    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
+    
+    # Generator and Discriminator losses
+    axes[0, 0].plot(g_losses, label='Generator Loss', color='blue', alpha=0.7)
+    axes[0, 0].plot(d_losses, label='Discriminator Loss', color='red', alpha=0.7)
+    axes[0, 0].set_xlabel('Epoch')
+    axes[0, 0].set_ylabel('Loss')
+    axes[0, 0].set_title('Generator and Discriminator Loss')
+    axes[0, 0].legend()
+    axes[0, 0].grid(True, alpha=0.3)
+    
+    # Discriminator scores
+    if real_scores and fake_scores:
+        axes[0, 1].plot(real_scores, label='Real Score', color='green', alpha=0.7)
+        axes[0, 1].plot(fake_scores, label='Fake Score', color='orange', alpha=0.7)
+        axes[0, 1].set_xlabel('Epoch')
+        axes[0, 1].set_ylabel('Score')
+        axes[0, 1].set_title('Discriminator Scores')
+        axes[0, 1].legend()
+        axes[0, 1].grid(True, alpha=0.3)
+    else:
+        axes[0, 1].axis('off')
+    
+    # Generator loss only
+    axes[1, 0].plot(g_losses, color='blue', alpha=0.7)
+    axes[1, 0].set_xlabel('Epoch')
+    axes[1, 0].set_ylabel('Generator Loss')
+    axes[1, 0].set_title('Generator Loss')
+    axes[1, 0].grid(True, alpha=0.3)
+    
+    # Discriminator loss only
+    axes[1, 1].plot(d_losses, color='red', alpha=0.7)
+    axes[1, 1].set_xlabel('Epoch')
+    axes[1, 1].set_ylabel('Discriminator Loss')
+    axes[1, 1].set_title('Discriminator Loss')
+    axes[1, 1].grid(True, alpha=0.3)
+    
+    plt.suptitle('GAN Training Progress', fontsize=16, fontweight='bold')
+    plt.tight_layout()
+    
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        print(f"GAN training curves saved to {save_path}")
+    
+    plt.show()
+
+
+def plot_image_comparison(original: np.ndarray, generated: np.ndarray,
+                         save_path: Optional[str] = None, title: str = "Original vs Generated") -> None:
+    """Plot comparison between original and generated images."""
+    fig, axes = plt.subplots(1, 2, figsize=(10, 5))
+    
+    # Original image
+    axes[0].imshow(original, cmap='gray', vmin=0, vmax=1)
+    axes[0].set_title('Original')
+    axes[0].axis('off')
+    
+    # Generated image
+    axes[1].imshow(generated, cmap='gray', vmin=0, vmax=1)
+    axes[1].set_title('Generated')
+    axes[1].axis('off')
+    
+    plt.suptitle(title, fontsize=16, fontweight='bold')
+    plt.tight_layout()
+    
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        print(f"Comparison saved to {save_path}")
+    
+    plt.show()
+
+
+def plot_feature_visualization(features: np.ndarray, save_path: Optional[str] = None) -> None:
+    """Visualize extracted features."""
+    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
+    
+    # Feature histogram
+    axes[0, 0].hist(features.flatten(), bins=50, alpha=0.7, color='blue')
+    axes[0, 0].set_title('Feature Distribution')
+    axes[0, 0].set_xlabel('Feature Value')
+    axes[0, 0].set_ylabel('Frequency')
+    axes[0, 0].grid(True, alpha=0.3)
+    
+    # Feature heatmap
+    if len(features.shape) == 2:
+        im = axes[0, 1].imshow(features, cmap='viridis', aspect='auto')
+        axes[0, 1].set_title('Feature Heatmap')
+        plt.colorbar(im, ax=axes[0, 1])
+    
+    # Feature statistics
+    stats_text = f"""
+    Mean: {features.mean():.4f}
+    Std: {features.std():.4f}
+    Min: {features.min():.4f}
+    Max: {features.max():.4f}
+    """
+    axes[1, 0].text(0.1, 0.5, stats_text, fontsize=12, verticalalignment='center')
+    axes[1, 0].set_title('Feature Statistics')
+    axes[1, 0].axis('off')
+    
+    # Feature correlation (if 2D)
+    if len(features.shape) == 2 and features.shape[1] > 1:
+        corr_matrix = np.corrcoef(features.T)
+        im = axes[1, 1].imshow(corr_matrix, cmap='coolwarm', vmin=-1, vmax=1)
+        axes[1, 1].set_title('Feature Correlation')
+        plt.colorbar(im, ax=axes[1, 1])
+    else:
+        axes[1, 1].axis('off')
+    
+    plt.suptitle('Feature Visualization', fontsize=16, fontweight='bold')
+    plt.tight_layout()
+    
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        print(f"Feature visualization saved to {save_path}")
+    
+    plt.show()
+
+
+def plot_symmetry_analysis(image: np.ndarray, save_path: Optional[str] = None) -> None:
+    """Analyze and visualize symmetry properties of an image."""
+    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
+    
+    # Original image
+    axes[0, 0].imshow(image, cmap='gray', vmin=0, vmax=1)
+    axes[0, 0].set_title('Original')
+    axes[0, 0].axis('off')
+    
+    # Horizontal flip
+    h_flipped = np.fliplr(image)
+    axes[0, 1].imshow(h_flipped, cmap='gray', vmin=0, vmax=1)
+    axes[0, 1].set_title('Horizontal Flip')
+    axes[0, 1].axis('off')
+    
+    # Vertical flip
+    v_flipped = np.flipud(image)
+    axes[0, 2].imshow(v_flipped, cmap='gray', vmin=0, vmax=1)
+    axes[0, 2].set_title('Vertical Flip')
+    axes[0, 2].axis('off')
+    
+    # Difference with horizontal flip
+    h_diff = np.abs(image - h_flipped)
+    axes[1, 0].imshow(h_diff, cmap='hot', vmin=0, vmax=1)
+    axes[1, 0].set_title(f'Horizontal Symmetry\n(Error: {h_diff.mean():.4f})')
+    axes[1, 0].axis('off')
+    
+    # Difference with vertical flip
+    v_diff = np.abs(image - v_flipped)
+    axes[1, 1].imshow(v_diff, cmap='hot', vmin=0, vmax=1)
+    axes[1, 1].set_title(f'Vertical Symmetry\n(Error: {v_diff.mean():.4f})')
+    axes[1, 1].axis('off')
+    
+    # Rotational symmetry
+    rotated = np.rot90(image, k=2)
+    r_diff = np.abs(image - rotated)
+    axes[1, 2].imshow(r_diff, cmap='hot', vmin=0, vmax=1)
+    axes[1, 2].set_title(f'Rotational Symmetry\n(Error: {r_diff.mean():.4f})')
+    axes[1, 2].axis('off')
+    
+    plt.suptitle('Symmetry Analysis', fontsize=16, fontweight='bold')
+    plt.tight_layout()
+    
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        print(f"Symmetry analysis saved to {save_path}")
+    
+    plt.show()
+
+
+def plot_metrics_comparison(metrics_dict: dict, save_path: Optional[str] = None) -> None:
+    """Plot comparison of different metrics."""
+    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
+    
+    # Extract metric names and values
+    metric_names = list(metrics_dict.keys())
+    metric_values = list(metrics_dict.values())
+    
+    # Bar plot of metrics
+    axes[0, 0].bar(metric_names, metric_values, alpha=0.7, color='skyblue')
+    axes[0, 0].set_title('Metrics Comparison')
+    axes[0, 0].set_ylabel('Value')
+    axes[0, 0].tick_params(axis='x', rotation=45)
+    axes[0, 0].grid(True, alpha=0.3)
+    
+    # Pie chart of relative importance
+    axes[0, 1].pie(metric_values, labels=metric_names, autopct='%1.1f%%', startangle=90)
+    axes[0, 1].set_title('Relative Metric Importance')
+    
+    # Line plot (if metrics are time series)
+    if len(metric_values) > 1:
+        axes[1, 0].plot(metric_names, metric_values, marker='o', linewidth=2, markersize=8)
+        axes[1, 0].set_title('Metrics Trend')
+        axes[1, 0].set_ylabel('Value')
+        axes[1, 0].tick_params(axis='x', rotation=45)
+        axes[1, 0].grid(True, alpha=0.3)
+    
+    # Summary statistics
+    stats_text = f"""
+    Best Metric: {metric_names[np.argmax(metric_values)]}
+    Worst Metric: {metric_names[np.argmin(metric_values)]}
+    Average: {np.mean(metric_values):.4f}
+    Std Dev: {np.std(metric_values):.4f}
+    """
+    axes[1, 1].text(0.1, 0.5, stats_text, fontsize=12, verticalalignment='center')
+    axes[1, 1].set_title('Summary Statistics')
+    axes[1, 1].axis('off')
+    
+    plt.suptitle('Metrics Analysis', fontsize=16, fontweight='bold')
+    plt.tight_layout()
+    
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        print(f"Metrics comparison saved to {save_path}")
+    
+    plt.show()
+
+
+def create_animation_frames(images: List[np.ndarray], save_dir: str) -> None:
+    """Create frames for animation from a list of images."""
+    save_dir = Path(save_dir)
+    save_dir.mkdir(parents=True, exist_ok=True)
+    
+    for i, img in enumerate(images):
+        plt.figure(figsize=(8, 8))
+        plt.imshow(img, cmap='gray', vmin=0, vmax=1)
+        plt.title(f'Frame {i+1}')
+        plt.axis('off')
+        plt.tight_layout()
+        plt.savefig(save_dir / f'frame_{i:04d}.png', dpi=150, bbox_inches='tight')
+        plt.close()
+    
+    print(f"Animation frames saved to {save_dir}")
+
+
+if __name__ == "__main__":
+    # Test visualization functions
+    print("Testing visualization functions...")
+    
+    # Create sample images
+    sample_images = [np.random.rand(64, 64) for _ in range(16)]
+    
+    # Test grid plotting
+    plot_kolam_grid(sample_images, title="Test Grid")
+    
+    # Test training curves
+    train_losses = [1.0 - i * 0.01 for i in range(100)]
+    val_losses = [1.1 - i * 0.009 for i in range(100)]
+    plot_training_curves(train_losses, val_losses, title="Test Training Curves")
+    
+    print("Visualization tests completed!")