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