cd4d313c423f7c40ed75916b52202ccd8420ee7c
[flatland.git] / sequence_generator.cc
1
2 /*
3  *  dyncnn is a deep-learning algorithm for the prediction of
4  *  interacting object dynamics
5  *
6  *  Copyright (c) 2016 Idiap Research Institute, http://www.idiap.ch/
7  *  Written by Francois Fleuret <francois.fleuret@idiap.ch>
8  *
9  *  This file is part of dyncnn.
10  *
11  *  dyncnn is free software: you can redistribute it and/or modify it
12  *  under the terms of the GNU General Public License version 3 as
13  *  published by the Free Software Foundation.
14  *
15  *  dyncnn is distributed in the hope that it will be useful, but
16  *  WITHOUT ANY WARRANTY; without even the implied warranty of
17  *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
18  *  General Public License for more details.
19  *
20  *  You should have received a copy of the GNU General Public License
21  *  along with dyncnn.  If not, see <http://www.gnu.org/licenses/>.
22  *
23  */
24
25 #include <iostream>
26 #include <fstream>
27 #include <cmath>
28 #include <stdio.h>
29 #include <stdlib.h>
30 #include <stdint.h>
31 #include <errno.h>
32 #include <string.h>
33 #include <sys/stat.h>
34 #include <sys/time.h>
35
36 using namespace std;
37
38 #include "misc.h"
39 #include "universe.h"
40 #include "canvas_cairo.h"
41
42 //////////////////////////////////////////////////////////////////////
43
44 void draw_universe_on_canvas(CanvasCairo *canvas, scalar_t scaling,
45                              Universe *universe) {
46   canvas->set_line_width(1.0 / scaling);
47   universe->draw(canvas);
48 }
49
50 void draw_grabbing_point_on_canvas(CanvasCairo *canvas, scalar_t scaling,
51                                    scalar_t xg, scalar_t yg,
52                                    scalar_t r, scalar_t g, scalar_t b) {
53   scalar_t radius = 1/scaling;
54   int n = 36;
55   scalar_t xp[n], yp[n];
56   for(int k = 0; k < n; k++) {
57     scalar_t alpha = 2 * M_PI * scalar_t(k) / scalar_t(n);
58     xp[k] = xg + radius * cos(alpha);
59     yp[k] = yg + radius * sin(alpha);
60   }
61   canvas->set_drawing_color(r, g, b);
62   canvas->set_line_width(2.0);
63   canvas->draw_polygon(1, n, xp, yp);
64 }
65
66 //////////////////////////////////////////////////////////////////////
67
68 extern "C" void fl_generate_sequence(int nb_images,
69                                      int width, int height,
70                                      int nb_shapes,
71                                      int random_shape_size, int random_colors,
72                                      int pulling,
73                                      unsigned char *output) {
74
75   const scalar_t super_definition = 8;
76   const scalar_t world_width = width * super_definition;
77   const scalar_t world_height = height * super_definition;
78   const scalar_t scaling = 1 / super_definition;
79
80   const scalar_t dt = 0.1;
81   const int nb_iterations_per_steps = 5;
82
83   //////////////////////////////////////////////////////////////////////
84
85   // We will generate images { 0, every_nth, 2 * every_nth, ..., k * every_nth < nb_simulated_frames }
86
87   // The framerate every_nth may be set to smaller value to generate
88   // nice materials for presentations or papers.
89
90   int every_nth = 16;
91   int nb_simulated_frames = 1 + (nb_images - 1) * every_nth;
92   int random_grasp = 1;
93
94   Universe *universe;
95   Polygon *grabbed_polygon;
96
97   universe = new Universe(nb_shapes, world_width, world_height);
98
99   const int nb_saved_frames = (nb_simulated_frames + every_nth - 1) / every_nth;
100   if(nb_saved_frames != nb_images) {
101     cerr << "It makes no sense." << endl;
102     abort();
103   }
104
105   CanvasCairo *canvases[nb_saved_frames * 2];
106
107   for(int s = 0; s < 2 * nb_saved_frames; s++) {
108     canvases[s] = new CanvasCairo(scaling, universe->width(), universe->height());
109   }
110
111   scalar_t gravity_fx = 0.0;
112   scalar_t gravity_fy = 1.0;
113
114   scalar_t grab_start_x, grab_start_y;
115
116   int failed;
117
118   int total_nb_attempts = 0;
119   const int max_total_nb_attempts = 1000000;
120
121   do {
122     if(pulling) {
123       if(random_grasp) {
124         grab_start_x = world_width * (0.1 + 0.8 * drand48());
125         grab_start_y = world_height * (0.1 + 0.8 * drand48());
126       } else {
127         grab_start_x = world_width * 0.5;
128         grab_start_y = world_height * 0.75;
129       }
130     }
131
132     do {
133       universe->clear();
134
135       const int nb_attempts_max = 100;
136       int nb_attempts = 0;
137
138       for(int u = 0; u < nb_shapes; u++) {
139         Polygon *pol = 0;
140
141         nb_attempts = 0;
142
143         scalar_t shape_size;
144
145         if(random_shape_size) {
146           shape_size = 40 + 80 * drand48();
147         } else {
148           shape_size = 80;
149         }
150
151         scalar_t red, green, blue;
152
153         if(random_colors) {
154           do {
155             red = drand48();
156             green = drand48();
157             blue = drand48();
158           } while(red < 0.9 and green < 0.9 and blue < 0.9 and
159                   red > 0.1 and green > 0.1 and blue > 0.1);
160         } else {
161           red = 1.0;
162           green = 1.0;
163           blue = 1.0;
164         }
165
166         do {
167           scalar_t x[] = { - shape_size * 0.4, + shape_size * 0.4,
168                            + shape_size * 0.4, - shape_size * 0.4 };
169
170           scalar_t y[] = { - shape_size * 0.6, - shape_size * 0.6,
171                            + shape_size * 0.6, + shape_size * 0.6 };
172
173           scalar_t object_center_x = world_width * drand48();
174           scalar_t object_center_y = world_height * drand48();
175
176           delete pol;
177           pol = new Polygon(0.5, red, green, blue, x, y, sizeof(x) / sizeof(scalar_t));
178           pol->set_position(object_center_x, object_center_y, M_PI * 2 * drand48());
179           pol->set_speed(0, 0, 0);
180
181           universe->initialize_polygon(pol);
182
183           nb_attempts++;
184         } while(nb_attempts < nb_attempts_max &&
185                 (universe->collide(pol) || universe->collide_with_borders(pol, 2.0 / scaling)));
186
187         if(nb_attempts == nb_attempts_max) {
188           delete pol;
189           u = -1;
190           universe->clear();
191           nb_attempts = 0;
192         } else {
193           universe->add_polygon(pol);
194         }
195       }
196
197       if(pulling) {
198         grabbed_polygon = universe->pick_polygon(grab_start_x, grab_start_y);
199       }
200     } while(pulling and !grabbed_polygon);
201
202     failed = 0;
203
204     scalar_t grab_relative_x, grab_relative_y;
205
206     if(pulling) {
207       grab_relative_x = grabbed_polygon->relative_x(grab_start_x, grab_start_y);
208       grab_relative_y = grabbed_polygon->relative_y(grab_start_x, grab_start_y);
209     }
210
211     for(int s = 0; !failed && s < nb_simulated_frames; s++) {
212       if(s % every_nth == 0) {
213         int t = s / every_nth;
214
215         canvases[t]->clear();
216         draw_universe_on_canvas(canvases[t], scaling, universe);
217       }
218
219       if(s < nb_simulated_frames - 1) {
220
221         // Run the simulation
222
223         for(int i = 0; i < nb_iterations_per_steps; i++) {
224           if(pulling) {
225             // Pulling the grabbed rectangle
226             scalar_t xf = grabbed_polygon->absolute_x(grab_relative_x, grab_relative_y);
227             scalar_t yf = grabbed_polygon->absolute_y(grab_relative_x, grab_relative_y);
228             if (xf < 0 || xf >= world_width || yf < 0 || yf >= world_height) {
229               failed = 1;
230             }
231             grabbed_polygon->apply_force(dt, xf, yf, 0.0, -1.0);
232           } else {
233             // Gravity
234             universe->apply_gravity(dt, gravity_fx, gravity_fy);
235           }
236
237           universe->update(dt, 1.0 / scaling);
238         }
239       }
240     }
241
242     total_nb_attempts++;
243
244     if(total_nb_attempts >= max_total_nb_attempts) {
245       cerr << "There was "
246            << max_total_nb_attempts
247            << " attempts at generating the sequences, aborting." << endl;
248       abort();
249     }
250
251   } while(failed);
252
253   for(int t = 0; t < nb_images; t++) {
254     unsigned char *src = canvases[t]->_data;
255     unsigned char *dst = output + t * width * height * 3;
256     for(int d = 0; d < 3; d++) {
257       for(int y = 0; y < height; y++) {
258         for(int x = 0; x < width; x++) {
259           dst[x + width * (y + height * d)] = src[d + 4 * (x + width * y)];
260         }
261       }
262     }
263   }
264
265   for(int t = 0; t < 2 * nb_saved_frames; t++) {
266     delete canvases[t];
267   }
268
269   delete universe;
270 }