sky.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 Sky(problem.Problem):
  21     colors = torch.tensor(
  22         [
  23             [255, 255, 255],
  24             [255, 0, 0],
  25             [0, 192, 0],
  26             [0, 0, 255],
  27             [255, 192, 0],
  28             [0, 255, 255],
  29             [255, 0, 255],
  30             [192, 255, 192],
  31             [255, 192, 192],
  32             [192, 192, 255],
  33             [192, 192, 192],
  34         ]
  35     )
  36
  37     token_background = 0
  38     first_bird_token = 1
  39     nb_bird_tokens = colors.size(0) - 1
  40
  41     token2char = (
  42         "_" + "".join([chr(ord("A") + n) for n in range(len(colors) - 1)]) + "><"
  43     )
  44
  45     def __init__(
  46         self,
  47         height=6,
  48         width=8,
  49         nb_birds=3,
  50         speed=2,
  51         nb_iterations=2,
  52         avoid_collision=True,
  53         max_nb_cached_chunks=None,
  54         chunk_size=None,
  55         nb_threads=-1,
  56     ):
  57         super().__init__(max_nb_cached_chunks, chunk_size, nb_threads)
  58         self.height = height
  59         self.width = width
  60         self.nb_birds = nb_birds
  61         self.speed = speed
  62         self.nb_iterations = nb_iterations
  63         self.avoid_collision = avoid_collision
  64
  65     def generate_frame_sequences(self, nb):
  66         frame_sequences = []
  67
  68         for _ in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"):
  69             i, j, vi, vj = (
  70                 torch.empty(self.nb_birds, dtype=torch.int64),
  71                 torch.empty(self.nb_birds, dtype=torch.int64),
  72                 torch.empty(self.nb_birds, dtype=torch.int64),
  73                 torch.empty(self.nb_birds, dtype=torch.int64),
  74             )
  75
  76             def collision_okay():
  77                 if not self.avoid_collision:
  78                     return True
  79
  80                 count = torch.zeros(self.height, self.width, dtype=torch.int64)
  81
  82                 for n in range(self.nb_birds):
  83                     count[i[n], j[n]] += 1
  84                     count[i[n] - vi[n], j[n]] += 1
  85                     count[i[n], j[n] - vj[n]] += 1
  86
  87                 return count.max() <= 1
  88
  89             col = (
  90                 torch.randperm(self.colors.size(0) - 1)[: self.nb_birds].sort().values
  91                 + 1
  92             )
  93
  94             while True:
  95                 while True:
  96                     for n in range(self.nb_birds):
  97                         while True:
  98                             i[n] = torch.randint(self.height, (1,))
  99                             j[n] = torch.randint(self.width, (1,))
 100                             vm = torch.randint(4, (1,))
 101                             vi[n], vj[n] = (vm % 2) * 2 - 1, (vm // 2) * 2 - 1
 102                             if (
 103                                 i[n] - vi[n] >= 0
 104                                 and i[n] - vi[n] < self.height
 105                                 and j[n] - vj[n] >= 0
 106                                 and j[n] - vj[n] < self.width
 107                             ):
 108                                 break
 109
 110                     if collision_okay():
 111                         break
 112
 113                 result = torch.zeros(
 114                     self.nb_iterations * self.speed,
 115                     self.height,
 116                     self.width,
 117                     dtype=torch.int64,
 118                 )
 119
 120                 fine = torch.empty(self.nb_iterations * self.speed)
 121
 122                 t_to_keep = (
 123                     torch.arange(self.nb_iterations, device=result.device) * self.speed
 124                 )
 125
 126                 for l in range(self.nb_iterations * self.speed):
 127                     fine[l] = collision_okay()
 128                     for n in range(self.nb_birds):
 129                         c = col[n]
 130                         result[l, i[n], j[n]] = c
 131                         result[l, i[n] - vi[n], j[n]] = c
 132                         result[l, i[n], j[n] - vj[n]] = c
 133
 134                         if (i[n] == 0 and vi[n] == -1) or (
 135                             i[n] == self.height - 1 and vi[n] == 1
 136                         ):
 137                             vi[n] = -vi[n]
 138
 139                         if (j[n] == 0 and vj[n] == -1) or (
 140                             j[n] == self.width - 1 and vj[n] == 1
 141                         ):
 142                             vj[n] = -vj[n]
 143
 144                         i[n] += vi[n]
 145                         j[n] += vj[n]
 146
 147                 result = result[t_to_keep]
 148                 fine = fine[t_to_keep]
 149
 150                 if fine[-1]:
 151                     break
 152
 153             frame_sequences.append(result)
 154
 155         return frame_sequences
 156
 157     ######################################################################
 158
 159     def frame2img(self, x, scale=15):
 160         x = x.reshape(x.size(0), self.height, -1)
 161         m = torch.logical_and(
 162             x >= 0, x < self.first_bird_token + self.nb_bird_tokens
 163         ).long()
 164         x = self.colors[x * m].permute(0, 3, 1, 2)
 165         s = x.shape
 166         x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
 167         x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
 168
 169         x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
 170         x[:, :, torch.arange(0, x.size(2), scale), :] = 0
 171         x = x[:, :, 1:, 1:]
 172
 173         for n in range(m.size(0)):
 174             for i in range(m.size(1)):
 175                 for j in range(m.size(2)):
 176                     if m[n, i, j] == 0:
 177                         for k in range(2, scale - 2):
 178                             for l in [0, 1]:
 179                                 x[n, :, i * scale + k, j * scale + k - l] = 0
 180                                 x[
 181                                     n, :, i * scale + scale - 1 - k, j * scale + k - l
 182                                 ] = 0
 183
 184         return x
 185
 186     def seq2str(self, seq):
 187         result = []
 188         for s in seq:
 189             result.append("".join([self.token2char[v] for v in s]))
 190         return result
 191
 192     def save_image(
 193         self,
 194         result_dir,
 195         filename,
 196         prompts,
 197         answers,
 198         predicted_prompts=None,
 199         predicted_answers=None,
 200     ):
 201         if predicted_prompts is None:
 202             predicted_prompts = 255
 203
 204         if predicted_answers is None:
 205             predicted_answers = 255
 206
 207         def add_frame(x, c, margin, bottom=False):
 208             if bottom:
 209                 h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0
 210             else:
 211                 h, w, di, dj = (
 212                     x.size(2) + 2 * margin,
 213                     x.size(3) + 2 * margin,
 214                     margin,
 215                     margin,
 216                 )
 217
 218             y = x.new_full((x.size(0), x.size(1), h, w), 0)
 219
 220             if type(c) is int:
 221                 y[...] = c
 222             else:
 223                 c = c.long()[:, None]
 224                 c = (
 225                     (c == 1).long() * torch.tensor([0, 255, 0], device=c.device)
 226                     + (c == 0).long() * torch.tensor([255, 255, 255], device=c.device)
 227                     + (c == -1).long() * torch.tensor([255, 0, 0], device=c.device)
 228                 )
 229                 y[...] = c[:, :, None, None]
 230
 231             y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x
 232
 233             return y
 234
 235         margin = 4
 236
 237         img_prompts = add_frame(self.frame2img(prompts.to("cpu")), c=0, margin=1)
 238         h = img_prompts.size(2)
 239         img_answers = add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1)
 240
 241         img_prompts = add_frame(img_prompts, c=255, margin=margin, bottom=True)
 242         img_answers = add_frame(img_answers, c=255, margin=margin, bottom=True)
 243
 244         img_prompts = add_frame(
 245             img_prompts, c=predicted_prompts, margin=margin, bottom=True
 246         )
 247         img_answers = add_frame(
 248             img_answers, c=predicted_answers, margin=margin, bottom=True
 249         )
 250
 251         marker_size = 16
 252
 253         separator = img_prompts.new_full(
 254             (
 255                 img_prompts.size(0),
 256                 img_prompts.size(1),
 257                 img_prompts.size(2),
 258                 marker_size,
 259             ),
 260             255,
 261         )
 262
 263         separator[:, :, 0] = 0
 264         separator[:, :, h - 1] = 0
 265
 266         for k in range(1, 2 * marker_size - 8):
 267             i = k - (marker_size - 4)
 268             j = marker_size - 5 - abs(i)
 269             separator[:, :, h // 2 - 1 + i, 2 + j] = 0
 270             separator[:, :, h // 2 - 1 + i + 1, 2 + j] = 0
 271
 272         img = torch.cat([img_prompts, separator, img_answers], dim=3)
 273
 274         image_name = os.path.join(result_dir, filename)
 275         torchvision.utils.save_image(
 276             img.float() / 255.0, image_name, nrow=6, padding=margin * 4, pad_value=1.0
 277         )
 278
 279     ######################################################################
 280
 281     def nb_token_values(self):
 282         return len(self.colors)
 283
 284     def generate_prompts_and_answers(self, nb):
 285         frame_sequences = self.generate_frame_sequences(nb)
 286         frame_sequences = torch.cat([x[None] for x in frame_sequences], dim=0)
 287
 288         prompts = frame_sequences[:, : frame_sequences.size(1) // 2].flatten(1)
 289
 290         answers = frame_sequences[:, frame_sequences.size(1) // 2 :].flatten(1)
 291
 292         # warnings.warn("dirty test with longer answer", RuntimeWarning)
 293         # answers = torch.cat(
 294         # [
 295         # frame_sequences[:, frame_sequences.size(1) // 2 :],
 296         # frame_sequences[:, frame_sequences.size(1) // 2 :],
 297         # ],
 298         # dim=3,
 299         # ).flatten(1)
 300
 301         return prompts, answers
 302
 303     def save_quiz_illustrations(
 304         self,
 305         result_dir,
 306         filename_prefix,
 307         prompts,
 308         answers,
 309         predicted_prompts=None,
 310         predicted_answers=None,
 311     ):
 312         self.save_image(
 313             result_dir,
 314             filename_prefix + ".png",
 315             prompts,
 316             answers,
 317             predicted_prompts,
 318             predicted_answers,
 319         )
 320
 321
 322 ######################################################################
 323
 324 if __name__ == "__main__":
 325     import time
 326
 327     sky = Sky(height=6, width=8, speed=1, nb_iterations=4)
 328
 329     prompts, answers = sky.generate_prompts_and_answers(4)
 330
 331     predicted_prompts = torch.randint(3, (prompts.size(0),)) - 1
 332     predicted_answers = torch.randint(3, (prompts.size(0),)) - 1
 333
 334     sky.save_quiz_illustrations(
 335         "/tmp", "test", prompts, answers, predicted_prompts, predicted_answers
 336     )
 337
 338     # start_time = time.perf_counter()
 339     # token_sequences = sky.generate_token_sequences(nb=64)
 340     # delay = time.perf_counter() - start_time
 341     # print(f"{token_sequences.size(0)/delay:02f} seq/s")
 342
 343     # print(sky.seq2str(seq[:4]))
 344
 345     # for t in range(len(it[0])):
 346     # img = torch.cat([sky.frame2img(f[t]) for f in it], dim=0)
 347     # torchvision.utils.save_image(
 348     # img.float() / 255.0,
 349     # f"/tmp/frame_{t:03d}.png",
 350     # nrow=8,
 351     # padding=6,
 352     # pad_value=0,
 353     # )
 354
 355     # m = (torch.rand(seq.size()) < 0.05).long()
 356     # seq = (1 - m) * seq + m * 23
 357
 358     # print(seq.size())
 359     # img = sky.seq2img(token_sequences)
 360     # print(img.size())
 361
 362     # torchvision.utils.save_image(
 363     # img.float() / 255.0, "/tmp/world.png", nrow=6, padding=6, pad_value=0
 364     # )