a194a1cf9116cdc4a40c86898a9afae06322e9d2
[mygpt.git] / 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 ######################################################################