3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
6 # Written by Francois Fleuret <francois@fleuret.org>
10 import torch, torchvision
13 from torch.nn import functional as F
15 ######################################################################
18 colors = torch.tensor(
46 nb_bird_tokens = colors.size(0) - 1
47 token_forward = first_bird_token + nb_bird_tokens
48 token_backward = token_forward + 1
50 token2char = "_" + "".join([chr(ord("A") + n) for n in range(len(colors) - 1)]) + "><"
62 for _ in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"):
63 f_start = torch.zeros(height, width, dtype=torch.int64)
66 torch.empty(nb_birds, dtype=torch.int64),
67 torch.empty(nb_birds, dtype=torch.int64),
68 torch.empty(nb_birds, dtype=torch.int64),
69 torch.empty(nb_birds, dtype=torch.int64),
72 col = torch.randperm(colors.size(0) - 1)[:nb_birds].sort().values + 1
74 for n in range(nb_birds):
79 torch.randint(height, (1,))[0],
80 torch.randint(width, (1,))[0],
82 vm = torch.randint(4, (1,))[0]
83 vi[n], vj[n] = (vm % 2) * 2 - 1, (vm // 2) * 2 - 1
86 and i[n] - vi[n] < height
88 and j[n] - vj[n] < width
89 and f_start[i[n], j[n]] == 0
90 and f_start[i[n] - vi[n], j[n]] == 0
91 and f_start[i[n], j[n] - vj[n]] == 0
95 f_start[i[n], j[n]] = c
96 f_start[i[n] - vi[n], j[n]] = c
97 f_start[i[n], j[n] - vj[n]] = c
99 f_end = f_start.clone()
101 for l in range(nb_iterations):
102 for n in range(nb_birds):
104 f_end[i[n], j[n]] = 0
105 f_end[i[n] - vi[n], j[n]] = 0
106 f_end[i[n], j[n] - vj[n]] = 0
108 pi, pj, pvi, pvj = i[n].item(), j[n].item(), vi[n].item(), vj[n].item()
111 f_end[i[n], j[n]] == 0
112 and f_end[i[n] - vi[n], j[n]] == 0
113 and f_end[i[n], j[n] - vj[n]] == 0
116 if (i[n] == 0 and vi[n] == -1) or (i[n] == height - 1 and vi[n] == 1):
118 if (j[n] == 0 and vj[n] == -1) or (j[n] == width - 1 and vj[n] == 1):
125 f_end[i[n], j[n]] == 0
126 and f_end[i[n] - vi[n], j[n]] == 0
127 and f_end[i[n], j[n] - vj[n]] == 0
129 i[n], j[n], vi[n], vj[n] = pi, pj, pvi, pvj
131 f_end[i[n], j[n]] = c
132 f_end[i[n] - vi[n], j[n]] = c
133 f_end[i[n], j[n] - vj[n]] = c
135 pairs.append((f_start, f_end))
139 if torch.rand(1) < 0.5:
142 [p[0].flatten(), torch.tensor([token_forward]), p[1].flatten()],
149 [p[1].flatten(), torch.tensor([token_backward]), p[0].flatten()],
154 return torch.cat(result, dim=0)
166 for n in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"):
167 f_start = torch.zeros(height, width, dtype=torch.int64)
168 f_end = torch.zeros(height, width, dtype=torch.int64)
169 n = torch.arange(f_start.size(0))
172 (torch.randperm(nb_bird_tokens) + first_bird_token)[:nb_birds].sort().values
175 torch.randint(height - 2, (1,))[0] + 1,
176 torch.randint(width - 2, (1,))[0] + 1,
178 vm = torch.randint(4, (1,))[0]
179 vi, vj = (vm // 2) * (2 * (vm % 2) - 1), (1 - vm // 2) * (2 * (vm % 2) - 1)
182 f_start[i - vi, j - vj] = c
183 f_start[i + vj, j - vi] = c
184 f_start[i - vj, j + vi] = c
186 for l in range(nb_iterations):
189 if i < 0 or i >= height or j < 0 or j >= width:
197 f_end[i - vi, j - vj] = c
198 f_end[i + vj, j - vi] = c
199 f_end[i - vj, j + vi] = c
201 pairs.append((f_start, f_end))
205 if torch.rand(1) < 0.5:
208 [p[0].flatten(), torch.tensor([token_forward]), p[1].flatten()],
215 [p[1].flatten(), torch.tensor([token_backward]), p[0].flatten()],
220 return torch.cat(result, dim=0)
223 def sample2img(seq, height, width, upscale=15):
224 f_first = seq[:, : height * width].reshape(-1, height, width)
225 f_second = seq[:, height * width + 1 :].reshape(-1, height, width)
226 direction = seq[:, height * width]
228 def mosaic(x, upscale):
229 x = x.reshape(-1, height, width)
230 m = torch.logical_and(x >= 0, x < first_bird_token + nb_bird_tokens).long()
231 x = colors[x * m].permute(0, 3, 1, 2)
233 x = x[:, :, :, None, :, None].expand(-1, -1, -1, upscale, -1, upscale)
234 x = x.reshape(s[0], s[1], s[2] * upscale, s[3] * upscale)
236 x[:, :, :, torch.arange(0, x.size(3), upscale)] = 0
237 x[:, :, torch.arange(0, x.size(2), upscale), :] = 0
240 for n in range(m.size(0)):
241 for i in range(m.size(1)):
242 for j in range(m.size(2)):
244 for k in range(2, upscale - 2):
245 x[n, :, i * upscale + k, j * upscale + k] = 0
246 x[n, :, i * upscale + upscale - 1 - k, j * upscale + k] = 0
250 direction_symbol = torch.full((direction.size(0), height * upscale - 1, upscale), 0)
251 direction_symbol = colors[direction_symbol].permute(0, 3, 1, 2)
252 separator = torch.full((direction.size(0), 3, height * upscale - 1, 1), 0)
254 for n in range(direction_symbol.size(0)):
255 if direction[n] == token_forward:
256 for k in range(upscale):
260 (height * upscale) // 2 - upscale // 2 + k,
261 3 + upscale // 2 - abs(k - upscale // 2),
263 elif direction[n] == token_backward:
264 for k in range(upscale):
268 (height * upscale) // 2 - upscale // 2 + k,
269 3 + abs(k - upscale // 2),
272 for k in range(2, upscale - 2):
274 n, :, (height * upscale) // 2 - upscale // 2 + k, k
277 n, :, (height * upscale) // 2 - upscale // 2 + k, upscale - 1 - k
282 mosaic(f_first, upscale),
286 mosaic(f_second, upscale),
295 result.append("".join([token2char[v] for v in s]))
299 ######################################################################
301 if __name__ == "__main__":
305 start_time = time.perf_counter()
306 seq = generate(nb=90, height=height, width=width)
307 delay = time.perf_counter() - start_time
308 print(f"{seq.size(0)/delay:02f} samples/s")
310 print(seq2str(seq[:4]))
312 # m = (torch.rand(seq.size()) < 0.05).long()
313 # seq = (1 - m) * seq + m * 23
315 img = sample2img(seq, height, width)
318 torchvision.utils.save_image(
319 img.float() / 255.0, "/tmp/world.png", nrow=6, padding=6, pad_value=0