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 Grids(problem.Problem):
22 ("white", [255, 255, 255]),
24 ("green", [0, 192, 0]),
25 ("blue", [0, 0, 255]),
26 ("yellow", [255, 224, 0]),
27 ("cyan", [0, 255, 255]),
28 ("violet", [224, 128, 255]),
29 ("lightgreen", [192, 255, 192]),
30 ("brown", [165, 42, 42]),
31 ("lightblue", [192, 192, 255]),
32 ("gray", [128, 128, 128]),
35 def __init__(self, device=torch.device("cpu")):
36 self.colors = torch.tensor([c for _, c in self.named_colors])
41 ######################################################################
43 def frame2img(self, x, scale=15):
44 x = x.reshape(x.size(0), self.height, -1)
45 m = torch.logical_and(x >= 0, x < self.nb_token_values()).long()
46 x = self.colors[x * m].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
55 for n in range(m.size(0)):
56 for i in range(m.size(1)):
57 for j in range(m.size(2)):
59 for k in range(2, scale - 2):
61 x[n, :, i * scale + k, j * scale + k - l] = 0
63 n, :, i * scale + scale - 1 - k, j * scale + k - l
74 predicted_prompts=None,
75 predicted_answers=None,
79 S = self.height * self.width
80 As = prompts[:, 0 * (S + 1) : 0 * (S + 1) + S].view(-1, self.height, self.width)
81 f_As = prompts[:, 1 * (S + 1) : 1 * (S + 1) + S].view(
82 -1, self.height, self.width
84 Bs = prompts[:, 2 * (S + 1) : 2 * (S + 1) + S].view(-1, self.height, self.width)
85 prompts = torch.cat([As, f_As, Bs], dim=2)
86 answers = answers.reshape(answers.size(0), self.height, self.width)
88 if predicted_prompts is None:
89 predicted_prompts = 255
91 if predicted_answers is None:
92 predicted_answers = 255
94 def add_frame(x, c, margin, bottom=False):
96 h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0
99 x.size(2) + 2 * margin,
100 x.size(3) + 2 * margin,
105 y = x.new_full((x.size(0), x.size(1), h, w), 0)
110 c = c.long()[:, None]
112 (1 - ((c == 1).long() + (c == 0).long() + (c == -1).long()))
113 * torch.tensor([64, 64, 64], device=c.device)
114 + (c == 1).long() * torch.tensor([0, 255, 0], device=c.device)
115 + (c == 0).long() * torch.tensor([255, 255, 255], device=c.device)
116 + (c == -1).long() * torch.tensor([255, 0, 0], device=c.device)
118 y[...] = c[:, :, None, None]
120 y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x
124 img_prompts = torch.cat(
127 add_frame(self.frame2img(x), c=0, margin=1),
131 for x in prompts.to("cpu").split(split_size=self.width, dim=2)
136 h = img_prompts.size(2)
137 img_answers = add_frame(
138 add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1),
143 separator_size = 2 * margin
145 separator = img_prompts.new_full(
155 marker = img_prompts.new_full(
165 # marker[:, :, 0] = 0
166 # marker[:, :, h - 1] = 0
168 for k in range(1, 2 * separator_size - 8):
169 i = k - (separator_size - 4)
170 j = separator_size - 5 - abs(i)
171 marker[:, :, h // 2 - 1 + i, 2 + j] = 0
172 marker[:, :, h // 2 - 1 + i + 1, 2 + j] = 0
183 image_name = os.path.join(result_dir, filename)
184 torchvision.utils.save_image(
192 ######################################################################
194 def nb_token_values(self):
195 return len(self.colors)
198 def rec_coo_(self, nb_rec, min_height=3, min_width=3):
200 def overlap(ia, ja, ib, jb):
202 ia[1] >= ib[0] and ia[0] <= ib[1] and ja[1] >= jb[0] and ja[0] <= jb[1]
207 i = torch.randint(self.height + 1, (nb_rec, 2)).sort(dim=1).values
208 j = torch.randint(self.width + 1, (nb_rec, 2)).sort(dim=1).values
211 overlap(i[0], j[0], i[1], j[1])
212 or overlap(i[0], j[0], i[2], j[2])
213 or overlap(i[1], j[1], i[2], j[2])
215 and (i[:, 1] - i[:, 0]).min() >= min_height
216 and (j[:, 1] - j[:, 0]).min() >= min_width
220 (i[0, 0], j[0, 0], i[0, 1], j[0, 1]),
221 (i[1, 0], j[1, 0], i[1, 1], j[1, 1]),
222 (i[2, 0], j[2, 0], i[2, 1], j[2, 1]),
225 # That's quite a tensorial spaghetti mess to sample
226 # non-overlapping rectangles quickly, but made the generation of
227 # 100k samples go from 1h50 with a lame pure python code to 3min30s
230 def rec_coo(self, nb_rec, min_height=3, min_width=3):
236 torch.rand(nb_trials * nb_rec, self.height + 1, device=self.device)
248 torch.rand(nb_trials * nb_rec, self.width + 1, device=self.device)
258 i = torch.logical_and(
259 v.sum(dim=-1) >= min_height, h.sum(dim=-1) >= min_width
263 v = v[: v.size(0) - v.size(0) % nb_rec]
264 h = h[: h.size(0) - h.size(0) % nb_rec]
265 v = v.reshape(v.size(0) // nb_rec, nb_rec, -1)
266 h = h.reshape(h.size(0) // nb_rec, nb_rec, -1)
268 r = v[:, :, :, None] * h[:, :, None, :]
270 valid = r.sum(dim=1).flatten(1).max(dim=-1).values == 1
278 av = torch.arange(v.size(2), device=self.device)[None, :]
279 ah = torch.arange(h.size(2), device=self.device)[None, :]
282 (i1.item(), j1.item(), i2.item() + 1, j2.item() + 1)
283 for i1, j1, i2, j2 in zip(
284 v.size(2) - (v[0] * (v.size(2) - av)).max(dim=-1).values,
285 h.size(2) - (h[0] * (h.size(2) - ah)).max(dim=-1).values,
286 (v[0] * av).max(dim=-1).values,
287 (h[0] * ah).max(dim=-1).values,
292 def rec_coo_(self, x, n, min_height=3, min_width=3):
293 collision = x.new(x.size())
299 i1, i2 = torch.randint(x.size(0), (2,))
300 if i1 + min_height <= i2:
303 j1, j2 = torch.randint(x.size(1), (2,))
304 if j1 + min_width <= j2:
306 collision[i1:i2, j1:j2] += 1
307 if collision.max() > 1:
309 result.append((i1, j1, i2, j2))
310 if collision.max() == 1:
314 ######################################################################
317 def task_replace_color(self, A, f_A, B, f_B):
319 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
320 for X, f_X in [(A, f_A), (B, f_B)]:
321 r = self.rec_coo(nb_rec)
322 for n in range(nb_rec):
323 i1, j1, i2, j2 = r[n]
324 X[i1:i2, j1:j2] = c[n]
325 f_X[i1:i2, j1:j2] = c[n if n > 0 else -1]
328 def task_translate(self, A, f_A, B, f_B):
329 di, dj = torch.randint(3, (2,)) - 1
331 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
332 for X, f_X in [(A, f_A), (B, f_B)]:
334 r = self.rec_coo(nb_rec)
335 i1, j1, i2, j2 = r[nb_rec - 1]
338 and i2 + di < X.size(0)
340 and j2 + dj < X.size(1)
344 for n in range(nb_rec):
345 i1, j1, i2, j2 = r[n]
346 X[i1:i2, j1:j2] = c[n]
348 f_X[i1 + di : i2 + di, j1 + dj : j2 + dj] = c[n]
350 f_X[i1:i2, j1:j2] = c[n]
353 def task_grow(self, A, f_A, B, f_B):
354 di, dj = torch.randint(2, (2,)) * 2 - 1
356 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
357 direction = torch.randint(2, (1,))
358 for X, f_X in [(A, f_A), (B, f_B)]:
360 r = self.rec_coo(nb_rec)
361 i1, j1, i2, j2 = r[nb_rec - 1]
362 if i1 + 3 < i2 and j1 + 3 < j2:
365 for n in range(nb_rec):
366 i1, j1, i2, j2 = r[n]
369 X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
370 f_X[i1:i2, j1:j2] = c[n]
372 X[i1:i2, j1:j2] = c[n]
373 f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
375 X[i1:i2, j1:j2] = c[n]
376 f_X[i1:i2, j1:j2] = c[n]
379 def task_color_grow(self, A, f_A, B, f_B):
380 di, dj = torch.randint(2, (2,)) * 2 - 1
382 c = torch.randperm(len(self.colors) - 1)[: 2 * nb_rec] + 1
383 direction = torch.randint(4, (1,))
384 for X, f_X in [(A, f_A), (B, f_B)]:
385 r = self.rec_coo(nb_rec)
386 for n in range(nb_rec):
387 i1, j1, i2, j2 = r[n]
388 X[i1:i2, j1:j2] = c[2 * n]
389 f_X[i1:i2, j1:j2] = c[2 * n]
390 # Not my proudest moment
393 X[i : i + 1, j1:j2] = c[2 * n + 1]
395 f_X[i:i2, j1:j2] = c[2 * n + 1]
397 f_X[i : i + 1, j1:j2] = c[2 * n + 1]
399 i = (i1 + i2 - 1) // 2
400 X[i : i + 1, j1:j2] = c[2 * n + 1]
402 f_X[i1 : i + 1, j1:j2] = c[2 * n + 1]
404 f_X[i : i + 1, j1:j2] = c[2 * n + 1]
407 X[i1:i2, j : j + 1] = c[2 * n + 1]
409 f_X[i1:i2, j:j2] = c[2 * n + 1]
411 f_X[i1:i2, j : j + 1] = c[2 * n + 1]
413 j = (j1 + j2 - 1) // 2
414 X[i1:i2, j : j + 1] = c[2 * n + 1]
416 f_X[i1:i2, j1 : j + 1] = c[2 * n + 1]
418 f_X[i1:i2, j : j + 1] = c[2 * n + 1]
421 def task_frame(self, A, f_A, B, f_B):
423 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
424 for X, f_X in [(A, f_A), (B, f_B)]:
425 r = self.rec_coo(nb_rec)
426 for n in range(nb_rec):
427 i1, j1, i2, j2 = r[n]
428 X[i1:i2, j1:j2] = c[n]
429 f_X[i1:i2, j1:j2] = c[n]
431 f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = 0
434 def task_detect(self, A, f_A, B, f_B):
436 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
437 for X, f_X in [(A, f_A), (B, f_B)]:
438 r = self.rec_coo(nb_rec)
439 for n in range(nb_rec):
440 i1, j1, i2, j2 = r[n]
441 X[i1:i2, j1:j2] = c[n]
446 def contact(self, X, i, j, q):
460 if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
461 if X[ii, jj] != 0 and X[ii, jj] != q:
470 if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
471 if X[ii, jj] == q and X[i, jj] != q and X[ii, j] != q:
474 for ii, jj in [(i - 1, j), (i, j - 1), (i, j + 1), (i + 1, j)]:
475 if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
479 return no, nq, nq_diag
482 def task_count(self, A, f_A, B, f_B):
483 N = (torch.randint(4, (1,)) + 2).item()
484 c = torch.randperm(len(self.colors) - 1)[:N] + 1
486 for X, f_X in [(A, f_A), (B, f_B)]:
487 nb = torch.zeros(N, dtype=torch.int64)
488 q = torch.randint(N, (self.height * self.width,))
489 k = torch.randperm(self.height * self.width)
490 for p in range(self.height * self.width):
491 i, j = k[p] % self.height, k[p] // self.height
492 no, nq, nq_diag = self.contact(X, i, j, c[q[p]])
493 if no == 0 and nq_diag == 0:
495 if nb[q[p]] < self.width:
502 for j in range(nb[n]):
506 def task_trajectory(self, A, f_A, B, f_B):
507 c = torch.randperm(len(self.colors) - 1)[:2] + 1
508 for X, f_X in [(A, f_A), (B, f_B)]:
510 di, dj = torch.randint(7, (2,)) - 3
511 i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
513 abs(di) + abs(dj) > 0
515 and i + 2 * di < self.height
517 and j + 2 * dj < self.width
524 and i + k * di < self.height
526 and j + k * dj < self.width
529 X[i + k * di, j + k * dj] = c[k]
530 f_X[i + k * di, j + k * dj] = c[min(k, 1)]
534 def task_bounce(self, A, f_A, B, f_B):
535 c = torch.randperm(len(self.colors) - 1)[:3] + 1
536 for X, f_X in [(A, f_A), (B, f_B)]:
551 for _ in range((self.height * self.width) // 10):
552 i, j = torch.randint(self.height, (1,)), torch.randint(
559 di, dj = torch.randint(7, (2,)) - 3
560 if abs(di) + abs(dj) == 1:
563 i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
571 if free(i + di, j + dj):
573 elif free(i - dj, j + di):
575 if free(i + dj, j - di):
576 if torch.rand(1) < 0.5:
578 elif free(i + dj, j - di):
583 i, j = i + di, j + dj
598 def task_scale(self, A, f_A, B, f_B):
599 c = torch.randperm(len(self.colors) - 1)[:2] + 1
601 i, j = torch.randint(self.height // 2, (1,)), torch.randint(
602 self.width // 2, (1,)
605 for X, f_X in [(A, f_A), (B, f_B)]:
608 i1, j1 = torch.randint(self.height // 2 + 1, (1,)), torch.randint(
609 self.width // 2 + 1, (1,)
611 i2, j2 = torch.randint(self.height // 2 + 1, (1,)), torch.randint(
612 self.width // 2 + 1, (1,)
614 if i1 < i2 and j1 < j2 and min(i2 - i1, j2 - j1) <= 3:
616 X[i + i1 : i + i2, j + j1 : j + j2] = c[0]
617 f_X[2 * i1 : 2 * i2, 2 * j1 : 2 * j2] = c[0]
623 def task_symbols(self, A, f_A, B, f_B):
625 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
627 for X, f_X in [(A, f_A), (B, f_B)]:
629 i, j = torch.randint(self.height - delta + 1, (nb_rec,)), torch.randint(
630 self.width - delta + 1, (nb_rec,)
632 d = (i[None, :] - i[:, None]).abs().max((j[None, :] - j[:, None]).abs())
633 d.fill_diagonal_(delta + 1)
637 for k in range(1, nb_rec):
638 X[i[k] : i[k] + delta, j[k] : j[k] + delta] = c[k]
640 ai, aj = i.float().mean(), j.float().mean()
642 q = torch.randint(3, (1,)) + 1
644 X[i[0] + delta // 2 - 1, j[0] + delta // 2 - 1] = c[0]
645 X[i[0] + delta // 2 - 1, j[0] + delta // 2 + 1] = c[0]
646 X[i[0] + delta // 2 + 1, j[0] + delta // 2 - 1] = c[0]
647 X[i[0] + delta // 2 + 1, j[0] + delta // 2 + 1] = c[0]
649 assert i[q] != ai and j[q] != aj
652 i[0] + delta // 2 + (i[q] - ai).sign().long(),
653 j[0] + delta // 2 + (j[q] - aj).sign().long(),
656 f_X[i[0] : i[0] + delta, j[0] : j[0] + delta] = c[q]
659 def task_ortho(self, A, f_A, B, f_B):
661 di, dj = torch.randint(3, (2,)) - 1
662 o = torch.tensor([[0.0, 1.0], [-1.0, 0.0]])
664 for _ in range(torch.randint(4, (1,))):
666 if torch.rand(1) < 0.5:
669 ci, cj = (self.height - 1) / 2, (self.width - 1) / 2
671 for X, f_X in [(A, f_A), (B, f_B)]:
676 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
678 for r in range(nb_rec):
680 i1, i2 = torch.randint(self.height - 2, (2,)) + 1
681 j1, j2 = torch.randint(self.width - 2, (2,)) + 1
685 and max(i2 - i1, j2 - j1) >= 2
686 and min(i2 - i1, j2 - j1) <= 3
689 X[i1 : i2 + 1, j1 : j2 + 1] = c[r]
691 i1, j1, i2, j2 = i1 - ci, j1 - cj, i2 - ci, j2 - cj
693 i1, j1 = m[0, 0] * i1 + m[0, 1] * j1, m[1, 0] * i1 + m[1, 1] * j1
694 i2, j2 = m[0, 0] * i2 + m[0, 1] * j2, m[1, 0] * i2 + m[1, 1] * j2
696 i1, j1, i2, j2 = i1 + ci, j1 + cj, i2 + ci, j2 + cj
697 i1, i2 = i1.long() + di, i2.long() + di
698 j1, j2 = j1.long() + dj, j2.long() + dj
704 f_X[i1 : i2 + 1, j1 : j2 + 1] = c[r]
706 n = F.one_hot(X.flatten()).sum(dim=0)[1:]
708 n.sum() > self.height * self.width // 4
709 and (n > 0).long().sum() == nb_rec
714 def task_islands(self, A, f_A, B, f_B):
717 # for X, f_X in [(A, f_A), (B, f_B)]:
718 # n = torch.arange(self.height * self.width).reshape(self.height, self.width)
719 # k = torch.randperm(self.height * self.width)
722 # i,j=q%self.height,q//self.height
725 ######################################################################
729 self.task_replace_color,
732 self.task_color_grow,
736 self.task_trajectory,
744 def trivial_prompts_and_answers(self, prompts, answers):
745 S = self.height * self.width
746 Bs = prompts[:, 2 * (S + 1) : 2 * (S + 1) + S]
748 return (Bs == f_Bs).long().min(dim=-1).values > 0
750 def generate_prompts_and_answers(
751 self, nb, tasks=None, progress_bar=False, device="cpu"
754 tasks = self.all_tasks()
756 S = self.height * self.width
757 prompts = torch.zeros(nb, 3 * S + 2, dtype=torch.int64)
758 answers = torch.zeros(nb, S, dtype=torch.int64)
760 bunch = zip(prompts, answers)
766 desc="world generation",
767 total=prompts.size(0),
770 for prompt, answer in bunch:
771 A = prompt[0 * (S + 1) : 0 * (S + 1) + S].view(self.height, self.width)
772 f_A = prompt[1 * (S + 1) : 1 * (S + 1) + S].view(self.height, self.width)
773 B = prompt[2 * (S + 1) : 2 * (S + 1) + S].view(self.height, self.width)
774 f_B = answer.view(self.height, self.width)
775 task = tasks[torch.randint(len(tasks), (1,))]
778 return prompts.flatten(1), answers.flatten(1)
786 predicted_prompts=None,
787 predicted_answers=None,
792 filename_prefix + ".png",
801 ######################################################################
803 if __name__ == "__main__":
809 # grids = problem.MultiThreadProblem(
810 # grids, max_nb_cached_chunks=50, chunk_size=100, nb_threads=1
813 # start_time = time.perf_counter()
814 # prompts, answers = grids.generate_prompts_and_answers(nb)
815 # delay = time.perf_counter() - start_time
816 # print(f"{prompts.size(0)/delay:02f} seq/s")
822 for t in grids.all_tasks():
823 # for t in [grids.task_ortho]:
825 prompts, answers = grids.generate_prompts_and_answers(nb, tasks=[t])
826 grids.save_quizzes("/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=4)
832 for t in grids.all_tasks():
833 start_time = time.perf_counter()
834 prompts, answers = grids.generate_prompts_and_answers(nb, tasks=[t])
835 delay = time.perf_counter() - start_time
836 print(f"{t.__name__} {prompts.size(0)/delay:02f} seq/s")
840 m = torch.randint(2, (prompts.size(0),))
841 predicted_prompts = m * (torch.randint(2, (prompts.size(0),)) * 2 - 1)
842 predicted_answers = (1 - m) * (torch.randint(2, (prompts.size(0),)) * 2 - 1)
849 # You can add a bool to put a frame around the predicted parts
850 predicted_prompts[:nb],
851 predicted_answers[:nb],