InteractiveMachineLearning/InteractiveNaiveBayes.py at master · sheepsy90/InteractiveMachineLearning · 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
import pygame
import math
from algorithms.NaiveBayes import NaiveBayes

#defining color variables
BLACK = (0, 0, 0)
WHITE = (190, 190, 190)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0,255)
YELLOW = (255, 255, 0)

colour_by_class = {
    0: RED,
    1: GREEN,
    2: BLUE,
    3: YELLOW
}

#window settings
size = (800, 800)
pygame.init()
screen = pygame.display.set_mode(size)
pygame.display.set_caption("Naive Bayes")
font = pygame.font.Font(None, 36)

#setting fps variable
clock = pygame.time.Clock()

# State Variables
done = False


naive_bayes_storage = NaiveBayes()
draw_area = False
draw_cursor = False
probabilistic_switch = False

while not done:
    pos = pygame.mouse.get_pos()

    for event in pygame.event.get():
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                done = True
            if event.key == pygame.K_c:
                naive_bayes_storage.clear()
            if event.key == pygame.K_1:
                naive_bayes_storage.add_data_point(0, pos)
                naive_bayes_storage.build_model()
                break
            if event.key == pygame.K_2:
                naive_bayes_storage.add_data_point(1, pos)
                naive_bayes_storage.build_model()
                break
            if event.key == pygame.K_3:
                naive_bayes_storage.add_data_point(2, pos)
                naive_bayes_storage.build_model()
                break
            if event.key == pygame.K_4:
                naive_bayes_storage.add_data_point(3, pos)
                naive_bayes_storage.build_model()
                break

            if event.key == pygame.K_m:
                break

            if event.key == pygame.K_a:
                if naive_bayes_storage.has_model():
                    draw_area = not draw_area
                break
            if event.key == pygame.K_l:
                if naive_bayes_storage.has_model():
                    draw_cursor = not draw_cursor
                break
            if event.key == pygame.K_p:
                if naive_bayes_storage.has_model():
                    probabilistic_switch = not probabilistic_switch
                break

    screen.fill(BLACK)

    if draw_area:
        for i in range(0, size[0], 20):
            for j in range(0, size[0], 20):
                x,y = i+10, j+10
                result = naive_bayes_storage.predict((i, j), probabilistic=probabilistic_switch)

                if result is not None and not math.isnan(result):
                    pygame.draw.circle(screen, colour_by_class[result], (i, j), 1)


    data_points = naive_bayes_storage.get_data_points()
    # First draw all points on screen
    for pnt in data_points:
        pygame.draw.circle(screen, colour_by_class[pnt[0]], pnt[1], 5)

    text = font.render(str("FPS: " + str(int(clock.get_fps()))), True, WHITE)
    screen.blit(text, (700,10))

    if draw_cursor:
        # Estimate the mouse cursors colour
        result = naive_bayes_storage.predict(pos)

        if result is not None:
            pygame.draw.circle(screen, colour_by_class[result], pos, 8)

    pygame.display.flip()
    clock.tick(60)

pygame.quit()