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