# -rw-r--r-- 1.9 KiB View raw 
                                                                                
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
import random
import math
import sys

from PIL import Image, ImageDraw
import numpy as np

xLim = 28
yLim = 28
radius = 4

generatePoint = lambda: (random.random()*(xLim - 2*radius) + radius, random.random()*(yLim - 2*radius) + radius)
distance = lambda p1, p2: math.sqrt((p1[0]-p2[0])**2 + (p1[1]-p2[1])**2)

def genCircles(num):
    tries = 0
    circles = []
    while num > 0:
        # In pathalogical cases, this can get stuck. Limit the number
        # of tries, and retry.
        tries += 1
        if tries > 100:
            return genCircles(num + len(circles))

        center = generatePoint()
        if any(map(lambda circle: distance(center, circle) < 2*radius, circles)):
            continue
        num -= 1
        circles.append(center)
    return circles

def drawCircles(centerPoints):
    im = Image.new("L", (xLim, yLim))
    draw = ImageDraw.Draw(im)
    for point in centerPoints:
        topLeft = (point[0] - radius, point[1] - radius)
        bottomRight = (point[0] + radius, point[1] + radius)
        draw.ellipse([topLeft, bottomRight], fill=255)
    return im

trainSize = 60000
testSize = 10000

trainImages = np.empty((trainSize, 28, 28), dtype=np.uint8)
trainLabels = np.empty((trainSize,), dtype=np.uint8)

testImages = np.empty((testSize, 28, 28), dtype=np.uint8)
testLabels = np.empty((testSize,), dtype=np.uint8)

for i in range(trainSize):
    num = random.randrange(1, 7)
    circles = genCircles(num)
    circles = drawCircles(circles)
    circles = np.array(circles)
    trainImages[i] = circles
    trainLabels[i] = num

for i in range(testSize):
    num = random.randrange(1, 7)
    circles = genCircles(num)
    circles = drawCircles(circles)
    circles = np.array(circles)
    testImages[i] = circles
    testLabels[i] = num

np.save("dataset/trainData.npy", trainImages)
np.save("dataset/trainLabels.npy", trainLabels)
np.save("dataset/testData.npy", testImages)
np.save("dataset/testLabels.npy", testLabels)