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>
8 import math, sys, tqdm, os, warnings
10 import torch, torchvision
13 from torch.nn import functional as F
15 ######################################################################
20 class Reasoning(problem.Problem):
22 ("white", [255, 255, 255]),
24 ("green", [0, 192, 0]),
25 ("blue", [0, 0, 255]),
26 ("orange", [255, 192, 0]),
27 ("cyan", [0, 255, 255]),
28 ("violet", [255, 0, 255]),
29 ("lightgreen", [192, 255, 192]),
30 ("pink", [255, 192, 192]),
31 ("lightblue", [192, 192, 255]),
32 ("gray", [192, 192, 192]),
35 def __init__(self, device=torch.device("cpu")):
36 self.colors = torch.tensor([c for _, c in self.named_colors])
37 self.name2color = dict([(p[0], i) for i, p in enumerate(self.named_colors)])
42 ######################################################################
44 def frame2img(self, x, scale=15):
45 x = x.reshape(x.size(0), self.height, -1)
46 x = self.colors[x].permute(0, 3, 1, 2)
48 x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
49 x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
51 x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
52 x[:, :, torch.arange(0, x.size(2), scale), :] = 0
63 predicted_prompts=None,
64 predicted_answers=None,
66 prompts = prompts.reshape(prompts.size(0), self.height, -1)
67 answers = answers.reshape(answers.size(0), self.height, -1)
69 if predicted_prompts is None:
70 predicted_prompts = 255
72 if predicted_answers is None:
73 predicted_answers = 255
75 def add_frame(x, c, margin, bottom=False):
77 h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0
80 x.size(2) + 2 * margin,
81 x.size(3) + 2 * margin,
86 y = x.new_full((x.size(0), x.size(1), h, w), 0)
92 c = c * torch.tensor([192, 192, 192], device=c.device) + (
94 ) * torch.tensor([255, 255, 255], device=c.device)
95 y[...] = c[:, :, None, None]
97 y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x
103 img_prompts = torch.cat(
106 add_frame(self.frame2img(x), c=0, margin=1),
110 for x in prompts.to("cpu").split(split_size=self.width, dim=2)
115 h = img_prompts.size(2)
116 img_answers = add_frame(
117 add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1),
122 separator_size = 2 * margin
124 separator = img_prompts.new_full(
134 marker = img_prompts.new_full(
144 # marker[:, :, 0] = 0
145 # marker[:, :, h - 1] = 0
147 for k in range(1, 2 * separator_size - 8):
148 i = k - (separator_size - 4)
149 j = separator_size - 5 - abs(i)
150 marker[:, :, h // 2 - 1 + i, 2 + j] = 0
151 marker[:, :, h // 2 - 1 + i + 1, 2 + j] = 0
162 image_name = os.path.join(result_dir, filename)
163 torchvision.utils.save_image(
164 img.float() / 255.0, image_name, nrow=4, padding=margin * 4, pad_value=1.0
167 ######################################################################
169 def nb_token_values(self):
170 return len(self.colors)
172 # That's quite a tensorial spaghetti mess to sample
173 # non-overlapping rectangles quickly, but made the generation of
174 # 100k samples go from 1h50 with a lame pure python code to 3min30s
176 def rec_coo(self, x, n, min_height=3, min_width=3):
183 torch.rand(N * K, self.height + 1, device=self.device)
195 torch.rand(N * K, self.width + 1, device=self.device)
205 i = torch.logical_and(
206 v.sum(dim=-1) >= min_height, h.sum(dim=-1) >= min_width
210 v = v[: v.size(0) - v.size(0) % K]
211 h = h[: h.size(0) - h.size(0) % K]
212 v = v.reshape(v.size(0) // K, K, -1)
213 h = h.reshape(h.size(0) // K, K, -1)
215 r = v[:, :, :, None] * h[:, :, None, :]
217 valid = r.sum(dim=1).flatten(1).max(dim=-1).values == 1
225 av = torch.arange(v.size(2), device=self.device)[None, :]
226 ah = torch.arange(h.size(2), device=self.device)[None, :]
229 (i1.item(), j1.item(), i2.item() + 1, j2.item() + 1)
230 for i1, j1, i2, j2 in zip(
231 v.size(2) - (v[0] * (v.size(2) - av)).max(dim=-1).values,
232 h.size(2) - (h[0] * (h.size(2) - ah)).max(dim=-1).values,
233 (v[0] * av).max(dim=-1).values,
234 (h[0] * ah).max(dim=-1).values,
238 def rec_coo_(self, x, n, min_height=3, min_width=3):
239 collision = x.new(x.size())
245 i1, i2 = torch.randint(x.size(0), (2,))
246 if i1 + min_height <= i2:
249 j1, j2 = torch.randint(x.size(1), (2,))
250 if j1 + min_width <= j2:
252 collision[i1:i2, j1:j2] += 1
253 if collision.max() > 1:
255 result.append((i1, j1, i2, j2))
256 if collision.max() == 1:
260 ######################################################################
262 def task_replace_color(self, A, f_A, B, f_B):
264 c = torch.randperm(len(self.colors) - 1)[: N + 1] + 1
265 for X, f_X in [(A, f_A), (B, f_B)]:
266 r = self.rec_coo(X, N)
268 i1, j1, i2, j2 = r[n]
269 X[i1:i2, j1:j2] = c[n]
270 f_X[i1:i2, j1:j2] = c[n if n > 0 else -1]
272 def task_move(self, A, f_A, B, f_B):
273 di, dj = torch.randint(2, (2,)) * 2 - 1
275 c = torch.randperm(len(self.colors) - 1)[:N] + 1
276 for X, f_X in [(A, f_A), (B, f_B)]:
278 r = self.rec_coo(X, N)
279 i1, j1, i2, j2 = r[N - 1]
282 and i2 + di < X.size(0)
284 and j2 + dj < X.size(1)
289 i1, j1, i2, j2 = r[n]
290 X[i1:i2, j1:j2] = c[n]
292 f_X[i1 + di : i2 + di, j1 + dj : j2 + dj] = c[n]
294 f_X[i1:i2, j1:j2] = c[n]
296 def task_grow(self, A, f_A, B, f_B):
297 di, dj = torch.randint(2, (2,)) * 2 - 1
299 c = torch.randperm(len(self.colors) - 1)[:N] + 1
300 direction = torch.randint(2, (1,))
301 for X, f_X in [(A, f_A), (B, f_B)]:
303 r = self.rec_coo(X, N)
304 i1, j1, i2, j2 = r[N - 1]
305 if i1 + 3 < i2 and j1 + 3 < j2:
309 i1, j1, i2, j2 = r[n]
312 X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
313 f_X[i1:i2, j1:j2] = c[n]
315 X[i1:i2, j1:j2] = c[n]
316 f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
318 X[i1:i2, j1:j2] = c[n]
319 f_X[i1:i2, j1:j2] = c[n]
321 def task_color_grow(self, A, f_A, B, f_B):
322 di, dj = torch.randint(2, (2,)) * 2 - 1
324 c = torch.randperm(len(self.colors) - 1)[: 2 * N] + 1
325 direction = torch.randint(4, (1,))
326 for X, f_X in [(A, f_A), (B, f_B)]:
327 r = self.rec_coo(X, N)
329 i1, j1, i2, j2 = r[n]
330 X[i1:i2, j1:j2] = c[2 * n]
331 f_X[i1:i2, j1:j2] = c[2 * n]
332 # Not my proudest moment
335 X[i : i + 1, j1:j2] = c[2 * n + 1]
337 f_X[i:i2, j1:j2] = c[2 * n + 1]
339 f_X[i : i + 1, j1:j2] = c[2 * n + 1]
341 i = (i1 + i2 - 1) // 2
342 X[i : i + 1, j1:j2] = c[2 * n + 1]
344 f_X[i1 : i + 1, j1:j2] = c[2 * n + 1]
346 f_X[i : i + 1, j1:j2] = c[2 * n + 1]
349 X[i1:i2, j : j + 1] = c[2 * n + 1]
351 f_X[i1:i2, j:j2] = c[2 * n + 1]
353 f_X[i1:i2, j : j + 1] = c[2 * n + 1]
355 j = (j1 + j2 - 1) // 2
356 X[i1:i2, j : j + 1] = c[2 * n + 1]
358 f_X[i1:i2, j1 : j + 1] = c[2 * n + 1]
360 f_X[i1:i2, j : j + 1] = c[2 * n + 1]
362 def task_frame(self, A, f_A, B, f_B):
364 c = torch.randperm(len(self.colors) - 1)[: N + 1] + 1
365 for X, f_X in [(A, f_A), (B, f_B)]:
366 r = self.rec_coo(X, N)
368 i1, j1, i2, j2 = r[n]
369 X[i1:i2, j1:j2] = c[n]
370 f_X[i1:i2, j1:j2] = c[n]
372 f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = 0
374 def task_detect(self, A, f_A, B, f_B):
376 c = torch.randperm(len(self.colors) - 1)[: N + 1] + 1
377 for X, f_X in [(A, f_A), (B, f_B)]:
378 r = self.rec_coo(X, N)
380 i1, j1, i2, j2 = r[n]
381 X[i1:i2, j1:j2] = c[n]
384 ######################################################################
386 def generate_prompts_and_answers(self, nb, device="cpu"):
388 self.task_replace_color,
391 self.task_color_grow,
395 prompts = torch.zeros(nb, self.height, self.width * 3, dtype=torch.int64)
396 answers = torch.zeros(nb, self.height, self.width, dtype=torch.int64)
399 for prompt, answer in tqdm.tqdm(
400 zip(prompts, answers),
402 desc="world generation",
403 total=prompts.size(0),
405 A = prompt[:, 0 * w : 1 * w]
406 f_A = prompt[:, 1 * w : 2 * w]
407 B = prompt[:, 2 * w : 3 * w]
409 task = tasks[torch.randint(len(tasks), (1,))]
412 return prompts.flatten(1), answers.flatten(1)
420 predicted_prompts=None,
421 predicted_answers=None,
425 filename_prefix + ".png",
433 ######################################################################
435 if __name__ == "__main__":
438 reasoning = Reasoning()
440 start_time = time.perf_counter()
441 prompts, answers = reasoning.generate_prompts_and_answers(100)
442 delay = time.perf_counter() - start_time
443 print(f"{prompts.size(0)/delay:02f} seq/s")
445 # predicted_prompts = torch.rand(prompts.size(0)) < 0.5
446 # predicted_answers = torch.logical_not(predicted_prompts)
448 reasoning.save_quizzes(
453 # You can add a bool to put a frame around the predicted parts
454 # predicted_prompts, predicted_answers