bfba036621e9e4e8896ceeb6772d5af696736429
[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 grab_start_x, grab_start_y;
112
113   int failed;
114
115   int total_nb_attempts = 0;
116   const int max_total_nb_attempts = 1000000;
117
118   do {
119     if(pulling) {
120       if(random_grasp) {
121         grab_start_x = world_width * (0.1 + 0.8 * drand48());
122         grab_start_y = world_height * (0.1 + 0.8 * drand48());
123       } else {
124         grab_start_x = world_width * 0.5;
125         grab_start_y = world_height * 0.75;
126       }
127     }
128
129     do {
130       universe->clear();
131
132       const int nb_attempts_max = 100;
133       int nb_attempts = 0;
134
135       for(int u = 0; u < nb_shapes; u++) {
136         Polygon *pol = 0;
137
138         nb_attempts = 0;
139
140         scalar_t shape_size;
141
142         if(random_shape_size) {
143           shape_size = 40 + 80 * drand48();
144         } else {
145           shape_size = 80;
146         }
147
148         scalar_t red, green, blue;
149
150         if(random_colors) {
151           do {
152             red = drand48();
153             green = drand48();
154             blue = drand48();
155           } while(red < 0.9 and green < 0.9 and blue < 0.9 and
156                   red > 0.1 and green > 0.1 and blue > 0.1);
157         } else {
158           red = 1.0;
159           green = 1.0;
160           blue = 1.0;
161         }
162
163         do {
164           scalar_t x[] = { - shape_size * 0.4, + shape_size * 0.4,
165                            + shape_size * 0.4, - shape_size * 0.4 };
166
167           scalar_t y[] = { - shape_size * 0.6, - shape_size * 0.6,
168                            + shape_size * 0.6, + shape_size * 0.6 };
169
170           scalar_t object_center_x = world_width * drand48();
171           scalar_t object_center_y = world_height * drand48();
172
173           delete pol;
174           pol = new Polygon(0.5, red, green, blue, x, y, sizeof(x) / sizeof(scalar_t));
175           pol->set_position(object_center_x, object_center_y, M_PI * 2 * drand48());
176           pol->set_speed(0, 0, 0);
177
178           universe->initialize_polygon(pol);
179
180           nb_attempts++;
181         } while(nb_attempts < nb_attempts_max &&
182                 (universe->collide(pol) || universe->collide_with_borders(pol, 2.0 / scaling)));
183
184         if(nb_attempts == nb_attempts_max) {
185           delete pol;
186           u = -1;
187           universe->clear();
188           nb_attempts = 0;
189         } else {
190           universe->add_polygon(pol);
191         }
192       }
193
194       if(pulling) {
195         grabbed_polygon = universe->pick_polygon(grab_start_x, grab_start_y);
196       }
197     } while(pulling and !grabbed_polygon);
198
199     failed = 0;
200
201     scalar_t grab_relative_x, grab_relative_y;
202
203     if(pulling) {
204       grab_relative_x = grabbed_polygon->relative_x(grab_start_x, grab_start_y);
205       grab_relative_y = grabbed_polygon->relative_y(grab_start_x, grab_start_y);
206     }
207
208     for(int s = 0; !failed && s < nb_simulated_frames; s++) {
209       if(s % every_nth == 0) {
210         int t = s / every_nth;
211
212         canvases[t]->clear();
213         draw_universe_on_canvas(canvases[t], scaling, universe);
214       }
215
216       if(s < nb_simulated_frames - 1) {
217         // Run the simulation
218         for(int i = 0; i < nb_iterations_per_steps; i++) {
219           if(pulling) {
220             scalar_t xf = grabbed_polygon->absolute_x(grab_relative_x, grab_relative_y);
221             scalar_t yf = grabbed_polygon->absolute_y(grab_relative_x, grab_relative_y);
222             if (xf < 0 || xf >= world_width || yf < 0 || yf >= world_height) {
223               failed = 1;
224             }
225             grabbed_polygon->apply_force(dt, xf, yf, 0.0, -1.0);
226           } else {
227             // Gravity
228           }
229           universe->update(dt, 1.0 / scaling);
230         }
231       }
232     }
233
234     total_nb_attempts++;
235
236     if(total_nb_attempts >= max_total_nb_attempts) {
237       cerr << "There was " << max_total_nb_attempts << " attempts at generating the sequences." << endl;
238       abort();
239     }
240
241   } while(failed);
242
243   for(int t = 0; t < nb_images; t++) {
244     unsigned char *src = canvases[t]->_data;
245     unsigned char *dst = output + t * width * height * 3;
246     for(int d = 0; d < 3; d++) {
247       for(int y = 0; y < height; y++) {
248         for(int x = 0; x < width; x++) {
249           dst[x + width * (y + height * d)] = src[d + 4 * (x + width * y)];
250         }
251       }
252     }
253   }
254
255   for(int t = 0; t < 2 * nb_saved_frames; t++) {
256     delete canvases[t];
257   }
258
259   delete universe;
260 }