world.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
   9
  10 import torch, torchvision
  11
  12 from torch import nn
  13 from torch.nn import functional as F
  14
  15 ######################################################################
  16
  17
  18 colors = torch.tensor(
  19     [
  20         [255, 255, 255],
  21         [255, 0, 0],
  22         [0, 128, 0],
  23         [0, 0, 255],
  24         [255, 200, 0],
  25         [192, 192, 192],
  26     ]
  27 )
  28
  29 token_background = 0
  30 first_fish_token = 1
  31 nb_fish_tokens = len(colors) - 1
  32 token_forward = first_fish_token + nb_fish_tokens
  33 token_backward = token_forward + 1
  34
  35 token2char = "_" + "".join([str(n) for n in range(len(colors) - 1)]) + "><"
  36
  37
  38 def generate(
  39     nb,
  40     height,
  41     width,
  42     max_nb_obj=2,
  43     nb_iterations=2,
  44 ):
  45     pairs = []
  46
  47     for n in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"):
  48         f_start = torch.zeros(height, width, dtype=torch.int64)
  49         f_end = torch.zeros(height, width, dtype=torch.int64)
  50         n = torch.arange(f_start.size(0))
  51
  52         nb_fish = torch.randint(max_nb_obj, (1,)).item() + 1
  53         for c in (
  54             (torch.randperm(nb_fish_tokens) + first_fish_token)[:nb_fish].sort().values
  55         ):
  56             i, j = (
  57                 torch.randint(height - 2, (1,))[0] + 1,
  58                 torch.randint(width - 2, (1,))[0] + 1,
  59             )
  60             vm = torch.randint(4, (1,))[0]
  61             vi, vj = (vm // 2) * (2 * (vm % 2) - 1), (1 - vm // 2) * (2 * (vm % 2) - 1)
  62
  63             f_start[i, j] = c
  64             f_start[i - vi, j - vj] = c
  65             f_start[i + vj, j - vi] = c
  66             f_start[i - vj, j + vi] = c
  67
  68             for l in range(nb_iterations):
  69                 i += vi
  70                 j += vj
  71                 if i < 0 or i >= height or j < 0 or j >= width:
  72                     i -= vi
  73                     j -= vj
  74                     vi, vj = -vi, -vj
  75                     i += vi
  76                     j += vj
  77
  78             f_end[i, j] = c
  79             f_end[i - vi, j - vj] = c
  80             f_end[i + vj, j - vi] = c
  81             f_end[i - vj, j + vi] = c
  82
  83         pairs.append((f_start, f_end))
  84
  85     result = []
  86     for p in pairs:
  87         if torch.rand(1) < 0.5:
  88             result.append(
  89                 torch.cat(
  90                     [p[0].flatten(), torch.tensor([token_forward]), p[1].flatten()],
  91                     dim=0,
  92                 )[None, :]
  93             )
  94         else:
  95             result.append(
  96                 torch.cat(
  97                     [p[1].flatten(), torch.tensor([token_backward]), p[0].flatten()],
  98                     dim=0,
  99                 )[None, :]
 100             )
 101
 102     return torch.cat(result, dim=0)
 103
 104
 105 def sample2img(seq, height, width, upscale=15):
 106     f_start = seq[:, : height * width].reshape(-1, height, width)
 107     f_end = seq[:, height * width + 1 :].reshape(-1, height, width)
 108
 109     def mosaic(x, upscale):
 110         x = x.reshape(-1, height, width)
 111         m = torch.logical_and(x >= 0, x < first_fish_token + nb_fish_tokens).long()
 112         x = colors[x * m].permute(0, 3, 1, 2)
 113         s = x.shape
 114         x = x[:, :, :, None, :, None].expand(-1, -1, -1, upscale, -1, upscale)
 115         x = x.reshape(s[0], s[1], s[2] * upscale, s[3] * upscale)
 116
 117         for n in range(m.size(0)):
 118             for i in range(m.size(1)):
 119                 for j in range(m.size(2)):
 120                     if m[n, i, j] == 0:
 121                         for k in range(2, upscale - 2):
 122                             x[n, :, i * upscale + k, j * upscale + k] = 0
 123                             x[n, :, i * upscale + upscale - 1 - k, j * upscale + k] = 0
 124
 125         return x
 126
 127     return torch.cat([mosaic(f_start, upscale), mosaic(f_end, upscale)], dim=3)
 128
 129
 130 def seq2str(seq):
 131     result = []
 132     for s in seq:
 133         result.append("".join([token2char[v] for v in s]))
 134     return result
 135
 136
 137 ######################################################################
 138
 139 if __name__ == "__main__":
 140     import time
 141
 142     height, width = 6, 8
 143     start_time = time.perf_counter()
 144     seq = generate(nb=64, height=height, width=width, max_nb_obj=3)
 145     delay = time.perf_counter() - start_time
 146     print(f"{seq.size(0)/delay:02f} samples/s")
 147
 148     print(seq2str(seq[:4]))
 149
 150     # m = (torch.rand(seq.size()) < 0.05).long()
 151     # seq = (1 - m) * seq + m * 23
 152
 153     img = sample2img(seq, height, width)
 154     print(img.size())
 155
 156     torchvision.utils.save_image(
 157         img.float() / 255.0, "/tmp/world.png", nrow=8, padding=2
 158     )