-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdata_gen.py
67 lines (61 loc) · 2.21 KB
/
data_gen.py
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
import numpy as np
#import matplotlib.pyplot as plt
import os
import shutil
import scipy.misc
#import matplotlib.cm as cm
from PIL import Image
image_size = 128
num_train_samples = 100000
num_valid_samples = 200
#os.mkdir('toy_data')
#os.mkdir('toy_dat/','images')
#os.mkdir('toy_data/annotations')
#os.mkdir('toy_data/images/training')
#os.mkdir('toy_data/images/validation')
#os.mkdir('toy_data/annotations/training')
#os.mkdir('toy_data/annotations/validation')
for i in range(num_train_samples):
pos = np.random.randint(0, image_size, 2)
im = np.zeros([image_size, image_size, 3], dtype=np.uint8)
label = np.random.randint(0, 2, 1)
for j in range(image_size):
for k in range(image_size):
if j <= pos[0] and k <= pos[1]:
im[j, k, 2*label] = 255
elif j > pos[0] and k <= pos[1]:
im[j, k, 2*(1-label)] = 255
elif j <= pos[0] and k > pos[1]:
im[j, k, 2*(1-label)] = 255
elif j > pos[0] and k > pos[1]:
im[j, k, 2*label] = 255
#plt.imshow(im)
im2 = Image.fromarray(im)
im2.save('toy_data/images/training/image_train_%07d.png' % i)
labels = im[:, :, 0]
labels[labels==255] = 1
im2 = Image.fromarray(labels)
im2.save('toy_data/annotations/training/image_train_%07d.png' % i)
#plt.imshow(labels)
for i in range(num_valid_samples):
pos = np.random.randint(0, image_size, 2)
im = np.zeros([image_size, image_size, 3], dtype=np.uint8)
label = np.random.randint(0, 2, 1)
for j in range(image_size):
for k in range(image_size):
if j <= pos[0] and k <= pos[1]:
im[j, k, 2*label] = 255
elif j > pos[0] and k <= pos[1]:
im[j, k, 2*(1-label)] = 255
elif j <= pos[0] and k > pos[1]:
im[j, k, 2*(1-label)] = 255
elif j > pos[0] and k > pos[1]:
im[j, k, 2*label] = 255
#plt.imshow(im)
im2 = Image.fromarray(im)
im2.save('toy_data/images/validation/image_val_%07d.png' % i)
labels = im[:, :, 0]
labels[labels==255] = 1
im2 = Image.fromarray(labels)
im2.save('toy_data/annotations/validation/image_val_%07d.png' % i)
#plt.imshow(labels)