forked from AniketTendulkar2510/IPARC_ML
-
Notifications
You must be signed in to change notification settings - Fork 0
/
GenerateCatB_Selection.py
180 lines (146 loc) · 6.84 KB
/
GenerateCatB_Selection.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
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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
import os
import numpy as np
import json
import pdb
from matplotlib import pyplot as plt
# from skimage.morphology import binary_dilation, binary_erosion
from scipy.ndimage import binary_dilation, binary_erosion, binary_hit_or_miss
import random
from ListSelEm import *
from Utils import Process, Change_Colour
def generate_inp_out_catB_Selection(list_se, **param):
"""
SE0/SE1 - Hit-Or-Miss
SE2/3 - Dilate (SE0)
SE2/3 - Erode (SE0)
SE4/5 - Dilate (SE1)
SE4/5 - Erode (SE1)
"""
sz = np.random.randint(2, 4)
# Select 1/2 pixels and dilate by SE0
base_img1 = np.zeros((param['img_size'], param['img_size']), dtype=np.int32)
idx1 = np.random.randint(0, param['img_size']//2, size=sz)
idx2 = np.random.randint(0, param['img_size']//2, size=sz)
base_img1[idx1, idx2] = 1
base_img1 = binary_dilation(base_img1, list_se_3x3[list_se[0]])
# Select 1/2 pixels and dilate by SE1
base_img2 = np.zeros((param['img_size'], param['img_size']), dtype=np.int32)
idx1 = np.random.randint(param['img_size']//2, param['img_size'], size=sz)
idx2 = np.random.randint(param['img_size']//2, param['img_size'], size=sz)
base_img2[idx1, idx2] = 1
base_img2 = binary_dilation(base_img2, list_se_3x3[list_se[1]])
# Combine the above images to get the base image.
base_img = np.logical_or(base_img1, base_img2)*1
# Copy the base_img for input/output
inp_img = np.array(base_img*1, copy=True)
out_img = np.array(base_img*1, copy=True)
# Next we have a hit_or_miss which selects a pixel and adds another color
tmp_img = binary_hit_or_miss(out_img, list_se_3x3[list_se[0]])
out_img[tmp_img] = 2 # Add another color
out_img = Process(out_img, num_colors=2)
# First color will be processed differently
out_img[:, :, 0] = binary_dilation(out_img[:, :, 0], list_se_3x3[list_se[2]])
out_img[:, :, 0] = binary_dilation(out_img[:, :, 0], list_se_3x3[list_se[3]])
out_img[:, :, 0] = binary_erosion(out_img[:, :, 0], list_se_3x3[list_se[2]])
out_img[:, :, 0] = binary_erosion(out_img[:, :, 0], list_se_3x3[list_se[3]])
# Second color will be processed differently
out_img[:, :, 1] = binary_dilation(out_img[:, :, 1], list_se_3x3[list_se[0]])
out_img[:, :, 1] = binary_dilation(out_img[:, :, 1], list_se_3x3[list_se[4]])
out_img[:, :, 1] = binary_dilation(out_img[:, :, 1], list_se_3x3[list_se[5]])
out_img[:, :, 1] = binary_erosion(out_img[:, :, 1], list_se_3x3[list_se[4]])
out_img[:, :, 1] = binary_erosion(out_img[:, :, 1], list_se_3x3[list_se[5]])
# Resolve the color by the rule
rule = np.array([[0, 0, 0], [0, 1, 2], [1, 0, 1], [1, 1, 2]], dtype=np.int32)
out_img = Change_Colour(out_img, rule)
return inp_img, out_img
def generate_one_task_CatB_Selection(**param):
"""
"""
k_example = 0
list_se_idx = np.random.randint(0, 8, size=6)
data = []
while k_example < param['no_examples_per_task']:
inp_img, out_img = generate_inp_out_catB_Selection(list_se_idx, **param)
# Check if both input and output images are non-trivial
FLAG = False
if np.all(inp_img*1 == 1) or np.all(inp_img*1 == 0):
FLAG = True
elif np.all(out_img*1 == 1) or np.all(out_img*1 == 0):
FLAG = True
if FLAG:
# If trivial regenerate the list of se's
# And reset all variables!!
data = []
list_se_idx = np.random.randint(0, 8, size=6)
k_example = -1
else:
data.append((inp_img, out_img))
# Increment k_example
k_example += 1
return data, list_se_idx
def write_dict_json_CatB_Selection(data, fname):
"""
"""
dict_data = []
for (inp, out) in data:
inp = [[int(y) for y in x] for x in inp]
out = [[int(y) for y in x] for x in out]
dict_data.append({"input": inp, "output": out})
with open(fname, "w") as f:
f.write(json.dumps(dict_data))
def write_solution_CatB_Selection(list_se_idx, fname):
"""
"""
color_rule = np.array([[0, 0, 0], [0, 1, 2], [1, 0, 1], [1, 1, 2]], dtype=np.int32)
with open(fname, 'w') as f:
f.write("Hit-Or-Miss SE{} \n".format(list_se_idx[0]))
f.write("Band 1 - Dilation SE{} \n".format(list_se_idx[2]+1))
f.write("Band 1 - Dilation SE{} \n".format(list_se_idx[3]+1))
f.write("Band 1 - Erosion SE{} \n".format(list_se_idx[2]+1))
f.write("Band 1 - Erosion SE{} \n".format(list_se_idx[3]+1))
f.write("Band 2 - Dilation SE{} \n".format(list_se_idx[0]+1))
f.write("Band 2 - Dilation SE{} \n".format(list_se_idx[4]+1))
f.write("Band 2 - Dilation SE{} \n".format(list_se_idx[5]+1))
f.write("Band 2 - Erosion SE{} \n".format(list_se_idx[4]+1))
f.write("Band 2 - Erosion SE{} \n".format(list_se_idx[5]+1))
f.write("Color rule : {}".format(json.dumps([[int(y) for y in x] for x in color_rule])))
f.write("\n")
def write_solution_CatB_Selection_json(list_se_idx, fname):
"""
"""
color_rule = np.array([[0, 0, 0], [0, 1, 2], [1, 0, 1], [1, 1, 2]], dtype=np.int32)
data = []
data.append((None, "Hit-Or-Miss", "SE{}".format(list_se_idx[0]+1)))
data.append((1, "Dilation", "SE{}".format(list_se_idx[2]+1)))
data.append((1, "Dilation", "SE{}".format(list_se_idx[3]+1)))
data.append((1, "Erosion", "SE{}".format(list_se_idx[2]+1)))
data.append((1, "Erosion", "SE{}".format(list_se_idx[3]+1)))
data.append((2, "Dilation", "SE{}".format(list_se_idx[0]+1)))
data.append((2, "Dilation", "SE{}".format(list_se_idx[4]+1)))
data.append((2, "Dilation", "SE{}".format(list_se_idx[5]+1)))
data.append((2, "Erosion", "SE{}".format(list_se_idx[4]+1)))
data.append((2, "Erosion", "SE{}".format(list_se_idx[5]+1)))
data.append((None, "change_color", [[int(y) for y in x] for x in color_rule]))
with open(fname, "w") as f:
f.write(json.dumps(data))
def generate_100_tasks_CatB_Selection(seed, **param):
"""
"""
np.random.seed(seed)
os.makedirs("./Dataset/CatB_Selection", exist_ok=True)
for task_no in range(100):
data, list_se_idx = generate_one_task_CatB_Selection(**param)
fname = './Dataset/CatB_Selection/Task{:03d}.json'.format(task_no)
write_dict_json_CatB_Selection(data, fname)
fname = './Dataset/CatB_Selection/Task{:03d}_soln.txt'.format(task_no)
write_solution_CatB_Selection(list_se_idx, fname)
fname = './Dataset/CatB_Selection/Task{:03d}_soln.json'.format(task_no)
write_solution_CatB_Selection_json(list_se_idx, fname)
if __name__ == "__main__":
param = {}
param['img_size'] = 15
param['se_size'] = 3 # Size of the structuring element
param['seq_length'] = 4 # Number of primitives would be 2*param['seq_length']
param['no_examples_per_task'] = 4
param['no_colors'] = 3
generate_100_tasks_CatB_Selection(32, **param)