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_first = seq[:, : height * width].reshape(-1, height, width)
 107     f_second = seq[:, height * width + 1 :].reshape(-1, height, width)
 108     direction = seq[:, height * width]
 109
 110     def mosaic(x, upscale):
 111         x = x.reshape(-1, height, width)
 112         m = torch.logical_and(x >= 0, x < first_fish_token + nb_fish_tokens).long()
 113         x = colors[x * m].permute(0, 3, 1, 2)
 114         s = x.shape
 115         x = x[:, :, :, None, :, None].expand(-1, -1, -1, upscale, -1, upscale)
 116         x = x.reshape(s[0], s[1], s[2] * upscale, s[3] * upscale)
 117
 118         for n in range(m.size(0)):
 119             for i in range(m.size(1)):
 120                 for j in range(m.size(2)):
 121                     if m[n, i, j] == 0:
 122                         for k in range(2, upscale - 2):
 123                             x[n, :, i * upscale + k, j * upscale + k] = 0
 124                             x[n, :, i * upscale + upscale - 1 - k, j * upscale + k] = 0
 125
 126         return x
 127
 128     direction_symbol = torch.full((direction.size(0), height * upscale, upscale), 0)
 129     direction_symbol = colors[direction_symbol].permute(0, 3, 1, 2)
 130     separator = torch.full((direction.size(0), 3, height * upscale, 1), 0)
 131
 132     for n in range(direction_symbol.size(0)):
 133         if direction[n] == token_forward:
 134             for k in range(upscale):
 135                 direction_symbol[
 136                     n,
 137                     :,
 138                     (height * upscale) // 2 - upscale // 2 + k,
 139                     3 + abs(k - upscale // 2),
 140                 ] = 0
 141         elif direction[n] == token_backward:
 142             for k in range(upscale):
 143                 direction_symbol[
 144                     n,
 145                     :,
 146                     (height * upscale) // 2 - upscale // 2 + k,
 147                     3 + upscale // 2 - abs(k - upscale // 2),
 148                 ] = 0
 149         else:
 150             for k in range(2, upscale - 2):
 151                 direction_symbol[
 152                     n, :, (height * upscale) // 2 - upscale // 2 + k, k
 153                 ] = 0
 154                 direction_symbol[
 155                     n, :, (height * upscale) // 2 - upscale // 2 + k, upscale - 1 - k
 156                 ] = 0
 157
 158     return torch.cat(
 159         [
 160             mosaic(f_first, upscale),
 161             separator,
 162             direction_symbol,
 163             separator,
 164             mosaic(f_second, upscale),
 165         ],
 166         dim=3,
 167     )
 168
 169
 170 def seq2str(seq):
 171     result = []
 172     for s in seq:
 173         result.append("".join([token2char[v] for v in s]))
 174     return result
 175
 176
 177 ######################################################################
 178
 179 if __name__ == "__main__":
 180     import time
 181
 182     height, width = 6, 8
 183     start_time = time.perf_counter()
 184     seq = generate(nb=90, height=height, width=width, max_nb_obj=3)
 185     delay = time.perf_counter() - start_time
 186     print(f"{seq.size(0)/delay:02f} samples/s")
 187
 188     print(seq2str(seq[:4]))
 189
 190     # m = (torch.rand(seq.size()) < 0.05).long()
 191     # seq = (1 - m) * seq + m * 23
 192
 193     img = sample2img(seq, height, width)
 194     print(img.size())
 195
 196     torchvision.utils.save_image(
 197         img.float() / 255.0, "/tmp/world.png", nrow=6, padding=4
 198     )