perturbing shape sizes and shape colors to increase difficulty

shapes_dataset/gen_dataset.py

@@ -4,7 +4,7 @@ import sys
 import os
 import json
 import random
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageColor
 from shape_globals import Globals
 
 SHAPES = Globals.SHAPES
@@ -33,17 +33,30 @@ def gen_image(block_array):
     cnt = 0;
     for blk in block_array:
         x1 = (cnt % 3)*100+1 + eps
-        x2 = x1+100 - 2*eps
-        y1 = cnt/3*100 +eps
-        y2 = y1+100 - 2*eps
+        x2 = x1+100 - 2*eps 
+        y1 = cnt/3*100 +eps 
+        y2 = y1+100 - 2*eps 
+        x1 +=random.randint(0, 15)
+        x2 -=random.randint(0, 15)
+        y1 +=random.randint(0, 15)
+        y2 -=random.randint(0, 15)
         if blk.shape > 0:
-
+            c1 = 255-random.randint(0, 50)
+            c2 = random.randint(0, 70)
+            c3 = random.randint(0, 70)
+            if blk.color == 0:
+                col = 'rgb(%d, %d, %d)' %(c1, c2, c3)
+            elif blk.color == 1:
+                col = 'rgb(%d, %d, %d)' %(c2, c1, c3)
+            elif blk.color == 2:
+                col = 'rgb(%d, %d, %d)' %(c2, c3, c1)
+                
             if blk.shape == 1: # square
-                draw.rectangle((x1,y1, x2, y2), fill = COLORS[blk.color], outline = 'black')                
+                draw.rectangle((x1,y1, x2, y2), fill = col, outline = 'black')                
             elif blk.shape == 2:
-                draw.polygon((x1,y1, x2, y1, (x1+x2)/2, y2), fill = COLORS[blk.color], outline = 'black')
+                draw.polygon((x1,y1, x2, y1, (x1+x2)/2, y2), fill = col, outline = 'black')
             else :
-                draw.ellipse((x1,y1, x2, y2), fill = COLORS[blk.color], outline = 'black')
+                draw.ellipse((x1,y1, x2, y2), fill = col, outline = 'black')
             #else: # diamond
                 #draw.polygon((x1,(y1+y2)/2, (x1+x2)/2, y1, x2, (y1+y2)/2, (x1+x2)/2, y2), fill = COLORS[blk.color], outline = 'black')
         cnt = cnt+1