NVisual/src/visualizer.hpp at main · TonSHd/NVisual · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#pragma once
#include "nn.hpp"
#include <ncurses.h>
#include <string>
#include <vector>

enum ColorPair {
    CP_DEFAULT  = 1,
    CP_TITLE    = 2,
    CP_ACTIVE   = 3,
    CP_DIM      = 4,
    CP_BALL     = 5,
    CP_EDGE     = 6,
    CP_ANSWER   = 7,
    CP_WARN     = 8,
    CP_BAR      = 9,
    CP_PROMPT   = 10,
    CP_LOSS     = 11,
    CP_LABEL    = 12,
};

struct BallPos {
    float x = 0.f;  // 0..1 between source and target column
    float y = 0.f;  // 0..1 between source and target neuron
    int   src_layer  = 0;
    int   src_neuron = 0;
    int   dst_layer  = 0;
    int   dst_neuron = 0;
};

class Visualizer {
public:
    Visualizer();
    ~Visualizer();

    void init();
    void cleanup();

    void draw_inference(const NeuralNet& nn,
                        const std::string& prompt,
                        const std::string& answer,
                        float confidence,
                        bool  warn,
                        const BallPos& ball,
                        int   active_layer);  // -1 = all dim

    void draw_training(const NeuralNet& nn,
                       int   epoch,
                       int   total_epochs,
                       float loss,
                       float accuracy,
                       int   dataset_size,
                       const std::string& status_msg);

    void draw_splash(const NeuralNet& nn, int dataset_size);

    void animate_ball(const NeuralNet& nn,
                      const std::string& prompt,
                      const std::string& answer,
                      float confidence,
                      bool  warn,
                      int   src_layer,
                      int   dst_layer,
                      int   frames = 20,
                      int   delay_ms = 30);

    int  get_key();
    void sleep_ms(int ms);
    bool is_quit_key(int ch);

private:
    bool initialized_ = false;
    int  rows_ = 0, cols_ = 0;

    int net_top_    = 0;
    int net_bottom_ = 0;
    int info_top_   = 0;
    int bar_y_      = 0;

    static constexpr int MAX_VIS_NEURONS = 8;

    std::vector<int> layer_col_x_;

    struct AmbientBubble {
        float x;
        float y;
        float vy;
        float r;
    };
    mutable std::vector<AmbientBubble> bubbles_;
    mutable float animation_time_ = 0.f;

    void setup_colors();
    void compute_layout();

    void draw_border();
    void draw_title_bar(const std::string& title);
    void draw_status_bar(const std::string& msg);

    void draw_network_panel(const NeuralNet& nn,
                             int   active_layer,
                             const BallPos& ball,
                             bool  ball_visible);

    void draw_layer(int col_x, int layer_idx,
                    const std::vector<float>& activations,
                    const std::string& label,
                    bool active);

    void draw_edges(int col_x1, int col_x2,
                    int n1, int n2,
                    const std::vector<float>& act1,
                    const std::vector<float>& act2);

    void draw_ball(const BallPos& ball, bool visible);

    void draw_info_panel(const std::string& prompt,
                          const std::string& answer,
                          float confidence,
                          const std::vector<std::pair<float,std::string>>& top_preds,
                          bool warn);

    void draw_training_panel(int epoch, int total,
                              float loss, float acc,
                              int dataset_size,
                              const std::string& msg);

    int  neuron_y(int layer_center_y, int neuron_idx, int total) const;
    int  neuron_y_abs(int layer_idx, int neuron_idx, int total) const;
    void attron_pair(ColorPair cp, bool bold = false);
    void attroff_pair(ColorPair cp, bool bold = false);
    std::string truncate(const std::string& s, int max_w) const;

    int  layer_center_y(int layer_idx, int n_neurons) const;
};