Cosmetics.
[dyncnn.git] / flatland.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
35 using namespace std;
36
37 #include "misc.h"
38 #include "universe.h"
39 #include "canvas_cairo.h"
40
41 FILE *safe_fopen(const char *name, const char *mode) {
42   FILE *file = fopen(name, mode);
43   if(!file) {
44     cerr << "Cannot open " << name << endl;
45     exit(1);
46   }
47   return file;
48 }
49
50 void print_help(const char *command) {
51   cerr << command << " <nb sequences to generate> [--dir <dir>] [--seed <seed>]]" << endl;
52   exit(1);
53 }
54
55 //////////////////////////////////////////////////////////////////////
56
57 void draw_universe_on_canvas(CanvasCairo *canvas, scalar_t scaling,
58                              Universe *universe) {
59   canvas->set_line_width(1.0 / scaling);
60   universe->draw(canvas);
61 }
62
63 void draw_grabbing_point_on_canvas(CanvasCairo *canvas, scalar_t scaling,
64                                    scalar_t xg, scalar_t yg,
65                                    scalar_t r, scalar_t g, scalar_t b) {
66   scalar_t radius = 1/scaling;
67   int n = 36;
68   scalar_t xp[n], yp[n];
69   for(int k = 0; k < n; k++) {
70     scalar_t alpha = 2 * M_PI * scalar_t(k) / scalar_t(n);
71     xp[k] = xg + radius * cos(alpha);
72     yp[k] = yg + radius * sin(alpha);
73   }
74   canvas->set_drawing_color(r, g, b);
75   canvas->set_line_width(2.0);
76   canvas->draw_polygon(1, n, xp, yp);
77 }
78
79 //////////////////////////////////////////////////////////////////////
80
81 int main(int argc, char **argv) {
82   const scalar_t world_width = 400;
83   const scalar_t world_height = 400;
84   const scalar_t scaling = 0.16; // So that 400 * 0.16 = 64
85
86   const scalar_t dt = 0.1;
87   const int nb_iterations_per_steps = 5;
88
89   //////////////////////////////////////////////////////////////////////
90
91   // We will generate images { 0, every_nth, 2 * every_nth, ..., k * every_nth < nb_frames }
92
93   // The framerate every_nth may be set to smaller value to generate
94   // nice materials for presentations or papers.
95
96   int every_nth = 4;
97   int nb_frames = 5;
98   int random_grasp = 0;
99   int random_shape_size = 0;
100   int nb_shapes = 1;
101   char data_dir[1024] = "/tmp/";
102   int multi_images = 0;
103   int show_grabbing_point = 0;
104   int skip = -1;
105
106   //////////////////////////////////////////////////////////////////////
107
108   if(argc < 2) {
109     print_help(argv[0]);
110   }
111
112   int nb_sequences = atoi(argv[1]);
113
114   int i = 2;
115   while(i < argc) {
116     if(strcmp(argv[i], "--dir") == 0) {
117       i++;
118       if(i == argc) { print_help(argv[0]);}
119       strncpy(data_dir, argv[i], sizeof(data_dir) / sizeof(char) - 1);
120       i++;
121     }
122
123     else if(strcmp(argv[i], "--seed") == 0) {
124       i++;
125       if(i == argc) { print_help(argv[0]); }
126       srand48(atoi(argv[i]));
127       i++;
128     }
129
130     else if(strcmp(argv[i], "--nb_shapes") == 0) {
131       i++;
132       if(i == argc) { print_help(argv[0]);}
133       nb_shapes = atoi(argv[i]);
134       i++;
135     }
136
137     else if(strcmp(argv[i], "--random_grasp") == 0) {
138       random_grasp = 1;
139       i++;
140     }
141
142     else if(strcmp(argv[i], "--random_shape_size") == 0) {
143       random_shape_size = 1;
144       i++;
145     }
146
147     else if(strcmp(argv[i], "--every_nth") == 0) {
148       i++;
149       if(i == argc) { print_help(argv[0]);}
150       every_nth = atoi(argv[i]);
151       i++;
152     }
153
154     else if(strcmp(argv[i], "--nb_frames") == 0) {
155       i++;
156       if(i == argc) { print_help(argv[0]);}
157       nb_frames = atoi(argv[i]);
158       i++;
159     }
160
161     else if(strcmp(argv[i], "--multi_images") == 0) {
162       multi_images = 1;
163       i++;
164     }
165
166     else if(strcmp(argv[i], "--show_grabbing_point") == 0) {
167       show_grabbing_point = 1;
168       i++;
169     }
170
171     else if(strcmp(argv[i], "--skip") == 0) {
172       i++;
173       if(i == argc) { print_help(argv[0]);}
174       skip = atoi(argv[i]);
175       i++;
176     }
177
178     else {
179       cerr << "Unknown option " << argv[i] << "." << endl;
180       abort();
181     }
182   }
183
184   if(nb_shapes < 1 || nb_shapes > 10) {
185     cerr << "nb_shapes has to be in {1, ..., 10}" << endl;
186     abort();
187   }
188
189   if(nb_frames < 0 || nb_frames > 50) {
190     cerr << "nb_frames has to be in {0, ..., 50}" << endl;
191     abort();
192   }
193
194   for(int n = 0; n < nb_sequences; n++) {
195
196     Universe *universe;
197     Polygon *grabbed_polygon;
198
199     universe = new Universe(nb_shapes, world_width, world_height);
200
201     const int nb_saved_frames = (nb_frames + every_nth - 1) / every_nth;
202
203     CanvasCairo *canvases[nb_saved_frames * 2];
204
205     for(int s = 0; s < 2 * nb_saved_frames; s++) {
206       canvases[s] = new CanvasCairo(scaling, universe->width(), universe->height());
207     }
208
209     scalar_t grab_start_x, grab_start_y;
210
211     if(random_grasp) {
212       grab_start_x = world_width * (0.1 + 0.8 * drand48());
213       grab_start_y = world_height * (0.1 + 0.8 * drand48());
214     } else {
215       grab_start_x = world_width * 0.5;
216       grab_start_y = world_height * 0.75;
217     }
218
219     if((n+1)%100 == 0) {
220       cout << "Created "
221            << n+1 << "/" << nb_sequences << " sequences in "
222            << data_dir
223            << "." << endl;
224     }
225
226     do {
227       universe->clear();
228
229       const int nb_attempts_max = 100;
230       int nb_attempts = 0;
231
232       for(int u = 0; u < nb_shapes; u++) {
233         Polygon *pol = 0;
234
235         nb_attempts = 0;
236
237         scalar_t shape_size;
238
239         if(random_shape_size) {
240           shape_size = 40 + 80 * drand48();
241         } else {
242           shape_size = 80;
243         }
244
245         do {
246           scalar_t x[] = { - shape_size * 0.4, + shape_size * 0.4,
247                            + shape_size * 0.4, - shape_size * 0.4 };
248
249           scalar_t y[] = { - shape_size * 0.6, - shape_size * 0.6,
250                            + shape_size * 0.6, + shape_size * 0.6 };
251
252           scalar_t delta = shape_size / sqrt(2.0);
253
254           scalar_t object_center_x = delta + (world_width - 2 * delta) * drand48();
255           scalar_t object_center_y = delta + (world_height - 2 * delta) * drand48();
256
257           delete pol;
258           pol = new Polygon(0.5, 1.0, 1.0, 1.0, x, y, sizeof(x)/sizeof(scalar_t));
259           pol->set_position(object_center_x, object_center_y, M_PI * 2 * drand48());
260           pol->set_speed(0, 0, 0);
261
262           universe->initialize_polygon(pol);
263
264           nb_attempts++;
265         } while(nb_attempts < nb_attempts_max && universe->collide(pol));
266
267         if(nb_attempts == nb_attempts_max) {
268           delete pol;
269           u = -1;
270           universe->clear();
271           nb_attempts = 0;
272         } else {
273           universe->add_polygon(pol);
274         }
275       }
276
277       grabbed_polygon = universe->pick_polygon(grab_start_x, grab_start_y);
278     } while(!grabbed_polygon);
279
280     if(n%1000 == 0) {
281       char buffer[1024];
282       sprintf(buffer, "%s/%03d/", data_dir, n / 1000);
283       mkdir(buffer, 0777);
284     }
285
286     if(skip < 0 || n >= skip) {
287
288       scalar_t grab_relative_x = grabbed_polygon->relative_x(grab_start_x, grab_start_y);
289       scalar_t grab_relative_y = grabbed_polygon->relative_y(grab_start_x, grab_start_y);
290
291       for(int s = 0; s < nb_frames; s++) {
292         if(s % every_nth == 0) {
293           int t = s / every_nth;
294           scalar_t xf = grabbed_polygon->absolute_x(grab_relative_x, grab_relative_y);
295           scalar_t yf = grabbed_polygon->absolute_y(grab_relative_x, grab_relative_y);
296
297           canvases[2 * t + 0]->clear();
298           draw_grabbing_point_on_canvas(canvases[2 * t + 0], scaling,
299                                         xf, yf, 0.0, 0.0, 0.0);
300           canvases[2 * t + 1]->clear();
301           draw_universe_on_canvas(canvases[2 * t + 1], scaling, universe);
302
303           if(show_grabbing_point) {
304             draw_grabbing_point_on_canvas(canvases[2 * t + 1], scaling,
305                                           xf, yf, 1.0, 0.0, 0.0);
306           }
307         }
308
309         if(s < nb_frames - 1) {
310           // Run the simulation
311           for(int i = 0; i < nb_iterations_per_steps; i++) {
312             scalar_t xf = grabbed_polygon->absolute_x(grab_relative_x, grab_relative_y);
313             scalar_t yf = grabbed_polygon->absolute_y(grab_relative_x, grab_relative_y);
314             grabbed_polygon->apply_force(dt, xf, yf, 0.0, -1.0);
315             universe->update(dt);
316           }
317         }
318       }
319
320       char buffer[1024];
321
322       if(multi_images) {
323         for(int j = 0; j < nb_saved_frames; j++) {
324           FILE *file;
325           sprintf(buffer, "%s/%03d/dyn_%06d_grab_%02d.png", data_dir, n / 1000, n, j);
326           file = safe_fopen(buffer, "w");
327           canvases[j * 2 + 0]->write_png(file);
328           fclose(file);
329           sprintf(buffer, "%s/%03d/dyn_%06d_state_%02d.png", data_dir, n / 1000, n, j);
330           file = safe_fopen(buffer, "w");
331           canvases[j * 2 + 1]->write_png(file);
332           fclose(file);
333         }
334       } else {
335         CanvasCairo main_canvas(scaling, nb_saved_frames, 2, canvases);
336         sprintf(buffer, "%s/%03d/dyn_%06d.png", data_dir, n / 1000, n);
337         FILE *file = safe_fopen(buffer, "w");
338         main_canvas.write_png(file);
339         fclose(file);
340       }
341     }
342
343     for(int t = 0; t < 2 * nb_saved_frames; t++) {
344       delete canvases[t];
345     }
346
347     delete universe;
348   }
349 }