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]),
37 max_nb_cached_chunks=None,
41 self.colors = torch.tensor([c for _, c in self.named_colors])
44 self.cache_rec_coo = {}
45 super().__init__(max_nb_cached_chunks, chunk_size, nb_threads)
47 ######################################################################
49 def frame2img(self, x, scale=15):
50 x = x.reshape(x.size(0), self.height, -1)
51 m = torch.logical_and(x >= 0, x < self.nb_token_values()).long()
52 x = self.colors[x * m].permute(0, 3, 1, 2)
54 x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
55 x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
57 x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
58 x[:, :, torch.arange(0, x.size(2), scale), :] = 0
61 for n in range(m.size(0)):
62 for i in range(m.size(1)):
63 for j in range(m.size(2)):
65 for k in range(2, scale - 2):
67 x[n, :, i * scale + k, j * scale + k - l] = 0
69 n, :, i * scale + scale - 1 - k, j * scale + k - l
80 predicted_prompts=None,
81 predicted_answers=None,
85 S = self.height * self.width
86 As = prompts[:, 0 * (S + 1) : 0 * (S + 1) + S].view(-1, self.height, self.width)
87 f_As = prompts[:, 1 * (S + 1) : 1 * (S + 1) + S].view(
88 -1, self.height, self.width
90 Bs = prompts[:, 2 * (S + 1) : 2 * (S + 1) + S].view(-1, self.height, self.width)
91 prompts = torch.cat([As, f_As, Bs], dim=2)
92 answers = answers.reshape(answers.size(0), self.height, self.width)
94 if predicted_prompts is None:
95 predicted_prompts = 255
97 if predicted_answers is None:
98 predicted_answers = 255
100 def add_frame(x, c, margin, bottom=False):
102 h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0
105 x.size(2) + 2 * margin,
106 x.size(3) + 2 * margin,
111 y = x.new_full((x.size(0), x.size(1), h, w), 0)
116 c = c.long()[:, None]
118 (1 - ((c == 1).long() + (c == 0).long() + (c == -1).long()))
119 * torch.tensor([64, 64, 64])
120 + (c == 1).long() * torch.tensor([0, 255, 0])
121 + (c == 0).long() * torch.tensor([255, 255, 255])
122 + (c == -1).long() * torch.tensor([255, 0, 0])
124 y[...] = c[:, :, None, None]
126 y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x
130 img_prompts = torch.cat(
133 add_frame(self.frame2img(x), c=0, margin=1),
137 for x in prompts.to("cpu").split(split_size=self.width, dim=2)
142 h = img_prompts.size(2)
143 img_answers = add_frame(
144 add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1),
149 separator_size = 2 * margin
151 separator = img_prompts.new_full(
161 marker = img_prompts.new_full(
171 # marker[:, :, 0] = 0
172 # marker[:, :, h - 1] = 0
174 for k in range(1, 2 * separator_size - 8):
175 i = k - (separator_size - 4)
176 j = separator_size - 5 - abs(i)
177 marker[:, :, h // 2 - 1 + i, 2 + j] = 0
178 marker[:, :, h // 2 - 1 + i + 1, 2 + j] = 0
189 image_name = os.path.join(result_dir, filename)
190 torchvision.utils.save_image(
198 ######################################################################
200 def nb_token_values(self):
201 return len(self.colors)
210 prevent_overlap=False,
212 if surface_max is None:
213 surface_max = self.height * self.width // 2
215 signature = (nb_rec, min_height, min_width, surface_max)
218 return self.cache_rec_coo[signature].pop()
227 i = torch.randint(self.height, (N * nb_rec, 2)).sort(dim=-1).values
228 j = torch.randint(self.width, (N * nb_rec, 2)).sort(dim=-1).values
231 (i[:, 1] >= i[:, 0] + min_height)
232 & (j[:, 1] >= j[:, 0] + min_height)
233 & ((i[:, 1] - i[:, 0]) * (j[:, 1] - j[:, 0]) <= surface_max)
236 i, j = i[big_enough], j[big_enough]
238 n = i.size(0) - i.size(0) % nb_rec
243 i = i[:n].reshape(n // nb_rec, nb_rec, -1)
244 j = j[:n].reshape(n // nb_rec, nb_rec, -1)
247 can_fit = ((i[:, :, 1] - i[:, :, 0]) * (j[:, :, 1] - j[:, :, 0])).sum(
249 ) <= self.height * self.width
250 i, j = i[can_fit], j[can_fit]
252 A_i1, A_i2, A_j1, A_j2 = (
258 B_i1, B_i2, B_j1, B_j2 = (
264 no_overlap = torch.logical_not(
270 i, j = i[no_overlap], j[no_overlap]
272 A_i1, A_i2, A_j1, A_j2 = (
278 B_i1, B_i2, B_j1, B_j2 = (
284 C_i1, C_i2, C_j1, C_j2 = (
310 i, j = (i[no_overlap], j[no_overlap])
317 self.cache_rec_coo[signature] = [
325 for k in range(nb_rec)
327 for n in range(i.size(0))
330 return self.cache_rec_coo[signature].pop()
332 ######################################################################
335 def task_replace_color(self, A, f_A, B, f_B):
337 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
338 for X, f_X in [(A, f_A), (B, f_B)]:
339 r = self.rec_coo(nb_rec, prevent_overlap=True)
340 for n in range(nb_rec):
341 i1, j1, i2, j2 = r[n]
342 X[i1:i2, j1:j2] = c[n]
343 f_X[i1:i2, j1:j2] = c[n if n > 0 else -1]
346 def task_translate(self, A, f_A, B, f_B):
348 di, dj = torch.randint(3, (2,)) - 1
349 if di.abs() + dj.abs() > 0:
353 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
354 for X, f_X in [(A, f_A), (B, f_B)]:
356 r = self.rec_coo(nb_rec, prevent_overlap=True)
357 i1, j1, i2, j2 = r[nb_rec - 1]
360 and i2 + di < X.size(0)
362 and j2 + dj < X.size(1)
366 for n in range(nb_rec):
367 i1, j1, i2, j2 = r[n]
368 X[i1:i2, j1:j2] = c[n]
370 f_X[i1 + di : i2 + di, j1 + dj : j2 + dj] = c[n]
372 f_X[i1:i2, j1:j2] = c[n]
375 def task_grow(self, A, f_A, B, f_B):
376 di, dj = torch.randint(2, (2,)) * 2 - 1
378 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
379 direction = torch.randint(2, (1,))
380 for X, f_X in [(A, f_A), (B, f_B)]:
382 r = self.rec_coo(nb_rec, prevent_overlap=True)
383 i1, j1, i2, j2 = r[nb_rec - 1]
384 if i1 + 3 < i2 and j1 + 3 < j2:
387 for n in range(nb_rec):
388 i1, j1, i2, j2 = r[n]
391 X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
392 f_X[i1:i2, j1:j2] = c[n]
394 X[i1:i2, j1:j2] = c[n]
395 f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
397 X[i1:i2, j1:j2] = c[n]
398 f_X[i1:i2, j1:j2] = c[n]
401 def task_color_grow(self, A, f_A, B, f_B):
402 di, dj = torch.randint(2, (2,)) * 2 - 1
404 c = torch.randperm(len(self.colors) - 1)[: 2 * nb_rec] + 1
405 direction = torch.randint(4, (1,))
406 for X, f_X in [(A, f_A), (B, f_B)]:
407 r = self.rec_coo(nb_rec, prevent_overlap=True)
408 for n in range(nb_rec):
409 i1, j1, i2, j2 = r[n]
410 X[i1:i2, j1:j2] = c[2 * n]
411 f_X[i1:i2, j1:j2] = c[2 * n]
412 # Not my proudest moment
415 X[i : i + 1, j1:j2] = c[2 * n + 1]
417 f_X[i:i2, j1:j2] = c[2 * n + 1]
419 f_X[i : i + 1, j1:j2] = c[2 * n + 1]
421 i = (i1 + i2 - 1) // 2
422 X[i : i + 1, j1:j2] = c[2 * n + 1]
424 f_X[i1 : i + 1, j1:j2] = c[2 * n + 1]
426 f_X[i : i + 1, j1:j2] = c[2 * n + 1]
429 X[i1:i2, j : j + 1] = c[2 * n + 1]
431 f_X[i1:i2, j:j2] = c[2 * n + 1]
433 f_X[i1:i2, j : j + 1] = c[2 * n + 1]
435 j = (j1 + j2 - 1) // 2
436 X[i1:i2, j : j + 1] = c[2 * n + 1]
438 f_X[i1:i2, j1 : j + 1] = c[2 * n + 1]
440 f_X[i1:i2, j : j + 1] = c[2 * n + 1]
443 def task_frame(self, A, f_A, B, f_B):
445 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
446 for X, f_X in [(A, f_A), (B, f_B)]:
447 r = self.rec_coo(nb_rec, prevent_overlap=True)
448 for n in range(nb_rec):
449 i1, j1, i2, j2 = r[n]
450 X[i1:i2, j1:j2] = c[n]
452 f_X[i1:i2, j1] = c[n]
453 f_X[i1:i2, j2 - 1] = c[n]
454 f_X[i1, j1:j2] = c[n]
455 f_X[i2 - 1, j1:j2] = c[n]
457 f_X[i1:i2, j1:j2] = c[n]
460 def task_detect(self, A, f_A, B, f_B):
462 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
463 for X, f_X in [(A, f_A), (B, f_B)]:
464 r = self.rec_coo(nb_rec, prevent_overlap=True)
465 for n in range(nb_rec):
466 i1, j1, i2, j2 = r[n]
467 X[i1:i2, j1:j2] = c[n]
472 def contact(self, X, i, j, q):
486 if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
487 if X[ii, jj] != 0 and X[ii, jj] != q:
496 if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
497 if X[ii, jj] == q and X[i, jj] != q and X[ii, j] != q:
500 for ii, jj in [(i - 1, j), (i, j - 1), (i, j + 1), (i + 1, j)]:
501 if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
505 return no, nq, nq_diag
507 def task_count(self, A, f_A, B, f_B):
508 N = (torch.randint(4, (1,)) + 2).item()
509 c = torch.randperm(len(self.colors) - 1)[:N] + 1
511 for X, f_X in [(A, f_A), (B, f_B)]:
512 l_q = torch.randperm(self.height * self.width)[
513 : self.height * self.width // 20
515 l_d = torch.randint(N, l_q.size())
516 nb = torch.zeros(N, dtype=torch.int64)
518 for q, e in zip(l_q, l_d):
520 i, j = q % self.height, q // self.height
523 and X[max(0, i - 1) : i + 2, max(0, j - 1) : j + 2] == 0
528 l_q = torch.randperm((self.height - 2) * (self.width - 2))[
529 : self.height * self.width // 2
531 l_d = torch.randint(N, l_q.size())
532 for q, e in zip(l_q, l_d):
534 i, j = q % (self.height - 2) + 1, q // (self.height - 2) + 1
535 a1, a2, a3 = X[i - 1, j - 1 : j + 2]
536 a8, a4 = X[i, j - 1], X[i, j + 1]
537 a7, a6, a5 = X[i + 1, j - 1 : j + 2]
540 and nb[e] < self.width
541 and (a2 == 0 or a2 == d)
542 and (a4 == 0 or a4 == d)
543 and (a6 == 0 or a6 == d)
544 and (a8 == 0 or a8 == d)
545 and (a1 == 0 or a2 == d or a8 == d)
546 and (a3 == 0 or a4 == d or a2 == d)
547 and (a5 == 0 or a6 == d or a4 == d)
548 and (a7 == 0 or a8 == d or a6 == d)
561 for j in range(nb[e]):
565 def task_trajectory(self, A, f_A, B, f_B):
566 c = torch.randperm(len(self.colors) - 1)[:2] + 1
567 for X, f_X in [(A, f_A), (B, f_B)]:
569 di, dj = torch.randint(7, (2,)) - 3
570 i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
572 abs(di) + abs(dj) > 0
574 and i + 2 * di < self.height
576 and j + 2 * dj < self.width
583 and i + k * di < self.height
585 and j + k * dj < self.width
588 X[i + k * di, j + k * dj] = c[k]
589 f_X[i + k * di, j + k * dj] = c[min(k, 1)]
593 def task_bounce(self, A, f_A, B, f_B):
594 c = torch.randperm(len(self.colors) - 1)[:3] + 1
595 for X, f_X in [(A, f_A), (B, f_B)]:
610 for _ in range((self.height * self.width) // 10):
611 i, j = torch.randint(self.height, (1,)), torch.randint(
618 di, dj = torch.randint(7, (2,)) - 3
619 if abs(di) + abs(dj) == 1:
622 i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
630 if free(i + di, j + dj):
632 elif free(i - dj, j + di):
634 if free(i + dj, j - di):
635 if torch.rand(1) < 0.5:
637 elif free(i + dj, j - di):
642 i, j = i + di, j + dj
657 def task_scale(self, A, f_A, B, f_B):
658 c = torch.randperm(len(self.colors) - 1)[:2] + 1
660 i, j = torch.randint(self.height // 2, (1,)), torch.randint(
661 self.width // 2, (1,)
664 for X, f_X in [(A, f_A), (B, f_B)]:
667 i1, j1 = torch.randint(self.height // 2 + 1, (1,)), torch.randint(
668 self.width // 2 + 1, (1,)
670 i2, j2 = torch.randint(self.height // 2 + 1, (1,)), torch.randint(
671 self.width // 2 + 1, (1,)
673 if i1 < i2 and j1 < j2 and min(i2 - i1, j2 - j1) <= 3:
675 X[i + i1 : i + i2, j + j1 : j + j2] = c[0]
676 f_X[2 * i1 : 2 * i2, 2 * j1 : 2 * j2] = c[0]
682 def task_symbols(self, A, f_A, B, f_B):
684 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
686 for X, f_X in [(A, f_A), (B, f_B)]:
688 i, j = torch.randint(self.height - delta + 1, (nb_rec,)), torch.randint(
689 self.width - delta + 1, (nb_rec,)
691 d = (i[None, :] - i[:, None]).abs().max((j[None, :] - j[:, None]).abs())
692 d.fill_diagonal_(delta + 1)
696 for k in range(1, nb_rec):
697 X[i[k] : i[k] + delta, j[k] : j[k] + delta] = c[k]
699 ai, aj = i.float().mean(), j.float().mean()
701 q = torch.randint(3, (1,)) + 1
703 X[i[0] + delta // 2 - 1, j[0] + delta // 2 - 1] = c[0]
704 X[i[0] + delta // 2 - 1, j[0] + delta // 2 + 1] = c[0]
705 X[i[0] + delta // 2 + 1, j[0] + delta // 2 - 1] = c[0]
706 X[i[0] + delta // 2 + 1, j[0] + delta // 2 + 1] = c[0]
708 assert i[q] != ai and j[q] != aj
711 i[0] + delta // 2 + (i[q] - ai).sign().long(),
712 j[0] + delta // 2 + (j[q] - aj).sign().long(),
715 f_X[i[0] : i[0] + delta, j[0] : j[0] + delta] = c[q]
718 def task_ortho(self, A, f_A, B, f_B):
720 di, dj = torch.randint(3, (2,)) - 1
721 o = torch.tensor([[0.0, 1.0], [-1.0, 0.0]])
723 for _ in range(torch.randint(4, (1,))):
725 if torch.rand(1) < 0.5:
728 ci, cj = (self.height - 1) / 2, (self.width - 1) / 2
730 for X, f_X in [(A, f_A), (B, f_B)]:
735 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
737 for r in range(nb_rec):
739 i1, i2 = torch.randint(self.height - 2, (2,)) + 1
740 j1, j2 = torch.randint(self.width - 2, (2,)) + 1
744 and max(i2 - i1, j2 - j1) >= 2
745 and min(i2 - i1, j2 - j1) <= 3
748 X[i1 : i2 + 1, j1 : j2 + 1] = c[r]
750 i1, j1, i2, j2 = i1 - ci, j1 - cj, i2 - ci, j2 - cj
752 i1, j1 = m[0, 0] * i1 + m[0, 1] * j1, m[1, 0] * i1 + m[1, 1] * j1
753 i2, j2 = m[0, 0] * i2 + m[0, 1] * j2, m[1, 0] * i2 + m[1, 1] * j2
755 i1, j1, i2, j2 = i1 + ci, j1 + cj, i2 + ci, j2 + cj
756 i1, i2 = i1.long() + di, i2.long() + di
757 j1, j2 = j1.long() + dj, j2.long() + dj
763 f_X[i1 : i2 + 1, j1 : j2 + 1] = c[r]
765 n = F.one_hot(X.flatten()).sum(dim=0)[1:]
767 n.sum() > self.height * self.width // 4
768 and (n > 0).long().sum() == nb_rec
772 def compute_distance(self, walls, goal_i, goal_j, start_i, start_j):
773 max_length = walls.numel()
774 dist = torch.full_like(walls, max_length)
776 dist[goal_i, goal_j] = 0
777 pred_dist = torch.empty_like(dist)
780 pred_dist.copy_(dist)
784 dist[None, 1:-1, 0:-2],
785 dist[None, 2:, 1:-1],
786 dist[None, 1:-1, 2:],
787 dist[None, 0:-2, 1:-1],
794 dist[1:-1, 1:-1].minimum_(d) # = torch.min(dist[1:-1, 1:-1], d)
795 dist = walls * max_length + (1 - walls) * dist
797 if dist[start_i, start_j] < max_length or dist.equal(pred_dist):
798 return dist * (1 - walls)
801 def task_path(self, A, f_A, B, f_B):
802 c = torch.randperm(len(self.colors) - 1)[:3] + 1
803 dist = torch.empty(self.height + 2, self.width + 2)
804 for X, f_X in [(A, f_A), (B, f_B)]:
805 nb_rec = torch.randint(3, (1,)) + 1
807 r = self.rec_coo(nb_rec, prevent_overlap=True)
810 for n in range(nb_rec):
811 i1, j1, i2, j2 = r[n]
812 X[i1:i2, j1:j2] = c[0]
813 f_X[i1:i2, j1:j2] = c[0]
815 i0, j0 = torch.randint(self.height, (1,)), torch.randint(
821 i1, j1 = torch.randint(self.height, (1,)), torch.randint(
827 dist[1:-1, 1:-1] = (X != 0).long()
828 dist[...] = self.compute_distance(dist, i1 + 1, j1 + 1, i0 + 1, j0 + 1)
829 if dist[i0 + 1, j0 + 1] >= 1 and dist[i0 + 1, j0 + 1] < self.height * 4:
832 dist[1:-1, 1:-1] += (X != 0).long() * self.height * self.width
833 dist[0, :] = self.height * self.width
834 dist[-1, :] = self.height * self.width
835 dist[:, 0] = self.height * self.width
836 dist[:, -1] = self.height * self.width
837 # dist += torch.rand(dist.size())
839 i, j = i0 + 1, j0 + 1
840 while i != i1 + 1 or j != j1 + 1:
841 f_X[i - 1, j - 1] = c[2]
864 # for X, f_X in [(A, f_A), (B, f_B)]:
865 # n = torch.arange(self.height * self.width).reshape(self.height, self.width)
866 # k = torch.randperm(self.height * self.width)
869 # i,j=q%self.height,q//self.height
872 ######################################################################
876 self.task_replace_color,
879 self.task_color_grow,
883 self.task_trajectory,
891 def trivial_prompts_and_answers(self, prompts, answers):
892 S = self.height * self.width
893 Bs = prompts[:, 2 * (S + 1) : 2 * (S + 1) + S]
895 return (Bs == f_Bs).long().min(dim=-1).values > 0
897 def generate_prompts_and_answers_(self, nb, tasks=None, progress_bar=False):
899 tasks = self.all_tasks()
901 S = self.height * self.width
902 prompts = torch.zeros(nb, 3 * S + 2, dtype=torch.int64)
903 answers = torch.zeros(nb, S, dtype=torch.int64)
905 bunch = zip(prompts, answers)
911 desc="world generation",
912 total=prompts.size(0),
915 for prompt, answer in bunch:
916 A = prompt[0 * (S + 1) : 0 * (S + 1) + S].view(self.height, self.width)
917 f_A = prompt[1 * (S + 1) : 1 * (S + 1) + S].view(self.height, self.width)
918 B = prompt[2 * (S + 1) : 2 * (S + 1) + S].view(self.height, self.width)
919 f_B = answer.view(self.height, self.width)
920 task = tasks[torch.randint(len(tasks), (1,))]
923 return prompts.flatten(1), answers.flatten(1)
931 predicted_prompts=None,
932 predicted_answers=None,
937 filename_prefix + ".png",
945 def save_some_examples(self, result_dir):
947 for t in self.all_tasks():
949 prompts, answers = self.generate_prompts_and_answers_(nb, tasks=[t])
951 result_dir, t.__name__, prompts[:nb], answers[:nb], nrow=nrow
955 ######################################################################
957 if __name__ == "__main__":
960 # grids = Grids(max_nb_cached_chunks=5, chunk_size=100, nb_threads=4)
964 # grids = problem.MultiThreadProblem(
965 # grids, max_nb_cached_chunks=50, chunk_size=100, nb_threads=1
968 # start_time = time.perf_counter()
969 # prompts, answers = grids.generate_prompts_and_answers(nb)
970 # delay = time.perf_counter() - start_time
971 # print(f"{prompts.size(0)/delay:02f} seq/s")
978 # for t in grids.all_tasks():
979 for t in [grids.task_path]:
981 prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
982 grids.save_quizzes("/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=nrow)
988 for t in grids.all_tasks():
989 # for t in [ grids.task_replace_color ]: #grids.all_tasks():
990 start_time = time.perf_counter()
991 prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
992 delay = time.perf_counter() - start_time
993 print(f"{t.__name__} {prompts.size(0)/delay:02f} seq/s")
997 m = torch.randint(2, (prompts.size(0),))
998 predicted_prompts = m * (torch.randint(2, (prompts.size(0),)) * 2 - 1)
999 predicted_answers = (1 - m) * (torch.randint(2, (prompts.size(0),)) * 2 - 1)
1006 # You can add a bool to put a frame around the predicted parts
1007 predicted_prompts[:nb],
1008 predicted_answers[:nb],