grids.py

   1 #!/usr/bin/env python
   2
   3 # Any copyright is dedicated to the Public Domain.
   4 # https://creativecommons.org/publicdomain/zero/1.0/
   5
   6 # Written by Francois Fleuret <francois@fleuret.org>
   7
   8 import math, sys, tqdm, os, warnings
   9
  10 import torch, torchvision
  11
  12 from torch import nn
  13 from torch.nn import functional as F
  14
  15 ######################################################################
  16
  17 import problem
  18
  19
  20 class Grids(problem.Problem):
  21     named_colors = [
  22         ("white", [255, 255, 255]),
  23         ("red", [255, 0, 0]),
  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         ("brown", [165, 42, 42]),
  31         ("lightblue", [192, 192, 255]),
  32         ("gray", [128, 128, 128]),
  33     ]
  34
  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)])
  38         self.height = 10
  39         self.width = 10
  40         self.device = device
  41
  42     ######################################################################
  43
  44     def frame2img(self, x, scale=15):
  45         x = x.reshape(x.size(0), self.height, -1)
  46         m = torch.logical_and(x >= 0, x < self.nb_token_values()).long()
  47         x = self.colors[x * m].permute(0, 3, 1, 2)
  48         s = x.shape
  49         x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
  50         x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
  51
  52         x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
  53         x[:, :, torch.arange(0, x.size(2), scale), :] = 0
  54         x = x[:, :, 1:, 1:]
  55
  56         for n in range(m.size(0)):
  57             for i in range(m.size(1)):
  58                 for j in range(m.size(2)):
  59                     if m[n, i, j] == 0:
  60                         for k in range(2, scale - 2):
  61                             for l in [0, 1]:
  62                                 x[n, :, i * scale + k, j * scale + k - l] = 0
  63                                 x[
  64                                     n, :, i * scale + scale - 1 - k, j * scale + k - l
  65                                 ] = 0
  66
  67         return x
  68
  69     def frame2img_(self, x, scale=15):
  70         x = x.reshape(x.size(0), self.height, -1)
  71         x = self.colors[x].permute(0, 3, 1, 2)
  72         s = x.shape
  73         x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
  74         x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
  75
  76         x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
  77         x[:, :, torch.arange(0, x.size(2), scale), :] = 0
  78         x = x[:, :, 1:, 1:]
  79
  80         return x
  81
  82     def save_image(
  83         self,
  84         result_dir,
  85         filename,
  86         prompts,
  87         answers,
  88         predicted_prompts=None,
  89         predicted_answers=None,
  90         nrow=4,
  91     ):
  92         S = self.height * self.width
  93         As = prompts[:, 0 * (S + 1) : 0 * (S + 1) + S].view(-1, self.height, self.width)
  94         f_As = prompts[:, 1 * (S + 1) : 1 * (S + 1) + S].view(
  95             -1, self.height, self.width
  96         )
  97         Bs = prompts[:, 2 * (S + 1) : 2 * (S + 1) + S].view(-1, self.height, self.width)
  98         prompts = torch.cat([As, f_As, Bs], dim=2)
  99         answers = answers.reshape(answers.size(0), self.height, self.width)
 100
 101         if predicted_prompts is None:
 102             predicted_prompts = 255
 103
 104         if predicted_answers is None:
 105             predicted_answers = 255
 106
 107         def add_frame(x, c, margin, bottom=False):
 108             if bottom:
 109                 h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0
 110             else:
 111                 h, w, di, dj = (
 112                     x.size(2) + 2 * margin,
 113                     x.size(3) + 2 * margin,
 114                     margin,
 115                     margin,
 116                 )
 117
 118             y = x.new_full((x.size(0), x.size(1), h, w), 0)
 119
 120             if type(c) is int:
 121                 y[...] = c
 122             else:
 123                 c = c.long()[:, None]
 124                 c = (
 125                     (1 - ((c == 1).long() + (c == 0).long() + (c == -1).long()))
 126                     * torch.tensor([64, 64, 64], device=c.device)
 127                     + (c == 1).long() * torch.tensor([0, 255, 0], device=c.device)
 128                     + (c == 0).long() * torch.tensor([255, 255, 255], device=c.device)
 129                     + (c == -1).long() * torch.tensor([255, 0, 0], device=c.device)
 130                 )
 131                 y[...] = c[:, :, None, None]
 132
 133             y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x
 134
 135             return y
 136
 137         margin = 8
 138
 139         img_prompts = torch.cat(
 140             [
 141                 add_frame(
 142                     add_frame(self.frame2img(x), c=0, margin=1),
 143                     c=predicted_prompts,
 144                     margin=margin,
 145                 )
 146                 for x in prompts.to("cpu").split(split_size=self.width, dim=2)
 147             ],
 148             dim=3,
 149         )
 150
 151         h = img_prompts.size(2)
 152         img_answers = add_frame(
 153             add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1),
 154             c=predicted_answers,
 155             margin=margin,
 156         )
 157
 158         separator_size = 2 * margin
 159
 160         separator = img_prompts.new_full(
 161             (
 162                 img_prompts.size(0),
 163                 img_prompts.size(1),
 164                 img_prompts.size(2),
 165                 separator_size,
 166             ),
 167             255,
 168         )
 169
 170         marker = img_prompts.new_full(
 171             (
 172                 img_prompts.size(0),
 173                 img_prompts.size(1),
 174                 img_prompts.size(2),
 175                 separator_size,
 176             ),
 177             255,
 178         )
 179
 180         # marker[:, :, 0] = 0
 181         # marker[:, :, h - 1] = 0
 182
 183         for k in range(1, 2 * separator_size - 8):
 184             i = k - (separator_size - 4)
 185             j = separator_size - 5 - abs(i)
 186             marker[:, :, h // 2 - 1 + i, 2 + j] = 0
 187             marker[:, :, h // 2 - 1 + i + 1, 2 + j] = 0
 188
 189         img = torch.cat(
 190             [
 191                 img_prompts,
 192                 marker,
 193                 img_answers,
 194             ],
 195             dim=3,
 196         )
 197
 198         image_name = os.path.join(result_dir, filename)
 199         torchvision.utils.save_image(
 200             img.float() / 255.0,
 201             image_name,
 202             nrow=nrow,
 203             padding=margin * 4,
 204             pad_value=1.0,
 205         )
 206
 207     ######################################################################
 208
 209     def nb_token_values(self):
 210         return len(self.colors)
 211
 212     # That's quite a tensorial spaghetti mess to sample
 213     # non-overlapping rectangles quickly, but made the generation of
 214     # 100k samples go from 1h50 with a lame pure python code to 3min30s
 215     # with this one.
 216     def rec_coo(self, nb_rec, min_height=3, min_width=3):
 217         nb_trials = 200
 218
 219         while True:
 220             v = (
 221                 (
 222                     torch.rand(nb_trials * nb_rec, self.height + 1, device=self.device)
 223                     .sort(dim=-1)
 224                     .indices
 225                     < 2
 226                 )
 227                 .long()
 228                 .cumsum(dim=1)
 229                 == 1
 230             ).long()
 231
 232             h = (
 233                 (
 234                     torch.rand(nb_trials * nb_rec, self.width + 1, device=self.device)
 235                     .sort(dim=-1)
 236                     .indices
 237                     < 2
 238                 )
 239                 .long()
 240                 .cumsum(dim=1)
 241                 == 1
 242             ).long()
 243
 244             i = torch.logical_and(
 245                 v.sum(dim=-1) >= min_height, h.sum(dim=-1) >= min_width
 246             )
 247
 248             v, h = v[i], h[i]
 249             v = v[: v.size(0) - v.size(0) % nb_rec]
 250             h = h[: h.size(0) - h.size(0) % nb_rec]
 251             v = v.reshape(v.size(0) // nb_rec, nb_rec, -1)
 252             h = h.reshape(h.size(0) // nb_rec, nb_rec, -1)
 253
 254             r = v[:, :, :, None] * h[:, :, None, :]
 255
 256             valid = r.sum(dim=1).flatten(1).max(dim=-1).values == 1
 257
 258             v = v[valid]
 259             h = h[valid]
 260
 261             if v.size(0) > 0:
 262                 break
 263
 264         av = torch.arange(v.size(2), device=self.device)[None, :]
 265         ah = torch.arange(h.size(2), device=self.device)[None, :]
 266
 267         return [
 268             (i1.item(), j1.item(), i2.item() + 1, j2.item() + 1)
 269             for i1, j1, i2, j2 in zip(
 270                 v.size(2) - (v[0] * (v.size(2) - av)).max(dim=-1).values,
 271                 h.size(2) - (h[0] * (h.size(2) - ah)).max(dim=-1).values,
 272                 (v[0] * av).max(dim=-1).values,
 273                 (h[0] * ah).max(dim=-1).values,
 274             )
 275         ]
 276
 277     def rec_coo_(self, x, n, min_height=3, min_width=3):
 278         collision = x.new(x.size())
 279         while True:
 280             collision[...] = 0
 281             result = []
 282             for _ in range(n):
 283                 while True:
 284                     i1, i2 = torch.randint(x.size(0), (2,))
 285                     if i1 + min_height <= i2:
 286                         break
 287                 while True:
 288                     j1, j2 = torch.randint(x.size(1), (2,))
 289                     if j1 + min_width <= j2:
 290                         break
 291                 collision[i1:i2, j1:j2] += 1
 292                 if collision.max() > 1:
 293                     break
 294                 result.append((i1, j1, i2, j2))
 295             if collision.max() == 1:
 296                 break
 297         return result
 298
 299     ######################################################################
 300
 301     def task_replace_color(self, A, f_A, B, f_B):
 302         nb_rec = 3
 303         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
 304         for X, f_X in [(A, f_A), (B, f_B)]:
 305             r = self.rec_coo(nb_rec)
 306             for n in range(nb_rec):
 307                 i1, j1, i2, j2 = r[n]
 308                 X[i1:i2, j1:j2] = c[n]
 309                 f_X[i1:i2, j1:j2] = c[n if n > 0 else -1]
 310
 311     def task_translate(self, A, f_A, B, f_B):
 312         di, dj = torch.randint(3, (2,)) - 1
 313         nb_rec = 3
 314         c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
 315         for X, f_X in [(A, f_A), (B, f_B)]:
 316             while True:
 317                 r = self.rec_coo(nb_rec)
 318                 i1, j1, i2, j2 = r[nb_rec - 1]
 319                 if (
 320                     i1 + di >= 0
 321                     and i2 + di < X.size(0)
 322                     and j1 + dj >= 0
 323                     and j2 + dj < X.size(1)
 324                 ):
 325                     break
 326
 327             for n in range(nb_rec):
 328                 i1, j1, i2, j2 = r[n]
 329                 X[i1:i2, j1:j2] = c[n]
 330                 if n == nb_rec - 1:
 331                     f_X[i1 + di : i2 + di, j1 + dj : j2 + dj] = c[n]
 332                 else:
 333                     f_X[i1:i2, j1:j2] = c[n]
 334
 335     def task_grow(self, A, f_A, B, f_B):
 336         di, dj = torch.randint(2, (2,)) * 2 - 1
 337         nb_rec = 3
 338         c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
 339         direction = torch.randint(2, (1,))
 340         for X, f_X in [(A, f_A), (B, f_B)]:
 341             while True:
 342                 r = self.rec_coo(nb_rec)
 343                 i1, j1, i2, j2 = r[nb_rec - 1]
 344                 if i1 + 3 < i2 and j1 + 3 < j2:
 345                     break
 346
 347             for n in range(nb_rec):
 348                 i1, j1, i2, j2 = r[n]
 349                 if n == nb_rec - 1:
 350                     if direction == 0:
 351                         X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
 352                         f_X[i1:i2, j1:j2] = c[n]
 353                     else:
 354                         X[i1:i2, j1:j2] = c[n]
 355                         f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
 356                 else:
 357                     X[i1:i2, j1:j2] = c[n]
 358                     f_X[i1:i2, j1:j2] = c[n]
 359
 360     def task_color_grow(self, A, f_A, B, f_B):
 361         di, dj = torch.randint(2, (2,)) * 2 - 1
 362         nb_rec = 3
 363         c = torch.randperm(len(self.colors) - 1)[: 2 * nb_rec] + 1
 364         direction = torch.randint(4, (1,))
 365         for X, f_X in [(A, f_A), (B, f_B)]:
 366             r = self.rec_coo(nb_rec)
 367             for n in range(nb_rec):
 368                 i1, j1, i2, j2 = r[n]
 369                 X[i1:i2, j1:j2] = c[2 * n]
 370                 f_X[i1:i2, j1:j2] = c[2 * n]
 371                 # Not my proudest moment
 372                 if direction == 0:
 373                     i = (i1 + i2) // 2
 374                     X[i : i + 1, j1:j2] = c[2 * n + 1]
 375                     if n == nb_rec - 1:
 376                         f_X[i:i2, j1:j2] = c[2 * n + 1]
 377                     else:
 378                         f_X[i : i + 1, j1:j2] = c[2 * n + 1]
 379                 elif direction == 1:
 380                     i = (i1 + i2 - 1) // 2
 381                     X[i : i + 1, j1:j2] = c[2 * n + 1]
 382                     if n == nb_rec - 1:
 383                         f_X[i1 : i + 1, j1:j2] = c[2 * n + 1]
 384                     else:
 385                         f_X[i : i + 1, j1:j2] = c[2 * n + 1]
 386                 elif direction == 2:
 387                     j = (j1 + j2) // 2
 388                     X[i1:i2, j : j + 1] = c[2 * n + 1]
 389                     if n == nb_rec - 1:
 390                         f_X[i1:i2, j:j2] = c[2 * n + 1]
 391                     else:
 392                         f_X[i1:i2, j : j + 1] = c[2 * n + 1]
 393                 elif direction == 3:
 394                     j = (j1 + j2 - 1) // 2
 395                     X[i1:i2, j : j + 1] = c[2 * n + 1]
 396                     if n == nb_rec - 1:
 397                         f_X[i1:i2, j1 : j + 1] = c[2 * n + 1]
 398                     else:
 399                         f_X[i1:i2, j : j + 1] = c[2 * n + 1]
 400
 401     def task_frame(self, A, f_A, B, f_B):
 402         nb_rec = 3
 403         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
 404         for X, f_X in [(A, f_A), (B, f_B)]:
 405             r = self.rec_coo(nb_rec)
 406             for n in range(nb_rec):
 407                 i1, j1, i2, j2 = r[n]
 408                 X[i1:i2, j1:j2] = c[n]
 409                 f_X[i1:i2, j1:j2] = c[n]
 410                 if n == nb_rec - 1:
 411                     f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = 0
 412
 413     def task_detect(self, A, f_A, B, f_B):
 414         nb_rec = 3
 415         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
 416         for X, f_X in [(A, f_A), (B, f_B)]:
 417             r = self.rec_coo(nb_rec)
 418             for n in range(nb_rec):
 419                 i1, j1, i2, j2 = r[n]
 420                 X[i1:i2, j1:j2] = c[n]
 421                 if n < nb_rec - 1:
 422                     f_X[i1, j1] = c[-1]
 423
 424     def contact(self, X, i, j, q):
 425         nq, nq_diag = 0, 0
 426         no = 0
 427
 428         for ii, jj in [
 429             (i - 1, j - 1),
 430             (i - 1, j),
 431             (i - 1, j + 1),
 432             (i, j - 1),
 433             (i, j + 1),
 434             (i + 1, j - 1),
 435             (i + 1, j),
 436             (i + 1, j + 1),
 437         ]:
 438             if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
 439                 if X[ii, jj] != 0 and X[ii, jj] != q:
 440                     no += 1
 441
 442         for ii, jj in [
 443             (i - 1, j - 1),
 444             (i - 1, j + 1),
 445             (i + 1, j - 1),
 446             (i + 1, j + 1),
 447         ]:
 448             if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
 449                 if X[ii, jj] == q and X[i, jj] != q and X[ii, j] != q:
 450                     nq_diag += 1
 451
 452         for ii, jj in [(i - 1, j), (i, j - 1), (i, j + 1), (i + 1, j)]:
 453             if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
 454                 if X[ii, jj] == q:
 455                     nq += 1
 456
 457         return no, nq, nq_diag
 458
 459     def task_count(self, A, f_A, B, f_B):
 460         N = torch.randint(4, (1,)) + 2
 461         c = torch.randperm(len(self.colors) - 1)[:N] + 1
 462
 463         for X, f_X in [(A, f_A), (B, f_B)]:
 464             nb = torch.zeros(N, dtype=torch.int64)
 465             q = torch.randint(N, (self.height * self.width,))
 466             k = torch.randperm(self.height * self.width)
 467             for p in range(self.height * self.width):
 468                 i, j = k[p] % self.height, k[p] // self.height
 469                 no, nq, nq_diag = self.contact(X, i, j, c[q[p]])
 470                 if no == 0 and nq_diag == 0:
 471                     if nq == 0:
 472                         if nb[q[p]] < self.width:
 473                             X[i, j] = c[q[p]]
 474                             nb[q[p]] += 1
 475                     if nq == 1:
 476                         X[i, j] = c[q[p]]
 477
 478             for n in range(N):
 479                 for j in range(nb[n]):
 480                     f_X[n, j] = c[n]
 481
 482     def task_trajectory(self, A, f_A, B, f_B):
 483         c = torch.randperm(len(self.colors) - 1)[:2] + 1
 484         for X, f_X in [(A, f_A), (B, f_B)]:
 485             while True:
 486                 di, dj = torch.randint(7, (2,)) - 3
 487                 i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
 488                 if (
 489                     abs(di) + abs(dj) > 0
 490                     and i + 2 * di >= 0
 491                     and i + 2 * di < self.height
 492                     and j + 2 * dj >= 0
 493                     and j + 2 * dj < self.width
 494                 ):
 495                     break
 496
 497             k = 0
 498             while (
 499                 i + k * di >= 0
 500                 and i + k * di < self.height
 501                 and j + k * dj >= 0
 502                 and j + k * dj < self.width
 503             ):
 504                 if k < 2:
 505                     X[i + k * di, j + k * dj] = c[k]
 506                 f_X[i + k * di, j + k * dj] = c[min(k, 1)]
 507                 k += 1
 508
 509     def task_bounce(self, A, f_A, B, f_B):
 510         c = torch.randperm(len(self.colors) - 1)[:3] + 1
 511         for X, f_X in [(A, f_A), (B, f_B)]:
 512
 513             def free(i, j):
 514                 return (
 515                     i >= 0
 516                     and i < self.height
 517                     and j >= 0
 518                     and j < self.width
 519                     and f_X[i, j] == 0
 520                 )
 521
 522             while True:
 523                 f_X[...] = 0
 524                 X[...] = 0
 525
 526                 for _ in range((self.height * self.width) // 10):
 527                     i, j = torch.randint(self.height, (1,)), torch.randint(
 528                         self.width, (1,)
 529                     )
 530                     X[i, j] = c[0]
 531                     f_X[i, j] = c[0]
 532
 533                 while True:
 534                     di, dj = torch.randint(7, (2,)) - 3
 535                     if abs(di) + abs(dj) == 1:
 536                         break
 537
 538                 i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
 539
 540                 X[i, j] = c[1]
 541                 f_X[i, j] = c[1]
 542                 l = 0
 543
 544                 while True:
 545                     l += 1
 546                     if free(i + di, j + dj):
 547                         pass
 548                     elif free(i - dj, j + di):
 549                         di, dj = -dj, di
 550                         if free(i + dj, j - di):
 551                             if torch.rand(1) < 0.5:
 552                                 di, dj = -di, -dj
 553                     elif free(i + dj, j - di):
 554                         di, dj = dj, -di
 555                     else:
 556                         break
 557
 558                     i, j = i + di, j + dj
 559                     f_X[i, j] = c[2]
 560                     if l <= 1:
 561                         X[i, j] = c[2]
 562
 563                     if l >= self.width:
 564                         break
 565
 566                 f_X[i, j] = c[1]
 567                 X[i, j] = c[1]
 568
 569                 if l > 3:
 570                     break
 571
 572     def task_scale(self, A, f_A, B, f_B):
 573         c = torch.randperm(len(self.colors) - 1)[:2] + 1
 574
 575         i, j = torch.randint(self.height // 2, (1,)), torch.randint(
 576             self.width // 2, (1,)
 577         )
 578
 579         for X, f_X in [(A, f_A), (B, f_B)]:
 580             for _ in range(3):
 581                 while True:
 582                     i1, j1 = torch.randint(self.height // 2 + 1, (1,)), torch.randint(
 583                         self.width // 2 + 1, (1,)
 584                     )
 585                     i2, j2 = torch.randint(self.height // 2 + 1, (1,)), torch.randint(
 586                         self.width // 2 + 1, (1,)
 587                     )
 588                     if i1 < i2 and j1 < j2 and min(i2 - i1, j2 - j1) <= 3:
 589                         break
 590                 X[i + i1 : i + i2, j + j1 : j + j2] = c[0]
 591                 f_X[2 * i1 : 2 * i2, 2 * j1 : 2 * j2] = c[0]
 592
 593             X[i, j] = c[1]
 594             f_X[0:2, 0:2] = c[1]
 595
 596     def task_islands(self, A, f_A, B, f_B):
 597         for X, f_X in [(A, f_A), (B, f_B)]:
 598             while True:
 599                 i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
 600                 if (
 601                     i == 0
 602                     or i == self.height - 1
 603                     or j == 0
 604                     or j == self.width - 1
 605                     or X[i, j] == 1
 606                 ):
 607                     break
 608             while True:
 609                 di, dj = torch.randint(3, (2,)) - 1
 610                 if abs(di) + abs(dj) > 0:
 611                     break
 612             X[i, j] = 1
 613             while True:
 614                 i, j = i + di, j + dj
 615                 if i < 0 or i >= self.height or j < 0 or j >= self.width:
 616                     break
 617                 b = (
 618                     i == 0
 619                     or i == self.height - 1
 620                     or j == 0
 621                     or j == self.width - 1
 622                     or X[i, j] == 1
 623                 )
 624                 X[i, j] = 1
 625                 if b:
 626                     break
 627
 628     ######################################################################
 629
 630     def all_tasks(self):
 631         return [
 632             self.task_replace_color,
 633             self.task_translate,
 634             self.task_grow,
 635             self.task_color_grow,
 636             self.task_frame,
 637             self.task_detect,
 638             self.task_count,
 639             self.task_trajectory,
 640             self.task_bounce,
 641             self.task_scale,
 642             # self.task_islands,
 643         ]
 644
 645     def generate_prompts_and_answers(self, nb, tasks=None, device="cpu"):
 646         if tasks is None:
 647             tasks = self.all_tasks()
 648
 649         S = self.height * self.width
 650         prompts = torch.zeros(nb, 3 * S + 2, dtype=torch.int64)
 651         answers = torch.zeros(nb, S, dtype=torch.int64)
 652
 653         for prompt, answer in tqdm.tqdm(
 654             zip(prompts, answers),
 655             dynamic_ncols=True,
 656             desc="world generation",
 657             total=prompts.size(0),
 658         ):
 659             A = prompt[0 * (S + 1) : 0 * (S + 1) + S].view(self.height, self.width)
 660             f_A = prompt[1 * (S + 1) : 1 * (S + 1) + S].view(self.height, self.width)
 661             B = prompt[2 * (S + 1) : 2 * (S + 1) + S].view(self.height, self.width)
 662             f_B = answer.view(self.height, self.width)
 663             task = tasks[torch.randint(len(tasks), (1,))]
 664             task(A, f_A, B, f_B)
 665
 666         return prompts.flatten(1), answers.flatten(1)
 667
 668     def save_quizzes(
 669         self,
 670         result_dir,
 671         filename_prefix,
 672         prompts,
 673         answers,
 674         predicted_prompts=None,
 675         predicted_answers=None,
 676         nrow=4,
 677     ):
 678         self.save_image(
 679             result_dir,
 680             filename_prefix + ".png",
 681             prompts,
 682             answers,
 683             predicted_prompts,
 684             predicted_answers,
 685             nrow,
 686         )
 687
 688
 689 ######################################################################
 690
 691 if __name__ == "__main__":
 692     import time
 693
 694     nb = 48
 695
 696     grids = Grids()
 697
 698     for t in grids.all_tasks():
 699         # for t in [grids.task_islands]:
 700         print(t.__name__)
 701         prompts, answers = grids.generate_prompts_and_answers(nb, tasks=[t])
 702         grids.save_quizzes("/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=4)
 703
 704     exit(0)
 705
 706     nb = 72
 707
 708     start_time = time.perf_counter()
 709     prompts, answers = grids.generate_prompts_and_answers(nb)
 710     delay = time.perf_counter() - start_time
 711     print(f"{prompts.size(0)/delay:02f} seq/s")
 712
 713     m = torch.randint(2, (prompts.size(0),))
 714     predicted_prompts = m * (torch.randint(2, (prompts.size(0),)) * 2 - 1)
 715     predicted_answers = (1 - m) * (torch.randint(2, (prompts.size(0),)) * 2 - 1)
 716
 717     grids.save_quizzes(
 718         "/tmp",
 719         "test",
 720         prompts[:nb],
 721         answers[:nb],
 722         # You can add a bool to put a frame around the predicted parts
 723         predicted_prompts[:nb],
 724         predicted_answers[:nb],
 725     )