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         img_answers = add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1)
 233
 234         img_prompts = add_frame(img_prompts, c=255, margin=margin, bottom=True)
 235         img_answers = add_frame(img_answers, c=255, margin=margin, bottom=True)
 236
 237         img_prompts = add_frame(
 238             img_prompts, c=predicted_prompts, margin=margin, bottom=True
 239         )
 240         img_answers = add_frame(
 241             img_answers, c=predicted_answers, margin=margin, bottom=True
 242         )
 243
 244         marker_size = 8
 245
 246         separator = img_prompts.new_full(
 247             (
 248                 img_prompts.size(0),
 249                 img_prompts.size(1),
 250                 img_prompts.size(2),
 251                 marker_size,
 252             ),
 253             255,
 254         )
 255
 256         for k in range(2, 2 * marker_size - 3):
 257             i = k + 1 - marker_size
 258             j = marker_size - 2 - abs(k - marker_size + 1)
 259             separator[:, :, separator.size(2) // 2 + i, j] = 0
 260             separator[:, :, separator.size(2) // 2 + i + 1, j] = 0
 261
 262         img = torch.cat([img_prompts, separator, img_answers], dim=3)
 263
 264         image_name = os.path.join(result_dir, filename)
 265         torchvision.utils.save_image(
 266             img.float() / 255.0, image_name, nrow=6, padding=margin * 2, pad_value=1.0
 267         )
 268
 269     ######################################################################
 270
 271     def nb_token_values(self):
 272         return len(self.colors)
 273
 274     def generate_prompts_and_answers(self, nb):
 275         frame_sequences = self.generate_frame_sequences(nb)
 276         frame_sequences = torch.cat([x[None] for x in frame_sequences], dim=0)
 277
 278         prompts = frame_sequences[:, : frame_sequences.size(1) // 2].flatten(1)
 279
 280         answers = frame_sequences[:, frame_sequences.size(1) // 2 :].flatten(1)
 281         # warnings.warn("dirty test with longer answer", RuntimeWarning)
 282         # answers = torch.cat(
 283         # [
 284         # frame_sequences[:, frame_sequences.size(1) // 2 :],
 285         # frame_sequences[:, frame_sequences.size(1) // 2 :],
 286         # ],
 287         # dim=3,
 288         # ).flatten(1)
 289
 290         return prompts, answers
 291
 292     def save_quizzes(
 293         self,
 294         result_dir,
 295         filename_prefix,
 296         prompts,
 297         answers,
 298         predicted_prompts=None,
 299         predicted_answers=None,
 300     ):
 301         self.save_image(
 302             result_dir,
 303             filename_prefix + ".png",
 304             prompts,
 305             answers,
 306             predicted_prompts,
 307             predicted_answers,
 308         )
 309
 310
 311 ######################################################################
 312
 313 if __name__ == "__main__":
 314     import time
 315
 316     sky = Sky(height=6, width=8, speed=1, nb_iterations=4)
 317
 318     prompts, answers = sky.generate_prompts_and_answers(4)
 319
 320     predicted_prompts = torch.rand(prompts.size(0)) < 0.5
 321     predicted_answers = torch.rand(answers.size(0)) < 0.5
 322
 323     sky.save_quizzes(
 324         "/tmp", "test", prompts, answers, predicted_prompts, predicted_answers
 325     )
 326
 327     # start_time = time.perf_counter()
 328     # token_sequences = sky.generate_token_sequences(nb=64)
 329     # delay = time.perf_counter() - start_time
 330     # print(f"{token_sequences.size(0)/delay:02f} seq/s")
 331
 332     # print(sky.seq2str(seq[:4]))
 333
 334     # for t in range(len(it[0])):
 335     # img = torch.cat([sky.frame2img(f[t]) for f in it], dim=0)
 336     # torchvision.utils.save_image(
 337     # img.float() / 255.0,
 338     # f"/tmp/frame_{t:03d}.png",
 339     # nrow=8,
 340     # padding=6,
 341     # pad_value=0,
 342     # )
 343
 344     # m = (torch.rand(seq.size()) < 0.05).long()
 345     # seq = (1 - m) * seq + m * 23
 346
 347     # print(seq.size())
 348     # img = sky.seq2img(token_sequences)
 349     # print(img.size())
 350
 351     # torchvision.utils.save_image(
 352     # img.float() / 255.0, "/tmp/world.png", nrow=6, padding=6, pad_value=0
 353     # )