Update.

diff --git a/Makefile b/Makefile
index 3483558..f639424 100644 (file)
--- a/Makefile
+++ b/Makefile
@@ -34,9 +34,9 @@ all: flatland_generator.so TAGS
 TAGS: *.cc *.h
        etags *.cc *.h
 
-flatland_generator.so: sequence_generator.o misc.o \
-       universe.o \
-       polygon.o \
+flatland_generator.so: \
+       sequence_generator.o misc.o \
+       polygon.o universe.o \
        canvas.o canvas_cairo.o
        $(CXX) $(CXXFLAGS) $(LDFLAGS) -shared -fPIC -o $@ $^
 

diff --git a/flatland.c b/flatland.c
index 77f89aa..a77d7c1 100644 (file)
--- a/flatland.c
+++ b/flatland.c
@@ -21,6 +21,7 @@
 THByteTensor *generate_sequence(long nb_sequences,
                                 long nb_images,
                                 long image_height, long image_width,
+                                int nb_shapes,
                                 int random_shape_size, int random_colors) {
 
   long nb_channels = 3;
@@ -46,7 +47,7 @@ THByteTensor *generate_sequence(long nb_sequences,
 
   for(s = 0; s < nb_sequences; s++) {
     a = THByteTensor_storage(result)->data + THByteTensor_storageOffset(result) + s * st0;
-    fl_generate_sequence(nb_images, image_width, image_height,
+    fl_generate_sequence(nb_images, image_width, image_height, nb_shapes,
                          random_shape_size, random_colors,
                          tmp_buffer);
     unsigned char *r = tmp_buffer;

diff --git a/flatland.h b/flatland.h
index e86533e..ba91f67 100644 (file)
--- a/flatland.h
+++ b/flatland.h
@@ -2,4 +2,5 @@
 THByteTensor *generate_sequence(long nb_sequences,
                                 long nb_images,
                                 long image_height, long image_width,
+                                int nb_shapes,
                                 int random_shape_size, int random_colors);

diff --git a/sequence_generator.cc b/sequence_generator.cc
index 0ebce0d..ddc337b 100644 (file)
--- a/sequence_generator.cc
+++ b/sequence_generator.cc
@@ -67,12 +67,14 @@ void draw_grabbing_point_on_canvas(CanvasCairo *canvas, scalar_t scaling,
 
 extern "C" void fl_generate_sequence(int nb_images,
                                      int width, int height,
+                                     int nb_shapes,
                                      int random_shape_size, int random_colors,
                                      unsigned char *output) {
 
-  const scalar_t world_width = width * 8;
-  const scalar_t world_height = height * 8;
-  const scalar_t scaling = 0.125;
+  const scalar_t super_definition = 8;
+  const scalar_t world_width = width * super_definition;
+  const scalar_t world_height = height * super_definition;
+  const scalar_t scaling = 1 / super_definition;
 
   const scalar_t dt = 0.1;
   const int nb_iterations_per_steps = 5;
@@ -87,7 +89,6 @@ extern "C" void fl_generate_sequence(int nb_images,
   int every_nth = 16;
   int nb_simulated_frames = 1 + (nb_images - 1) * every_nth;
   int random_grasp = 1;
-  int nb_shapes = 10;
 
   Universe *universe;
   Polygon *grabbed_polygon;
@@ -110,6 +111,9 @@ extern "C" void fl_generate_sequence(int nb_images,
 
   int failed;
 
+  int total_nb_attempts = 0;
+  const int max_total_nb_attempts = 1000000;
+
   do {
     if(random_grasp) {
       grab_start_x = world_width * (0.1 + 0.8 * drand48());
@@ -185,6 +189,7 @@ extern "C" void fl_generate_sequence(int nb_images,
       }
 
       grabbed_polygon = universe->pick_polygon(grab_start_x, grab_start_y);
+
     } while(!grabbed_polygon);
 
     failed = 0;
@@ -226,6 +231,14 @@ extern "C" void fl_generate_sequence(int nb_images,
         }
       }
     }
+
+    total_nb_attempts++;
+
+    if(total_nb_attempts >= max_total_nb_attempts) {
+      cerr << "There was " << max_total_nb_attempts << " attempts at generating the sequences." << endl;
+      abort();
+    }
+
   } while(failed);
 
   for(int t = 0; t < nb_images; t++) {

diff --git a/sequence_generator.h b/sequence_generator.h
index 76f1d33..7bfc7ea 100644 (file)
--- a/sequence_generator.h
+++ b/sequence_generator.h
@@ -5,6 +5,7 @@ extern "C" {
 
 void fl_generate_sequence(int nb_images,
                           int width, int height,
+                          int nb_shapes,
                           int random_shape_size, int random_colors,
                           unsigned char *output);
 

diff --git a/test.py b/test.py
index 0a065fa..ac3eb7e 100755 (executable)
--- a/test.py
+++ b/test.py
@@ -2,13 +2,57 @@
 
 import torch
 import torchvision
-from torchvision import datasets
+import argparse
 
 from _ext import flatland
 
 ######################################################################
 
-def sequences_to_image(x, gap=1, gap_color = (0, 128, 255)):
+parser = argparse.ArgumentParser(
+    description='Dummy test of the flatland sequence generation.',
+    formatter_class=argparse.ArgumentDefaultsHelpFormatter
+)
+
+parser.add_argument('--seed',
+                    type = int, default = 0,
+                    help = 'Random seed, < 0 is no seeding')
+
+parser.add_argument('--width',
+                    type = int, default = 80,
+                    help = 'Image width')
+
+parser.add_argument('--height',
+                    type = int, default = 80,
+                    help = 'Image height')
+
+parser.add_argument('--nb_shapes',
+                    type = int, default = 10,
+                    help = 'Image height')
+
+parser.add_argument('--nb_sequences',
+                    type = int, default = 1,
+                    help = 'How many sequences to generate')
+
+parser.add_argument('--nb_images_per_sequences',
+                    type = int, default = 3,
+                    help = 'How many images per sequence')
+
+parser.add_argument('--randomize_colors',
+                    action='store_true', default=False,
+                    help = 'Should the shapes be of different colors')
+
+parser.add_argument('--randomize_shape_size',
+                    action='store_true', default=False,
+                    help = 'Should the shapes be of different size')
+
+args = parser.parse_args()
+
+if args.seed >= 0:
+    torch.manual_seed(args.seed)
+
+######################################################################
+
+def sequences_to_image(x, gap = 1, gap_color = (0, 128, 255)):
     from PIL import Image
 
     nb_sequences = x.size(0)
@@ -26,9 +70,7 @@ def sequences_to_image(x, gap=1, gap_color = (0, 128, 255)):
                               gap + nb_sequences * (height + gap),
                               gap + nb_images_per_sequences * (width + gap))
 
-    result[0].fill_(gap_color[0])
-    result[1].fill_(gap_color[1])
-    result[2].fill_(gap_color[2])
+    result.copy_(torch.Tensor(gap_color).view(-1, 1, 1).expand_as(result))
 
     for s in range(0, nb_sequences):
         for i in range(0, nb_images_per_sequences):
@@ -42,6 +84,11 @@ def sequences_to_image(x, gap=1, gap_color = (0, 128, 255)):
 
 ######################################################################
 
-x = flatland.generate_sequence(10, 6, 80, 80, True, True)
+x = flatland.generate_sequence(args.nb_sequences,
+                               args.nb_images_per_sequences,
+                               args.height, args.width,
+                               args.nb_shapes,
+                               args.randomize_colors,
+                               args.randomize_shape_size)
 
-sequences_to_image(x, gap = 2, gap_color = (0, 0, 0)).save('sequences.png')
+sequences_to_image(x, gap = 1, gap_color = (0, 0, 0)).save('sequences.png')