lang.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 Lang(problem.Problem):
  21     named_colors = [
  22         ("white", [255, 255, 255]),
  23         ("red", [255, 0, 0]),
  24         ("green", [0, 192, 0]),
  25         ("blue", [0, 0, 255]),
  26         ("orange", [255, 192, 0]),
  27         ("cyan", [0, 255, 255]),
  28         ("violet", [255, 0, 255]),
  29         ("lightgreen", [192, 255, 192]),
  30         ("pink", [255, 192, 192]),
  31         ("lightblue", [192, 192, 255]),
  32         ("gray", [192, 192, 192]),
  33     ]
  34
  35     def __init__(
  36         self,
  37         nb_iterations=2,
  38     ):
  39         self.colors = torch.tensor([c for _, c in self.named_colors])
  40         self.name2color = dict([(p[0], i) for i, p in enumerate(self.named_colors)])
  41         self.height = 10
  42         self.width = 10
  43         self.nb_iterations = nb_iterations
  44
  45     ######################################################################
  46
  47     def frame2img(self, x, scale=15):
  48         x = x.reshape(x.size(0), self.height, -1)
  49         x = self.colors[x].permute(0, 3, 1, 2)
  50         s = x.shape
  51         x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
  52         x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
  53
  54         x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
  55         x[:, :, torch.arange(0, x.size(2), scale), :] = 0
  56         x = x[:, :, 1:, 1:]
  57
  58         return x
  59
  60     def save_image(
  61         self,
  62         result_dir,
  63         filename,
  64         prompts,
  65         answers,
  66         predicted_prompts=None,
  67         predicted_answers=None,
  68     ):
  69         if predicted_prompts is None:
  70             predicted_prompts = 255
  71
  72         if predicted_answers is None:
  73             predicted_answers = 255
  74
  75         def add_frame(x, c, margin, bottom=False):
  76             print(f"{type(x)=} {type(c)=}")
  77             if bottom:
  78                 h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0
  79             else:
  80                 h, w, di, dj = (
  81                     x.size(2) + 2 * margin,
  82                     x.size(3) + 2 * margin,
  83                     margin,
  84                     margin,
  85                 )
  86
  87             y = x.new_full((x.size(0), x.size(1), h, w), 0)
  88
  89             if type(c) is int:
  90                 y[...] = c
  91             else:
  92                 c = c.long()[:, None]
  93                 c = c * torch.tensor([192, 192, 192], device=c.device) + (
  94                     1 - c
  95                 ) * torch.tensor([255, 255, 255], device=c.device)
  96                 y[...] = c[:, :, None, None]
  97
  98             y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x
  99
 100             return y
 101
 102         margin = 8
 103
 104         img_prompts = torch.cat(
 105             [
 106                 add_frame(
 107                     add_frame(self.frame2img(x), c=0, margin=1),
 108                     c=predicted_prompts,
 109                     margin=margin,
 110                 )
 111                 for x in prompts.to("cpu").split(split_size=self.width, dim=2)
 112             ],
 113             dim=3,
 114         )
 115
 116         h = img_prompts.size(2)
 117         img_answers = add_frame(
 118             add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1),
 119             c=predicted_answers,
 120             margin=margin,
 121         )
 122
 123         separator_size = 2 * margin
 124
 125         separator = img_prompts.new_full(
 126             (
 127                 img_prompts.size(0),
 128                 img_prompts.size(1),
 129                 img_prompts.size(2),
 130                 separator_size,
 131             ),
 132             255,
 133         )
 134
 135         marker = img_prompts.new_full(
 136             (
 137                 img_prompts.size(0),
 138                 img_prompts.size(1),
 139                 img_prompts.size(2),
 140                 separator_size,
 141             ),
 142             255,
 143         )
 144
 145         # marker[:, :, 0] = 0
 146         # marker[:, :, h - 1] = 0
 147
 148         for k in range(1, 2 * separator_size - 8):
 149             i = k - (separator_size - 4)
 150             j = separator_size - 5 - abs(i)
 151             marker[:, :, h // 2 - 1 + i, 2 + j] = 0
 152             marker[:, :, h // 2 - 1 + i + 1, 2 + j] = 0
 153
 154         img = torch.cat(
 155             [
 156                 img_prompts,
 157                 marker,
 158                 img_answers,
 159             ],
 160             dim=3,
 161         )
 162
 163         image_name = os.path.join(result_dir, filename)
 164         torchvision.utils.save_image(
 165             img.float() / 255.0, image_name, nrow=4, padding=margin * 4, pad_value=1.0
 166         )
 167
 168     ######################################################################
 169
 170     def nb_token_values(self):
 171         return len(self.colors)
 172
 173     def rec_coo(self, x):
 174         while True:
 175             i1, i2 = torch.randint(x.size(0), (2,))
 176             if i1 < i2 - 1:
 177                 break
 178         while True:
 179             j1, j2 = torch.randint(x.size(1), (2,))
 180             if j1 < j2 - 1:
 181                 break
 182         return i1, j1, i2, j2
 183
 184     def task_red_to_green(self, A, f_A, B, f_B):
 185         i1, j1, i2, j2 = self.rec_coo(A)
 186         A[i1:i2, j1:j2] = self.name2color["red"]
 187         f_A[i1:i2, j1:j2] = self.name2color["green"]
 188         i1, j1, i2, j2 = self.rec_coo(B)
 189         B[i1:i2, j1:j2] = self.name2color["red"]
 190         f_B[i1:i2, j1:j2] = self.name2color["green"]
 191
 192     def generate_prompts_and_answers(self, nb):
 193         prompts = torch.zeros(nb, self.height, self.width * 3, dtype=torch.int64)
 194         answers = torch.zeros(nb, self.height, self.width, dtype=torch.int64)
 195         w = self.width
 196         for prompt, answer in zip(prompts, answers):
 197             self.task_red_to_green(
 198                 prompt[:, 0 * w : 1 * w],
 199                 prompt[:, 1 * w : 2 * w],
 200                 prompt[:, 2 * w : 3 * w],
 201                 answer,
 202             )
 203         return prompts, answers
 204
 205     def save_quizzes(
 206         self,
 207         result_dir,
 208         filename_prefix,
 209         prompts,
 210         answers,
 211         predicted_prompts=None,
 212         predicted_answers=None,
 213     ):
 214         self.save_image(
 215             result_dir,
 216             filename_prefix + ".png",
 217             prompts,
 218             answers,
 219             predicted_prompts,
 220             predicted_answers,
 221         )
 222
 223
 224 ######################################################################
 225
 226 if __name__ == "__main__":
 227     import time
 228
 229     lang = Lang(nb_iterations=4)
 230
 231     prompts, answers = lang.generate_prompts_and_answers(24)
 232
 233     predicted_prompts = torch.rand(prompts.size(0)) < 0.5
 234     predicted_answers = torch.logical_not(predicted_prompts)
 235
 236     lang.save_quizzes(
 237         "/tmp", "test", prompts, answers, predicted_prompts, predicted_answers
 238     )
 239
 240     # start_time = time.perf_counter()
 241     # token_sequences = lang.generate_token_sequences(nb=64)
 242     # delay = time.perf_counter() - start_time
 243     # print(f"{token_sequences.size(0)/delay:02f} seq/s")
 244
 245     # print(lang.seq2str(seq[:4]))
 246
 247     # for t in range(len(it[0])):
 248     # img = torch.cat([lang.frame2img(f[t]) for f in it], dim=0)
 249     # torchvision.utils.save_image(
 250     # img.float() / 255.0,
 251     # f"/tmp/frame_{t:03d}.png",
 252     # nrow=8,
 253     # padding=6,
 254     # pad_value=0,
 255     # )
 256
 257     # m = (torch.rand(seq.size()) < 0.05).long()
 258     # seq = (1 - m) * seq + m * 23
 259
 260     # print(seq.size())
 261     # img = lang.seq2img(token_sequences)
 262     # print(img.size())
 263
 264     # torchvision.utils.save_image(
 265     # img.float() / 255.0, "/tmp/world.png", nrow=6, padding=6, pad_value=0
 266     # )