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
  10 colors = [
  11     [ 255, 255, 255 ],
  12     [ 255,   0,   0 ],
  13     [   0, 255,   0 ],
  14     [   0,   0, 255 ],
  15     [ 255, 255,   0 ],
  16     [   0,   0,   0 ],
  17 ]
  18
  19 color_names = [
  20     'white',
  21     'red',
  22     'green',
  23     'blue',
  24     'yellow',
  25     'black',
  26 ]
  27
  28 color_tokens = dict( [ (n, c) for n, c in zip(color_names, colors) ] )
  29
  30 def generate(nb, height = 6, width = 8, max_nb_statements = 10):
  31
  32     descr = [ ]
  33
  34     for n in range(nb):
  35         nb = torch.randint(5, (1,)) + 1
  36         shape_position = torch.randperm(height * width)[:nb]
  37         shape_c = torch.randperm(5)[:nb] + 1
  38         shape_i = shape_position.div(width, rounding_mode = 'floor')
  39         shape_j = shape_position % width
  40
  41         img = [ 0 ] * height * width
  42         for k in range(nb): img[shape_position[k]] = shape_c[k]
  43
  44         s = [ ]
  45
  46         for r, c in [ (k, color_names[shape_c[k]]) for k in range(nb) ]:
  47             s += [ f'there is {c}' ]
  48
  49             if shape_i[r] >= height - height/4: s += [ f'{c} bottom' ]
  50             if shape_i[r] < height/4: s += [ f'{c} top' ]
  51             if shape_j[r] >= width - width/4: s += [ f'{c} right' ]
  52             if shape_j[r] < width/4: s += [ f'{c} left' ]
  53
  54             for t, d in [ (k, color_names[shape_c[k]]) for k in range(nb) ]:
  55                 if shape_i[r] > shape_i[t]: s += [ f'{c} below {d}' ]
  56                 if shape_i[r] < shape_i[t]: s += [ f'{c} above {d}' ]
  57                 if shape_j[r] > shape_j[t]: s += [ f'{c} right of {d}' ]
  58                 if shape_j[r] < shape_j[t]: s += [ f'{c} left of {d}' ]
  59
  60         nb_statements = torch.randint(max_nb_statements, (1,)) + 1
  61         s = ' <sep> '.join([ s[k] for k in torch.randperm(len(s))[:nb_statements] ] )
  62         s += ' <img> ' + ' '.join([ f'{color_names[n]}' for n in img ])
  63         descr += [ s ]
  64
  65     return descr
  66
  67 ######################################################################
  68
  69 def descr2img(descr, height = 6, width = 8):
  70
  71     def token2color(t):
  72         try:
  73             return color_tokens[t]
  74         except KeyError:
  75             return [ 128, 128, 128 ]
  76
  77     def img_descr(x):
  78         u = x.split('<img>', 1)
  79         return u[1] if len(u) > 1 else ''
  80
  81     img = torch.full((len(descr), 3, height, width), 255)
  82     d = [ img_descr(x) for x in descr ]
  83     d = [ u.strip().split(' ')[:height * width] for u in d ]
  84     d = [ u + [ '<unk>' ] * (height * width - len(u)) for u in d ]
  85     d = [ [ token2color(t) for t in u ] for u in d ]
  86     img = torch.tensor(d).permute(0, 2, 1)
  87     img = img.reshape(img.size(0), 3, height, width)
  88
  89     return img
  90
  91 ######################################################################
  92
  93 if __name__ == '__main__':
  94     descr = generate(5)
  95     img = descr2img(descr)
  96     print(descr, img.size())
  97     torchvision.utils.save_image(img / 255.,
  98                                  'example.png', nrow = 16, pad_value = 0.8)
  99
 100     import time
 101
 102     start_time = time.perf_counter()
 103     descr = generate(10000)
 104     end_time = time.perf_counter()
 105     print(f'{len(descr) / (end_time - start_time):.02f} samples per second')
 106
 107 ######################################################################