draw/create_data.py at master · numberlearning/draw · 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
import sys
import numpy as np
import random
min_edge = 2
max_edge = 5
min_blobs = 0
max_blobs = 9
n_labels = max_blobs - min_blobs + 1

c = min_blobs
total = 0
while (c < max_blobs + 1):
   if c == 0:
       total = total + 1000
       c = c + 1
       continue
   cN = int(10000/c)
   total = total + cN
   c = c + 1


def set_size(num_blobs, even=None):
    if even is None:
        if num_blobs == 0:
            return 1000
        else:
            return int(10000/num_blobs)
    else:
        return 1000

def generate_data(even=None):
    num_blobs = min_blobs
    i = 0
    if even is None:
        train = np.zeros([total, 10000])
        label = np.zeros([total, n_labels])
    else:
        train = np.zeros([1000 * n_labels, 10000])
        label = np.zeros([1000 * n_labels, n_labels])

    while (num_blobs < max_blobs + 1):

        nOfItem = set_size(num_blobs, even)

        if num_blobs == 0:
           for n in range(nOfItem):
               train[i] = np.zeros(10000)
               label[i, num_blobs] = 1
               i = i +1
           num_blobs = num_blobs + 1
           continue

        for n in range(nOfItem):
            a = np.zeros(10000)
            count = 0
            used = np.zeros((num_blobs, 4))

            while count < num_blobs:
                height = random.randint(min_edge, max_edge)
                width = random.randint(min_edge, max_edge)
                cX = random.randint(1, 99-width)
                cY = random.randint(1, 99-height)
                h = height
                w = width

                index = 0

                while index < count:
                    if cX+width+1 <= used[index, 0] or used[index, 0]+1+used[index, 2] <= cX or used[index, 1]+1+used[index,3] <= cY or cY+height+1<=used[index,1]:
                        index = index + 1
                    else:
                        cX = random.randint(1, 99-width)
                        cY = random.randint(1, 99-height)
                        index = 0

                used[index, 0] = cX
                used[index, 1] = cY
                used[index, 2] = width
                used[index, 3] = height

                for p in range(cX, cX+width):
                    for q in range(cY, cY+height):
                        a[p*100+q] = 255
                count = count + 1

            train[i] = a
            if min_blobs == 0:
                label[i, num_blobs] = 1
            elif min_blobs == 1:
                label[i, num_blobs - 1] = 1
            else:
                print("Hey, you! Min blobs should be one or zero!!!!")
            i = i + 1

        num_blobs = num_blobs + 1

    np.set_printoptions(threshold=np.nan)

#    if empty_img:
#        train = np.vstack((train, np.zeros((1000, 10000))))
#        for empty_idx in range(1000):
#            empty_label = np.zeros(max_blobs + 1)
#            empty_label[0] = 1
#            label = np.vstack((label, empty_label))

    return train, label