picoclvr.py

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