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