visualizer.py

# Main Visualizer

import pygame
import numpy as np
from renderer import NetworkRenderer, StatsRenderer
from config import *

class NeuralNetworkVisualizer:
    def __init__(self, network):
        self.network = network
     
        # Initialize Pygame
        pygame.init()
        self.screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
        pygame.display.set_caption("Neural Network Visualizer - CRAZY MODE")
        self.clock = pygame.time.Clock()
     
        # Create Renderers
        self.network_renderer = NetworkRenderer(self.screen, network.layer_sizes)
        self.stats_renderer = StatsRenderer(self.screen)
     
        # State
        self.running = True
        self.paused = False
        self.current_epoch = 0
     
    def handle_events(self):
        # Handle Keyboard and mouse events
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    self.paused = not self.paused
                elif event.key == pygame.K_ESCAPE:
                    self.running = False
                elif event.key == pygame.K_s:
                    
                    # Save Model
                    self.network.save("model_visualized.pkl")
                    print("Model Saved")
     
    def render_frame(self, activations=None, epoch=0, loss=0, accuracy=0):
        # Render a single frame
        
        # Clear screen
        self.screen.fill(get_color("bg"))
     
        # Draw stats panel background
        self.stats_renderer.draw_background()
     
        # Draw Network
        self.network_renderer.draw_connections(activations, self.network.weights)
        self.network_renderer.draw_neurons(activations)
        self.network_renderer.draw_layer_labels()
     
        # Draw stats
        self.stats_renderer.draw_metrics(epoch, loss, accuracy, self.network.learning_rate)
        self.stats_renderer.draw_graphs()
     
        # Draw Title
        status = "PAUSED" if self.paused else "TRAINING"
        self.network_renderer.draw_title(f"Neural Network - {status}")
     
        # Update Display
        pygame.display.flip()
        self.clock.tick(FPS)
     
    def train_with_visualization(self, X_train, y_train, X_test, y_test,
                                 epochs=500, batch_size=32):
        
        # Convert labels to one hot
     
        train_activations, _ = self.network.forward(X_train[:100])
      
        if len(y_train.shape) == 1:
            y_one_hot = np.zeros((y_train.shape[0], self.network.layer_sizes[-1]))
            y_one_hot[np.arange(y_train.shape[0]), y_train] = 1
        else:
            y_one_hot = y_train
         
        n_samples = X_train.shape[0]
         
        for epoch in range(epochs):
            if not self.running:
                break
              
            self.current_epoch = epoch + 1
         
            self.handle_events()
         
            if self.paused:
                # Still render but don't train
                self.render_frame()
                continue
           
            # Shuffle data
            indices = np.random.permutation(n_samples)
            X_shuffled = X_train[indices]
            y_shuffled = y_one_hot[indices]
         
            # Mini Batch Training
            for i in range(0, n_samples, batch_size):
                X_batch = X_shuffled[i:i + batch_size]
                y_batch = y_shuffled[i:i + batch_size]
             
                # forward pass
                activations, z_values = self.network.forward(X_batch)
             
                # Backward pass
                weight_gradients, bias_gradients = self.network.backward(
                    X_batch, y_batch, activations, z_values)
             
                # Update weights
                for j in range(len(self.network.weights)):
                    self.network.weights[j] -= self.network.learning_rate * weight_gradients[j]
                    self.network.biases[j] -= self.network.learning_rate * bias_gradients[j]
             
            # Update visualization
            if (epoch + 1) % TRAIN_UPDATE_FREQUENCY == 0 or epoch == 0:
                
                # Calculate metrics on batch
                batch_activations, _ = self.network.forward(X_train[:100])
                loss = self.network._cross_entropy_loss(y_one_hot[:100], batch_activations[-1])
                accuracy = self.network.accuracy(X_test, y_test)

                # Single sample activations for visualization
                sample_idx = np.random.randint(min(100, len(X_train)))
                vis_activations, _ = self.network.forward(X_train[sample_idx:sample_idx+1])

                # Update stats
                self.stats_renderer.update_data(loss, accuracy)

                # Render
                self.render_frame(vis_activations, epoch + 1, loss, accuracy)
             
                print(f"Epoch {epoch + 1}/{epochs} - Loss: {loss:.4f} - Accuracy: {accuracy:.4f}")
                
                # Added delay to slow down visualization
                pygame.time.delay(500)
         
        print("\nTraining Complete")
     
        # Keep window open
        while self.running:
            self.handle_events()
            self.render_frame()
     
        pygame.quit()
     
    def visualize_prediction(self, X_sample, y_true=None, hold_time=2000):
        # Visualize a single prediction
        activations, _ = self.network.forward(X_sample.reshape(1, -1))
        prediction = np.argmax(activations[-1])
        confidence = np.max(activations[-1])
     
        # Render for hold_time milliseconds
        start_time = pygame.time.get_ticks()
        while pygame.time.get_ticks() - start_time < hold_time and self.running:
            self.handle_events()
            self.render_frame(activations)
         
            # Show prediction info
            text = f"Prediction: {prediction} ({confidence:.1%})"
            if y_true is not None:
                text += f" | True: {y_true}"
            
            # Draw prediction text
            font = pygame.font.Font(None, 32)
            pred_text = font.render(text, True, get_color("accent"))
            self.screen.blit(pred_text, (20, WINDOW_HEIGHT - 80))
            pygame.display.flip()