X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=picoclvr.py;h=2342016a05689be937107d391a2162a1fb5695e0;hb=e80417768332b6a9810ab3aaa252337810fed3c2;hp=4906c65d610020a84c9c874dcd5399d07ce9bf46;hpb=063e25c1e1442c406746a39220f3c3590882cf51;p=mygpt.git

diff --git a/picoclvr.py b/picoclvr.py
index 4906c65..2342016 100755
--- a/picoclvr.py
+++ b/picoclvr.py
@@ -34,18 +34,18 @@ def generate(nb, height = 6, width = 8, max_nb_statements = 10):
     descr = [ ]
 
     for n in range(nb):
-        nb = torch.randint(5, (1,)) + 1
-        shape_position = torch.randperm(height * width)[:nb]
-        shape_c = torch.randperm(5)[:nb] + 1
+        nb_shapes = torch.randint(len(color_tokens) - 1, (1,)) + 1
+        shape_position = torch.randperm(height * width)[:nb_shapes]
+        shape_c = torch.randperm(len(color_tokens) - 1)[:nb_shapes] + 1
         shape_i = shape_position.div(width, rounding_mode = 'floor')
         shape_j = shape_position % width
 
         img = [ 0 ] * height * width
-        for k in range(nb): img[shape_position[k]] = shape_c[k]
+        for k in range(nb_shapes): img[shape_position[k]] = shape_c[k]
 
         s = [ ]
 
-        for r, c in [ (k, color_names[shape_c[k]]) for k in range(nb) ]:
+        for r, c in [ (k, color_names[shape_c[k]]) for k in range(nb_shapes) ]:
             s += [ f'there is {c}' ]
 
             if shape_i[r] >= height - height//3: s += [ f'{c} bottom' ]
@@ -53,7 +53,7 @@ def generate(nb, height = 6, width = 8, max_nb_statements = 10):
             if shape_j[r] >= width - width//3: s += [ f'{c} right' ]
             if shape_j[r] < width//3: s += [ f'{c} left' ]
 
-            for t, d in [ (k, color_names[shape_c[k]]) for k in range(nb) ]:
+            for t, d in [ (k, color_names[shape_c[k]]) for k in range(nb_shapes) ]:
                 if shape_i[r] > shape_i[t]: s += [ f'{c} below {d}' ]
                 if shape_i[r] < shape_i[t]: s += [ f'{c} above {d}' ]
                 if shape_j[r] > shape_j[t]: s += [ f'{c} right of {d}' ]