picoclvr.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 torch, torchvision
   9 import torch.nn.functional as F
  10
  11 color_name2rgb = {
  12     "white": [255, 255, 255],
  13     "red": [255, 0, 0],
  14     "green": [0, 128, 0],
  15     "blue": [0, 0, 255],
  16     "yellow": [255, 255, 0],
  17     "black": [0, 0, 0],
  18     "maroon": [128, 0, 0],
  19     "dark_red": [139, 0, 0],
  20     "brown": [165, 42, 42],
  21     "firebrick": [178, 34, 34],
  22     "crimson": [220, 20, 60],
  23     "tomato": [255, 99, 71],
  24     "coral": [255, 127, 80],
  25     "indian_red": [205, 92, 92],
  26     "light_coral": [240, 128, 128],
  27     "dark_salmon": [233, 150, 122],
  28     "salmon": [250, 128, 114],
  29     "light_salmon": [255, 160, 122],
  30     "orange_red": [255, 69, 0],
  31     "dark_orange": [255, 140, 0],
  32     "orange": [255, 165, 0],
  33     "gold": [255, 215, 0],
  34     "dark_golden_rod": [184, 134, 11],
  35     "golden_rod": [218, 165, 32],
  36     "pale_golden_rod": [238, 232, 170],
  37     "dark_khaki": [189, 183, 107],
  38     "khaki": [240, 230, 140],
  39     "olive": [128, 128, 0],
  40     "yellow_green": [154, 205, 50],
  41     "dark_olive_green": [85, 107, 47],
  42     "olive_drab": [107, 142, 35],
  43     "lawn_green": [124, 252, 0],
  44     "chartreuse": [127, 255, 0],
  45     "green_yellow": [173, 255, 47],
  46     "dark_green": [0, 100, 0],
  47     "forest_green": [34, 139, 34],
  48     "lime": [0, 255, 0],
  49     "lime_green": [50, 205, 50],
  50     "light_green": [144, 238, 144],
  51     "pale_green": [152, 251, 152],
  52     "dark_sea_green": [143, 188, 143],
  53     "medium_spring_green": [0, 250, 154],
  54     "spring_green": [0, 255, 127],
  55     "sea_green": [46, 139, 87],
  56     "medium_aqua_marine": [102, 205, 170],
  57     "medium_sea_green": [60, 179, 113],
  58     "light_sea_green": [32, 178, 170],
  59     "dark_slate_gray": [47, 79, 79],
  60     "teal": [0, 128, 128],
  61     "dark_cyan": [0, 139, 139],
  62     "aqua": [0, 255, 255],
  63     "cyan": [0, 255, 255],
  64     "light_cyan": [224, 255, 255],
  65     "dark_turquoise": [0, 206, 209],
  66     "turquoise": [64, 224, 208],
  67     "medium_turquoise": [72, 209, 204],
  68     "pale_turquoise": [175, 238, 238],
  69     "aqua_marine": [127, 255, 212],
  70     "powder_blue": [176, 224, 230],
  71     "cadet_blue": [95, 158, 160],
  72     "steel_blue": [70, 130, 180],
  73     "corn_flower_blue": [100, 149, 237],
  74     "deep_sky_blue": [0, 191, 255],
  75     "dodger_blue": [30, 144, 255],
  76     "light_blue": [173, 216, 230],
  77     "sky_blue": [135, 206, 235],
  78     "light_sky_blue": [135, 206, 250],
  79     "midnight_blue": [25, 25, 112],
  80     "navy": [0, 0, 128],
  81     "dark_blue": [0, 0, 139],
  82     "medium_blue": [0, 0, 205],
  83     "royal_blue": [65, 105, 225],
  84     "blue_violet": [138, 43, 226],
  85     "indigo": [75, 0, 130],
  86     "dark_slate_blue": [72, 61, 139],
  87     "slate_blue": [106, 90, 205],
  88     "medium_slate_blue": [123, 104, 238],
  89     "medium_purple": [147, 112, 219],
  90     "dark_magenta": [139, 0, 139],
  91     "dark_violet": [148, 0, 211],
  92     "dark_orchid": [153, 50, 204],
  93     "medium_orchid": [186, 85, 211],
  94     "purple": [128, 0, 128],
  95     "thistle": [216, 191, 216],
  96     "plum": [221, 160, 221],
  97     "violet": [238, 130, 238],
  98     "magenta": [255, 0, 255],
  99     "orchid": [218, 112, 214],
 100     "medium_violet_red": [199, 21, 133],
 101     "pale_violet_red": [219, 112, 147],
 102     "deep_pink": [255, 20, 147],
 103     "hot_pink": [255, 105, 180],
 104     "light_pink": [255, 182, 193],
 105     "pink": [255, 192, 203],
 106     "antique_white": [250, 235, 215],
 107     "beige": [245, 245, 220],
 108     "bisque": [255, 228, 196],
 109     "blanched_almond": [255, 235, 205],
 110     "wheat": [245, 222, 179],
 111     "corn_silk": [255, 248, 220],
 112     "lemon_chiffon": [255, 250, 205],
 113     "light_golden_rod_yellow": [250, 250, 210],
 114     "light_yellow": [255, 255, 224],
 115     "saddle_brown": [139, 69, 19],
 116     "sienna": [160, 82, 45],
 117     "chocolate": [210, 105, 30],
 118     "peru": [205, 133, 63],
 119     "sandy_brown": [244, 164, 96],
 120     "burly_wood": [222, 184, 135],
 121     "tan": [210, 180, 140],
 122     "rosy_brown": [188, 143, 143],
 123     "moccasin": [255, 228, 181],
 124     "navajo_white": [255, 222, 173],
 125     "peach_puff": [255, 218, 185],
 126     "misty_rose": [255, 228, 225],
 127     "lavender_blush": [255, 240, 245],
 128     "linen": [250, 240, 230],
 129     "old_lace": [253, 245, 230],
 130     "papaya_whip": [255, 239, 213],
 131     "sea_shell": [255, 245, 238],
 132     "mint_cream": [245, 255, 250],
 133     "slate_gray": [112, 128, 144],
 134     "light_slate_gray": [119, 136, 153],
 135     "light_steel_blue": [176, 196, 222],
 136     "lavender": [230, 230, 250],
 137     "floral_white": [255, 250, 240],
 138     "alice_blue": [240, 248, 255],
 139     "ghost_white": [248, 248, 255],
 140     "honeydew": [240, 255, 240],
 141     "ivory": [255, 255, 240],
 142     "azure": [240, 255, 255],
 143     "snow": [255, 250, 250],
 144     "silver": [192, 192, 192],
 145     "gainsboro": [220, 220, 220],
 146     "white_smoke": [245, 245, 245],
 147 }
 148
 149 color_name2id = dict([(n, k) for k, n in enumerate(color_name2rgb.keys())])
 150 color_id2name = dict([(k, n) for k, n in enumerate(color_name2rgb.keys())])
 151
 152 ######################################################################
 153
 154
 155 def all_properties(height, width, nb_squares, square_i, square_j, square_c):
 156     s = []
 157
 158     for r, c_r in [(k, color_id2name[square_c[k].item()]) for k in range(nb_squares)]:
 159         s += [f"there is {c_r}"]
 160
 161         if square_i[r] >= height - height // 3:
 162             s += [f"{c_r} bottom"]
 163         if square_i[r] < height // 3:
 164             s += [f"{c_r} top"]
 165         if square_j[r] >= width - width // 3:
 166             s += [f"{c_r} right"]
 167         if square_j[r] < width // 3:
 168             s += [f"{c_r} left"]
 169
 170         for t, c_t in [
 171             (k, color_id2name[square_c[k].item()]) for k in range(nb_squares)
 172         ]:
 173             if square_i[r] > square_i[t]:
 174                 s += [f"{c_r} below {c_t}"]
 175             if square_i[r] < square_i[t]:
 176                 s += [f"{c_r} above {c_t}"]
 177             if square_j[r] > square_j[t]:
 178                 s += [f"{c_r} right of {c_t}"]
 179             if square_j[r] < square_j[t]:
 180                 s += [f"{c_r} left of {c_t}"]
 181
 182     return s
 183
 184
 185 ######################################################################
 186
 187 # Generates sequences
 188
 189
 190 def generate(
 191     nb,
 192     height,
 193     width,
 194     max_nb_squares=5,
 195     max_nb_properties=10,
 196     nb_colors=5,
 197     pruner=None,
 198 ):
 199     assert nb_colors >= max_nb_squares and nb_colors <= len(color_name2rgb) - 1
 200
 201     descr = []
 202
 203     for n in range(nb):
 204         nb_squares = torch.randint(max_nb_squares, (1,)) + 1
 205         square_position = torch.randperm(height * width)[:nb_squares]
 206
 207         # color 0 is white and reserved for the background
 208         square_c = torch.randperm(nb_colors)[:nb_squares] + 1
 209         square_i = square_position.div(width, rounding_mode="floor")
 210         square_j = square_position % width
 211
 212         img = torch.zeros(height * width, dtype=torch.int64)
 213         for k in range(nb_squares):
 214             img[square_position[k]] = square_c[k]
 215
 216         # generates all the true properties
 217
 218         s = all_properties(height, width, nb_squares, square_i, square_j, square_c)
 219
 220         if pruner is not None:
 221             s = list(filter(pruner, s))
 222
 223         # picks at most max_nb_properties at random
 224
 225         nb_properties = torch.randint(max_nb_properties, (1,)) + 1
 226         s = (
 227             " <sep> ".join([s[k] for k in torch.randperm(len(s))[:nb_properties]])
 228             + " <img> "
 229             + " ".join([f"{color_id2name[n.item()]}" for n in img])
 230         )
 231
 232         descr += [s]
 233
 234     return descr
 235
 236
 237 ######################################################################
 238
 239 # Extracts the image after <img> in descr as a 1x3xHxW tensor
 240
 241
 242 def descr2img(descr, height, width):
 243     result = []
 244
 245     def token2color(t):
 246         try:
 247             return color_name2rgb[t]
 248         except KeyError:
 249             return [128, 128, 128]
 250
 251     for d in descr:
 252         d = d.split("<img>")[1]
 253         d = d.strip().split(" ")[: height * width]
 254         d = d + ["<unk>"] * (height * width - len(d))
 255         d = [token2color(t) for t in d]
 256         img = torch.tensor(d).permute(1, 0).reshape(1, 3, height, width)
 257         result.append(img)
 258
 259     return torch.cat(result, 0)
 260
 261
 262 ######################################################################
 263
 264 # Returns all the properties of the image after <img> in descr
 265
 266
 267 def descr2properties(descr, height, width):
 268     if type(descr) == list:
 269         return [descr2properties(d, height, width) for d in descr]
 270
 271     d = descr.split("<img>")
 272     img_tokens = d[-1] if len(d) > 1 else ""
 273     img_tokens = img_tokens.strip().split(" ")[: height * width]
 274     if len(img_tokens) != height * width:
 275         return []
 276
 277     seen = {}
 278     for k, x in enumerate(img_tokens):
 279         if x != color_id2name[0]:
 280             if x in color_name2rgb:
 281                 if x in seen:
 282                     return []
 283             else:
 284                 return []
 285             seen[x] = (color_name2id[x], k // width, k % width)
 286
 287     square_infos = tuple(zip(*seen.values()))
 288
 289     if square_infos:
 290         square_c = torch.tensor(square_infos[0])
 291         square_i = torch.tensor(square_infos[1])
 292         square_j = torch.tensor(square_infos[2])
 293     else:
 294         square_c = torch.tensor([])
 295         square_i = torch.tensor([])
 296         square_j = torch.tensor([])
 297
 298     s = all_properties(height, width, len(seen), square_i, square_j, square_c)
 299
 300     return s
 301
 302
 303 ######################################################################
 304
 305 # Returns a triplet composed of (1) the total number of properties
 306 # before <img> in descr, (2) the total number of properties the image
 307 # after <img> verifies, and (3) the number of properties in (1) not in
 308 # (2)
 309
 310
 311 def nb_properties(descr, height, width, pruner=None):
 312     if type(descr) == list:
 313         return [nb_properties(d, height, width, pruner) for d in descr]
 314
 315     d = descr.split("<img>", 1)
 316     if len(d) == 0:
 317         return 0
 318     d = d[0].strip().split("<sep>")
 319     d = [x.strip() for x in d]
 320
 321     all_properties = set(descr2properties(descr, height, width))
 322
 323     if pruner is None:
 324         requested_properties = set(d)
 325     else:
 326         requested_properties = set(filter(pruner, d))
 327
 328     missing_properties = requested_properties - all_properties
 329
 330     return (len(requested_properties), len(all_properties), len(missing_properties))
 331
 332
 333 ######################################################################
 334
 335 if __name__ == "__main__":
 336     for n in range(16):
 337         descr = generate(nb=1, height=12, width=16)
 338
 339         print(nb_properties(descr, height=12, width=16))
 340
 341         with open(f"picoclvr_example_{n:02d}.txt", "w") as f:
 342             for d in descr:
 343                 f.write(f"{d}\n\n")
 344
 345         img = descr2img(descr, height=12, width=16)
 346         if img.size(0) == 1:
 347             img = F.pad(img, (1, 1, 1, 1), value=64)
 348
 349         torchvision.utils.save_image(
 350             img / 255.0,
 351             f"picoclvr_example_{n:02d}.png",
 352             padding=1,
 353             nrow=4,
 354             pad_value=0.8,
 355         )
 356
 357     import time
 358
 359     start_time = time.perf_counter()
 360     descr = generate(nb=1000, height=12, width=16)
 361     end_time = time.perf_counter()
 362     print(f"{len(descr) / (end_time - start_time):.02f} samples per second")
 363
 364 ######################################################################