git.sesse.net Git - nageru/blob - flow.cpp

   1 #define NO_SDL_GLEXT 1
   2
   3 #include <epoxy/gl.h>
   4
   5 #include <SDL2/SDL.h>
   6 #include <SDL2/SDL_error.h>
   7 #include <SDL2/SDL_events.h>
   8 #include <SDL2/SDL_image.h>
   9 #include <SDL2/SDL_keyboard.h>
  10 #include <SDL2/SDL_mouse.h>
  11 #include <SDL2/SDL_video.h>
  12
  13 #include <assert.h>
  14 #include <getopt.h>
  15 #include <stdio.h>
  16 #include <unistd.h>
  17
  18 #include "util.h"
  19
  20 #include <algorithm>
  21 #include <memory>
  22 #include <map>
  23 #include <vector>
  24
  25 #define BUFFER_OFFSET(i) ((char *)nullptr + (i))
  26
  27 using namespace std;
  28
  29 // Operating point 3 (10 Hz on CPU, excluding preprocessing).
  30 constexpr float patch_overlap_ratio = 0.75f;
  31 constexpr unsigned coarsest_level = 5;
  32 constexpr unsigned finest_level = 1;
  33 constexpr unsigned patch_size_pixels = 12;
  34
  35 // Weighting constants for the different parts of the variational refinement.
  36 // These don't correspond 1:1 to the values given in the DIS paper,
  37 // since we have different normalizations and ranges in some cases.
  38 float vr_gamma = 10.0f, vr_delta = 5.0f, vr_alpha = 10.0f;
  39
  40 // Some global OpenGL objects.
  41 GLuint nearest_sampler, linear_sampler, smoothness_sampler;
  42 GLuint vertex_vbo;
  43
  44 string read_file(const string &filename)
  45 {
  46         FILE *fp = fopen(filename.c_str(), "r");
  47         if (fp == nullptr) {
  48                 perror(filename.c_str());
  49                 exit(1);
  50         }
  51
  52         int ret = fseek(fp, 0, SEEK_END);
  53         if (ret == -1) {
  54                 perror("fseek(SEEK_END)");
  55                 exit(1);
  56         }
  57
  58         int size = ftell(fp);
  59
  60         ret = fseek(fp, 0, SEEK_SET);
  61         if (ret == -1) {
  62                 perror("fseek(SEEK_SET)");
  63                 exit(1);
  64         }
  65
  66         string str;
  67         str.resize(size);
  68         ret = fread(&str[0], size, 1, fp);
  69         if (ret == -1) {
  70                 perror("fread");
  71                 exit(1);
  72         }
  73         if (ret == 0) {
  74                 fprintf(stderr, "Short read when trying to read %d bytes from %s\n",
  75                                 size, filename.c_str());
  76                 exit(1);
  77         }
  78         fclose(fp);
  79
  80         return str;
  81 }
  82
  83
  84 GLuint compile_shader(const string &shader_src, GLenum type)
  85 {
  86         GLuint obj = glCreateShader(type);
  87         const GLchar* source[] = { shader_src.data() };
  88         const GLint length[] = { (GLint)shader_src.size() };
  89         glShaderSource(obj, 1, source, length);
  90         glCompileShader(obj);
  91
  92         GLchar info_log[4096];
  93         GLsizei log_length = sizeof(info_log) - 1;
  94         glGetShaderInfoLog(obj, log_length, &log_length, info_log);
  95         info_log[log_length] = 0;
  96         if (strlen(info_log) > 0) {
  97                 fprintf(stderr, "Shader compile log: %s\n", info_log);
  98         }
  99
 100         GLint status;
 101         glGetShaderiv(obj, GL_COMPILE_STATUS, &status);
 102         if (status == GL_FALSE) {
 103                 // Add some line numbers to easier identify compile errors.
 104                 string src_with_lines = "/*   1 */ ";
 105                 size_t lineno = 1;
 106                 for (char ch : shader_src) {
 107                         src_with_lines.push_back(ch);
 108                         if (ch == '\n') {
 109                                 char buf[32];
 110                                 snprintf(buf, sizeof(buf), "/* %3zu */ ", ++lineno);
 111                                 src_with_lines += buf;
 112                         }
 113                 }
 114
 115                 fprintf(stderr, "Failed to compile shader:\n%s\n", src_with_lines.c_str());
 116                 exit(1);
 117         }
 118
 119         return obj;
 120 }
 121
 122 GLuint load_texture(const char *filename, unsigned *width_ret, unsigned *height_ret)
 123 {
 124         SDL_Surface *surf = IMG_Load(filename);
 125         if (surf == nullptr) {
 126                 fprintf(stderr, "IMG_Load(%s): %s\n", filename, IMG_GetError());
 127                 exit(1);
 128         }
 129
 130         // For whatever reason, SDL doesn't support converting to YUV surfaces
 131         // nor grayscale, so we'll do it (slowly) ourselves.
 132         SDL_Surface *rgb_surf = SDL_ConvertSurfaceFormat(surf, SDL_PIXELFORMAT_RGBA8888, /*flags=*/0);
 133         if (rgb_surf == nullptr) {
 134                 fprintf(stderr, "SDL_ConvertSurfaceFormat(%s): %s\n", filename, SDL_GetError());
 135                 exit(1);
 136         }
 137
 138         SDL_FreeSurface(surf);
 139
 140         unsigned width = rgb_surf->w, height = rgb_surf->h;
 141         const uint8_t *sptr = (uint8_t *)rgb_surf->pixels;
 142         unique_ptr<uint8_t[]> pix(new uint8_t[width * height]);
 143
 144         // Extract the Y component, and convert to bottom-left origin.
 145         for (unsigned y = 0; y < height; ++y) {
 146                 unsigned y2 = height - 1 - y;
 147                 for (unsigned x = 0; x < width; ++x) {
 148                         uint8_t r = sptr[(y2 * width + x) * 4 + 3];
 149                         uint8_t g = sptr[(y2 * width + x) * 4 + 2];
 150                         uint8_t b = sptr[(y2 * width + x) * 4 + 1];
 151
 152                         // Rec. 709.
 153                         pix[y * width + x] = lrintf(r * 0.2126f + g * 0.7152f + b * 0.0722f);
 154                 }
 155         }
 156         SDL_FreeSurface(rgb_surf);
 157
 158         int levels = 1;
 159         for (int w = width, h = height; w > 1 || h > 1; ) {
 160                 w >>= 1;
 161                 h >>= 1;
 162                 ++levels;
 163         }
 164
 165         GLuint tex;
 166         glCreateTextures(GL_TEXTURE_2D, 1, &tex);
 167         glTextureStorage2D(tex, levels, GL_R8, width, height);
 168         glTextureSubImage2D(tex, 0, 0, 0, width, height, GL_RED, GL_UNSIGNED_BYTE, pix.get());
 169         glGenerateTextureMipmap(tex);
 170
 171         *width_ret = width;
 172         *height_ret = height;
 173
 174         return tex;
 175 }
 176
 177 GLuint link_program(GLuint vs_obj, GLuint fs_obj)
 178 {
 179         GLuint program = glCreateProgram();
 180         glAttachShader(program, vs_obj);
 181         glAttachShader(program, fs_obj);
 182         glLinkProgram(program);
 183         GLint success;
 184         glGetProgramiv(program, GL_LINK_STATUS, &success);
 185         if (success == GL_FALSE) {
 186                 GLchar error_log[1024] = {0};
 187                 glGetProgramInfoLog(program, 1024, nullptr, error_log);
 188                 fprintf(stderr, "Error linking program: %s\n", error_log);
 189                 exit(1);
 190         }
 191         return program;
 192 }
 193
 194 GLuint generate_vbo(GLint size, GLsizeiptr data_size, const GLvoid *data)
 195 {
 196         GLuint vbo;
 197         glCreateBuffers(1, &vbo);
 198         glBufferData(GL_ARRAY_BUFFER, data_size, data, GL_STATIC_DRAW);
 199         glNamedBufferData(vbo, data_size, data, GL_STATIC_DRAW);
 200         return vbo;
 201 }
 202
 203 GLuint fill_vertex_attribute(GLuint vao, GLuint glsl_program_num, const string &attribute_name, GLint size, GLenum type, GLsizeiptr data_size, const GLvoid *data)
 204 {
 205         int attrib = glGetAttribLocation(glsl_program_num, attribute_name.c_str());
 206         if (attrib == -1) {
 207                 return -1;
 208         }
 209
 210         GLuint vbo = generate_vbo(size, data_size, data);
 211
 212         glBindBuffer(GL_ARRAY_BUFFER, vbo);
 213         glEnableVertexArrayAttrib(vao, attrib);
 214         glVertexAttribPointer(attrib, size, type, GL_FALSE, 0, BUFFER_OFFSET(0));
 215         glBindBuffer(GL_ARRAY_BUFFER, 0);
 216
 217         return vbo;
 218 }
 219
 220 void bind_sampler(GLuint program, GLint location, GLuint texture_unit, GLuint tex, GLuint sampler)
 221 {
 222         if (location == -1) {
 223                 return;
 224         }
 225
 226         glBindTextureUnit(texture_unit, tex);
 227         glBindSampler(texture_unit, sampler);
 228         glProgramUniform1i(program, location, texture_unit);
 229 }
 230
 231 // A class that caches FBOs that render to a given set of textures.
 232 // It never frees anything, so it is only suitable for rendering to
 233 // the same (small) set of textures over and over again.
 234 template<size_t num_elements>
 235 class PersistentFBOSet {
 236 public:
 237         void render_to(const array<GLuint, num_elements> &textures);
 238
 239         // Convenience wrappers.
 240         void render_to(GLuint texture0, enable_if<num_elements == 1> * = nullptr) {
 241                 render_to({{texture0}});
 242         }
 243
 244         void render_to(GLuint texture0, GLuint texture1, enable_if<num_elements == 2> * = nullptr) {
 245                 render_to({{texture0, texture1}});
 246         }
 247
 248         void render_to(GLuint texture0, GLuint texture1, GLuint texture2, enable_if<num_elements == 3> * = nullptr) {
 249                 render_to({{texture0, texture1, texture2}});
 250         }
 251
 252         void render_to(GLuint texture0, GLuint texture1, GLuint texture2, GLuint texture3, enable_if<num_elements == 4> * = nullptr) {
 253                 render_to({{texture0, texture1, texture2, texture3}});
 254         }
 255
 256 private:
 257         // TODO: Delete these on destruction.
 258         map<array<GLuint, num_elements>, GLuint> fbos;
 259 };
 260
 261 template<size_t num_elements>
 262 void PersistentFBOSet<num_elements>::render_to(const array<GLuint, num_elements> &textures)
 263 {
 264         auto it = fbos.find(textures);
 265         if (it != fbos.end()) {
 266                 glBindFramebuffer(GL_FRAMEBUFFER, it->second);
 267                 return;
 268         }
 269
 270         GLuint fbo;
 271         glCreateFramebuffers(1, &fbo);
 272         GLenum bufs[num_elements];
 273         for (size_t i = 0; i < num_elements; ++i) {
 274                 glNamedFramebufferTexture(fbo, GL_COLOR_ATTACHMENT0 + i, textures[i], 0);
 275                 bufs[i] = GL_COLOR_ATTACHMENT0 + i;
 276         }
 277         glNamedFramebufferDrawBuffers(fbo, num_elements, bufs);
 278
 279         fbos[textures] = fbo;
 280         glBindFramebuffer(GL_FRAMEBUFFER, fbo);
 281 }
 282
 283 // Compute gradients in every point, used for the motion search.
 284 // The DIS paper doesn't actually mention how these are computed,
 285 // but seemingly, a 3x3 Sobel operator is used here (at least in
 286 // later versions of the code), while a [1 -8 0 8 -1] kernel is
 287 // used for all the derivatives in the variational refinement part
 288 // (which borrows code from DeepFlow). This is inconsistent,
 289 // but I guess we're better off with staying with the original
 290 // decisions until we actually know having different ones would be better.
 291 class Sobel {
 292 public:
 293         Sobel();
 294         void exec(GLint tex0_view, GLint grad0_tex, int level_width, int level_height);
 295
 296 private:
 297         PersistentFBOSet<1> fbos;
 298         GLuint sobel_vs_obj;
 299         GLuint sobel_fs_obj;
 300         GLuint sobel_program;
 301         GLuint sobel_vao;
 302
 303         GLuint uniform_tex, uniform_image_size;
 304 };
 305
 306 Sobel::Sobel()
 307 {
 308         sobel_vs_obj = compile_shader(read_file("vs.vert"), GL_VERTEX_SHADER);
 309         sobel_fs_obj = compile_shader(read_file("sobel.frag"), GL_FRAGMENT_SHADER);
 310         sobel_program = link_program(sobel_vs_obj, sobel_fs_obj);
 311
 312         // Set up the VAO containing all the required position/texcoord data.
 313         glCreateVertexArrays(1, &sobel_vao);
 314         glBindVertexArray(sobel_vao);
 315
 316         GLint position_attrib = glGetAttribLocation(sobel_program, "position");
 317         glEnableVertexArrayAttrib(sobel_vao, position_attrib);
 318         glVertexAttribPointer(position_attrib, 2, GL_FLOAT, GL_FALSE, 0, BUFFER_OFFSET(0));
 319
 320         uniform_tex = glGetUniformLocation(sobel_program, "tex");
 321 }
 322
 323 void Sobel::exec(GLint tex0_view, GLint grad0_tex, int level_width, int level_height)
 324 {
 325         glUseProgram(sobel_program);
 326         glBindTextureUnit(0, tex0_view);
 327         glBindSampler(0, nearest_sampler);
 328         glProgramUniform1i(sobel_program, uniform_tex, 0);
 329
 330         glViewport(0, 0, level_width, level_height);
 331         fbos.render_to(grad0_tex);
 332         glBindVertexArray(sobel_vao);
 333         glUseProgram(sobel_program);
 334         glDisable(GL_BLEND);
 335         glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
 336 }
 337
 338 // Motion search to find the initial flow. See motion_search.frag for documentation.
 339 class MotionSearch {
 340 public:
 341         MotionSearch();
 342         void exec(GLuint tex0_view, GLuint tex1_view, GLuint grad0_tex, GLuint flow_tex, GLuint flow_out_tex, int level_width, int level_height, int prev_level_width, int prev_level_height, int width_patches, int height_patches);
 343
 344 private:
 345         PersistentFBOSet<1> fbos;
 346
 347         GLuint motion_vs_obj;
 348         GLuint motion_fs_obj;
 349         GLuint motion_search_program;
 350         GLuint motion_search_vao;
 351
 352         GLuint uniform_image_size, uniform_inv_image_size, uniform_inv_prev_level_size;
 353         GLuint uniform_image0_tex, uniform_image1_tex, uniform_grad0_tex, uniform_flow_tex;
 354 };
 355
 356 MotionSearch::MotionSearch()
 357 {
 358         motion_vs_obj = compile_shader(read_file("motion_search.vert"), GL_VERTEX_SHADER);
 359         motion_fs_obj = compile_shader(read_file("motion_search.frag"), GL_FRAGMENT_SHADER);
 360         motion_search_program = link_program(motion_vs_obj, motion_fs_obj);
 361
 362         // Set up the VAO containing all the required position/texcoord data.
 363         glCreateVertexArrays(1, &motion_search_vao);
 364         glBindVertexArray(motion_search_vao);
 365         glBindBuffer(GL_ARRAY_BUFFER, vertex_vbo);
 366
 367         GLint position_attrib = glGetAttribLocation(motion_search_program, "position");
 368         glEnableVertexArrayAttrib(motion_search_vao, position_attrib);
 369         glVertexAttribPointer(position_attrib, 2, GL_FLOAT, GL_FALSE, 0, BUFFER_OFFSET(0));
 370
 371         uniform_image_size = glGetUniformLocation(motion_search_program, "image_size");
 372         uniform_inv_image_size = glGetUniformLocation(motion_search_program, "inv_image_size");
 373         uniform_inv_prev_level_size = glGetUniformLocation(motion_search_program, "inv_prev_level_size");
 374         uniform_image0_tex = glGetUniformLocation(motion_search_program, "image0_tex");
 375         uniform_image1_tex = glGetUniformLocation(motion_search_program, "image1_tex");
 376         uniform_grad0_tex = glGetUniformLocation(motion_search_program, "grad0_tex");
 377         uniform_flow_tex = glGetUniformLocation(motion_search_program, "flow_tex");
 378 }
 379
 380 void MotionSearch::exec(GLuint tex0_view, GLuint tex1_view, GLuint grad0_tex, GLuint flow_tex, GLuint flow_out_tex, int level_width, int level_height, int prev_level_width, int prev_level_height, int width_patches, int height_patches)
 381 {
 382         glUseProgram(motion_search_program);
 383
 384         bind_sampler(motion_search_program, uniform_image0_tex, 0, tex0_view, nearest_sampler);
 385         bind_sampler(motion_search_program, uniform_image1_tex, 1, tex1_view, linear_sampler);
 386         bind_sampler(motion_search_program, uniform_grad0_tex, 2, grad0_tex, nearest_sampler);
 387         bind_sampler(motion_search_program, uniform_flow_tex, 3, flow_tex, linear_sampler);
 388
 389         glProgramUniform2f(motion_search_program, uniform_image_size, level_width, level_height);
 390         glProgramUniform2f(motion_search_program, uniform_inv_image_size, 1.0f / level_width, 1.0f / level_height);
 391         glProgramUniform2f(motion_search_program, uniform_inv_prev_level_size, 1.0f / prev_level_width, 1.0f / prev_level_height);
 392
 393         glViewport(0, 0, width_patches, height_patches);
 394         fbos.render_to(flow_out_tex);
 395         glBindVertexArray(motion_search_vao);
 396         glUseProgram(motion_search_program);
 397         glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
 398 }
 399
 400 // Do “densification”, ie., upsampling of the flow patches to the flow field
 401 // (the same size as the image at this level). We draw one quad per patch
 402 // over its entire covered area (using instancing in the vertex shader),
 403 // and then weight the contributions in the pixel shader by post-warp difference.
 404 // This is equation (3) in the paper.
 405 //
 406 // We accumulate the flow vectors in the R/G channels (for u/v) and the total
 407 // weight in the B channel. Dividing R and G by B gives the normalized values.
 408 class Densify {
 409 public:
 410         Densify();
 411         void exec(GLuint tex0_view, GLuint tex1_view, GLuint flow_tex, GLuint dense_flow_tex, int level_width, int level_height, int width_patches, int height_patches);
 412
 413 private:
 414         PersistentFBOSet<1> fbos;
 415
 416         GLuint densify_vs_obj;
 417         GLuint densify_fs_obj;
 418         GLuint densify_program;
 419         GLuint densify_vao;
 420
 421         GLuint uniform_width_patches, uniform_patch_size, uniform_patch_spacing;
 422         GLuint uniform_image0_tex, uniform_image1_tex, uniform_flow_tex;
 423 };
 424
 425 Densify::Densify()
 426 {
 427         densify_vs_obj = compile_shader(read_file("densify.vert"), GL_VERTEX_SHADER);
 428         densify_fs_obj = compile_shader(read_file("densify.frag"), GL_FRAGMENT_SHADER);
 429         densify_program = link_program(densify_vs_obj, densify_fs_obj);
 430
 431         // Set up the VAO containing all the required position/texcoord data.
 432         glCreateVertexArrays(1, &densify_vao);
 433         glBindVertexArray(densify_vao);
 434         glBindBuffer(GL_ARRAY_BUFFER, vertex_vbo);
 435
 436         GLint position_attrib = glGetAttribLocation(densify_program, "position");
 437         glEnableVertexArrayAttrib(densify_vao, position_attrib);
 438         glVertexAttribPointer(position_attrib, 2, GL_FLOAT, GL_FALSE, 0, BUFFER_OFFSET(0));
 439
 440         uniform_width_patches = glGetUniformLocation(densify_program, "width_patches");
 441         uniform_patch_size = glGetUniformLocation(densify_program, "patch_size");
 442         uniform_patch_spacing = glGetUniformLocation(densify_program, "patch_spacing");
 443         uniform_image0_tex = glGetUniformLocation(densify_program, "image0_tex");
 444         uniform_image1_tex = glGetUniformLocation(densify_program, "image1_tex");
 445         uniform_flow_tex = glGetUniformLocation(densify_program, "flow_tex");
 446 }
 447
 448 void Densify::exec(GLuint tex0_view, GLuint tex1_view, GLuint flow_tex, GLuint dense_flow_tex, int level_width, int level_height, int width_patches, int height_patches)
 449 {
 450         glUseProgram(densify_program);
 451
 452         bind_sampler(densify_program, uniform_image0_tex, 0, tex0_view, nearest_sampler);
 453         bind_sampler(densify_program, uniform_image1_tex, 1, tex1_view, linear_sampler);
 454         bind_sampler(densify_program, uniform_flow_tex, 2, flow_tex, nearest_sampler);
 455
 456         glProgramUniform1i(densify_program, uniform_width_patches, width_patches);
 457         glProgramUniform2f(densify_program, uniform_patch_size,
 458                 float(patch_size_pixels) / level_width,
 459                 float(patch_size_pixels) / level_height);
 460
 461         float patch_spacing_x = float(level_width - patch_size_pixels) / (width_patches - 1);
 462         float patch_spacing_y = float(level_height - patch_size_pixels) / (height_patches - 1);
 463         if (width_patches == 1) patch_spacing_x = 0.0f;  // Avoid infinities.
 464         if (height_patches == 1) patch_spacing_y = 0.0f;
 465         glProgramUniform2f(densify_program, uniform_patch_spacing,
 466                 patch_spacing_x / level_width,
 467                 patch_spacing_y / level_height);
 468
 469         glViewport(0, 0, level_width, level_height);
 470         glEnable(GL_BLEND);
 471         glBlendFunc(GL_ONE, GL_ONE);
 472         glBindVertexArray(densify_vao);
 473         fbos.render_to(dense_flow_tex);
 474         glDrawArraysInstanced(GL_TRIANGLE_STRIP, 0, 4, width_patches * height_patches);
 475 }
 476
 477 // Warp I_1 to I_w, and then compute the mean (I) and difference (I_t) of
 478 // I_0 and I_w. The prewarping is what enables us to solve the variational
 479 // flow for du,dv instead of u,v.
 480 //
 481 // Also calculates the normalized flow, ie. divides by z (this is needed because
 482 // Densify works by additive blending) and multiplies by the image size.
 483 //
 484 // See variational_refinement.txt for more information.
 485 class Prewarp {
 486 public:
 487         Prewarp();
 488         void exec(GLuint tex0_view, GLuint tex1_view, GLuint flow_tex, GLuint normalized_flow_tex, GLuint I_tex, GLuint I_t_tex, int level_width, int level_height);
 489
 490 private:
 491         PersistentFBOSet<3> fbos;
 492
 493         GLuint prewarp_vs_obj;
 494         GLuint prewarp_fs_obj;
 495         GLuint prewarp_program;
 496         GLuint prewarp_vao;
 497
 498         GLuint uniform_image0_tex, uniform_image1_tex, uniform_flow_tex;
 499         GLuint uniform_image_size;
 500 };
 501
 502 Prewarp::Prewarp()
 503 {
 504         prewarp_vs_obj = compile_shader(read_file("vs.vert"), GL_VERTEX_SHADER);
 505         prewarp_fs_obj = compile_shader(read_file("prewarp.frag"), GL_FRAGMENT_SHADER);
 506         prewarp_program = link_program(prewarp_vs_obj, prewarp_fs_obj);
 507
 508         // Set up the VAO containing all the required position/texcoord data.
 509         glCreateVertexArrays(1, &prewarp_vao);
 510         glBindVertexArray(prewarp_vao);
 511         glBindBuffer(GL_ARRAY_BUFFER, vertex_vbo);
 512
 513         GLint position_attrib = glGetAttribLocation(prewarp_program, "position");
 514         glEnableVertexArrayAttrib(prewarp_vao, position_attrib);
 515         glVertexAttribPointer(position_attrib, 2, GL_FLOAT, GL_FALSE, 0, BUFFER_OFFSET(0));
 516
 517         uniform_image0_tex = glGetUniformLocation(prewarp_program, "image0_tex");
 518         uniform_image1_tex = glGetUniformLocation(prewarp_program, "image1_tex");
 519         uniform_flow_tex = glGetUniformLocation(prewarp_program, "flow_tex");
 520
 521         uniform_image_size = glGetUniformLocation(prewarp_program, "image_size");
 522 }
 523
 524 void Prewarp::exec(GLuint tex0_view, GLuint tex1_view, GLuint flow_tex, GLuint I_tex, GLuint I_t_tex, GLuint normalized_flow_tex, int level_width, int level_height)
 525 {
 526         glUseProgram(prewarp_program);
 527
 528         bind_sampler(prewarp_program, uniform_image0_tex, 0, tex0_view, nearest_sampler);
 529         bind_sampler(prewarp_program, uniform_image1_tex, 1, tex1_view, linear_sampler);
 530         bind_sampler(prewarp_program, uniform_flow_tex, 2, flow_tex, nearest_sampler);
 531
 532         glProgramUniform2f(prewarp_program, uniform_image_size, level_width, level_height);
 533
 534         glViewport(0, 0, level_width, level_height);
 535         glDisable(GL_BLEND);
 536         glBindVertexArray(prewarp_vao);
 537         fbos.render_to(I_tex, I_t_tex, normalized_flow_tex);
 538         glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
 539 }
 540
 541 // From I, calculate the partial derivatives I_x and I_y. We use a four-tap
 542 // central difference filter, since apparently, that's tradition (I haven't
 543 // measured quality versus a more normal 0.5 (I[x+1] - I[x-1]).)
 544 // The coefficients come from
 545 //
 546 //   https://en.wikipedia.org/wiki/Finite_difference_coefficient
 547 //
 548 // Also computes β_0, since it depends only on I_x and I_y.
 549 class Derivatives {
 550 public:
 551         Derivatives();
 552         void exec(GLuint input_tex, GLuint I_x_y_tex, GLuint beta_0_tex, int level_width, int level_height);
 553
 554 private:
 555         PersistentFBOSet<2> fbos;
 556
 557         GLuint derivatives_vs_obj;
 558         GLuint derivatives_fs_obj;
 559         GLuint derivatives_program;
 560         GLuint derivatives_vao;
 561
 562         GLuint uniform_tex;
 563 };
 564
 565 Derivatives::Derivatives()
 566 {
 567         derivatives_vs_obj = compile_shader(read_file("vs.vert"), GL_VERTEX_SHADER);
 568         derivatives_fs_obj = compile_shader(read_file("derivatives.frag"), GL_FRAGMENT_SHADER);
 569         derivatives_program = link_program(derivatives_vs_obj, derivatives_fs_obj);
 570
 571         // Set up the VAO containing all the required position/texcoord data.
 572         glCreateVertexArrays(1, &derivatives_vao);
 573         glBindVertexArray(derivatives_vao);
 574         glBindBuffer(GL_ARRAY_BUFFER, vertex_vbo);
 575
 576         GLint position_attrib = glGetAttribLocation(derivatives_program, "position");
 577         glEnableVertexArrayAttrib(derivatives_vao, position_attrib);
 578         glVertexAttribPointer(position_attrib, 2, GL_FLOAT, GL_FALSE, 0, BUFFER_OFFSET(0));
 579
 580         uniform_tex = glGetUniformLocation(derivatives_program, "tex");
 581 }
 582
 583 void Derivatives::exec(GLuint input_tex, GLuint I_x_y_tex, GLuint beta_0_tex, int level_width, int level_height)
 584 {
 585         glUseProgram(derivatives_program);
 586
 587         bind_sampler(derivatives_program, uniform_tex, 0, input_tex, nearest_sampler);
 588
 589         glViewport(0, 0, level_width, level_height);
 590         glDisable(GL_BLEND);
 591         glBindVertexArray(derivatives_vao);
 592         fbos.render_to(I_x_y_tex, beta_0_tex);
 593         glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
 594 }
 595
 596 // Calculate the smoothness constraints between neighboring pixels;
 597 // s_x(x,y) stores smoothness between pixel (x,y) and (x+1,y),
 598 // and s_y(x,y) stores between (x,y) and (x,y+1). We'll sample with
 599 // border color (0,0) later, so that there's zero diffusion out of
 600 // the border.
 601 //
 602 // See variational_refinement.txt for more information.
 603 class ComputeSmoothness {
 604 public:
 605         ComputeSmoothness();
 606         void exec(GLuint flow_tex, GLuint diff_flow_tex, GLuint smoothness_x_tex, GLuint smoothness_y_tex, int level_width, int level_height);
 607
 608 private:
 609         PersistentFBOSet<2> fbos;
 610
 611         GLuint smoothness_vs_obj;
 612         GLuint smoothness_fs_obj;
 613         GLuint smoothness_program;
 614         GLuint smoothness_vao;
 615
 616         GLuint uniform_flow_tex, uniform_diff_flow_tex;
 617         GLuint uniform_alpha;
 618 };
 619
 620 ComputeSmoothness::ComputeSmoothness()
 621 {
 622         smoothness_vs_obj = compile_shader(read_file("vs.vert"), GL_VERTEX_SHADER);
 623         smoothness_fs_obj = compile_shader(read_file("smoothness.frag"), GL_FRAGMENT_SHADER);
 624         smoothness_program = link_program(smoothness_vs_obj, smoothness_fs_obj);
 625
 626         // Set up the VAO containing all the required position/texcoord data.
 627         glCreateVertexArrays(1, &smoothness_vao);
 628         glBindVertexArray(smoothness_vao);
 629         glBindBuffer(GL_ARRAY_BUFFER, vertex_vbo);
 630
 631         GLint position_attrib = glGetAttribLocation(smoothness_program, "position");
 632         glEnableVertexArrayAttrib(smoothness_vao, position_attrib);
 633         glVertexAttribPointer(position_attrib, 2, GL_FLOAT, GL_FALSE, 0, BUFFER_OFFSET(0));
 634
 635         uniform_flow_tex = glGetUniformLocation(smoothness_program, "flow_tex");
 636         uniform_diff_flow_tex = glGetUniformLocation(smoothness_program, "diff_flow_tex");
 637         uniform_alpha = glGetUniformLocation(smoothness_program, "alpha");
 638 }
 639
 640 void ComputeSmoothness::exec(GLuint flow_tex, GLuint diff_flow_tex, GLuint smoothness_x_tex, GLuint smoothness_y_tex, int level_width, int level_height)
 641 {
 642         glUseProgram(smoothness_program);
 643
 644         bind_sampler(smoothness_program, uniform_flow_tex, 0, flow_tex, nearest_sampler);
 645         bind_sampler(smoothness_program, uniform_diff_flow_tex, 1, diff_flow_tex, nearest_sampler);
 646         glProgramUniform1f(smoothness_program, uniform_alpha, vr_alpha);
 647
 648         glViewport(0, 0, level_width, level_height);
 649
 650         glDisable(GL_BLEND);
 651         glBindVertexArray(smoothness_vao);
 652         fbos.render_to(smoothness_x_tex, smoothness_y_tex);
 653         glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
 654
 655         // Make sure the smoothness on the right and upper borders is zero.
 656         // We could have done this by making (W-1)xH and Wx(H-1) textures instead
 657         // (we're sampling smoothness with all-zero border color), but we'd
 658         // have to adjust the sampling coordinates, which is annoying.
 659         glClearTexSubImage(smoothness_x_tex, 0,  level_width - 1, 0, 0,   1, level_height, 1,  GL_RED, GL_FLOAT, nullptr);
 660         glClearTexSubImage(smoothness_y_tex, 0,  0, level_height - 1, 0,  level_width, 1, 1,   GL_RED, GL_FLOAT, nullptr);
 661 }
 662
 663 // Set up the equations set (two equations in two unknowns, per pixel).
 664 // We store five floats; the three non-redundant elements of the 2x2 matrix (A)
 665 // as 32-bit floats, and the two elements on the right-hand side (b) as 16-bit
 666 // floats. (Actually, we store the inverse of the diagonal elements, because
 667 // we only ever need to divide by them.) This fits into four u32 values;
 668 // R, G, B for the matrix (the last element is symmetric) and A for the two b values.
 669 // All the values of the energy term (E_I, E_G, E_S), except the smoothness
 670 // terms that depend on other pixels, are calculated in one pass.
 671 //
 672 // See variational_refinement.txt for more information.
 673 class SetupEquations {
 674 public:
 675         SetupEquations();
 676         void exec(GLuint I_x_y_tex, GLuint I_t_tex, GLuint diff_flow_tex, GLuint flow_tex, GLuint beta_0_tex, GLuint smoothness_x_tex, GLuint smoothness_y_tex, GLuint equation_tex, int level_width, int level_height);
 677
 678 private:
 679         PersistentFBOSet<1> fbos;
 680
 681         GLuint equations_vs_obj;
 682         GLuint equations_fs_obj;
 683         GLuint equations_program;
 684         GLuint equations_vao;
 685
 686         GLuint uniform_I_x_y_tex, uniform_I_t_tex;
 687         GLuint uniform_diff_flow_tex, uniform_base_flow_tex;
 688         GLuint uniform_beta_0_tex;
 689         GLuint uniform_smoothness_x_tex, uniform_smoothness_y_tex;
 690         GLuint uniform_gamma, uniform_delta;
 691 };
 692
 693 SetupEquations::SetupEquations()
 694 {
 695         equations_vs_obj = compile_shader(read_file("vs.vert"), GL_VERTEX_SHADER);
 696         equations_fs_obj = compile_shader(read_file("equations.frag"), GL_FRAGMENT_SHADER);
 697         equations_program = link_program(equations_vs_obj, equations_fs_obj);
 698
 699         // Set up the VAO containing all the required position/texcoord data.
 700         glCreateVertexArrays(1, &equations_vao);
 701         glBindVertexArray(equations_vao);
 702         glBindBuffer(GL_ARRAY_BUFFER, vertex_vbo);
 703
 704         GLint position_attrib = glGetAttribLocation(equations_program, "position");
 705         glEnableVertexArrayAttrib(equations_vao, position_attrib);
 706         glVertexAttribPointer(position_attrib, 2, GL_FLOAT, GL_FALSE, 0, BUFFER_OFFSET(0));
 707
 708         uniform_I_x_y_tex = glGetUniformLocation(equations_program, "I_x_y_tex");
 709         uniform_I_t_tex = glGetUniformLocation(equations_program, "I_t_tex");
 710         uniform_diff_flow_tex = glGetUniformLocation(equations_program, "diff_flow_tex");
 711         uniform_base_flow_tex = glGetUniformLocation(equations_program, "base_flow_tex");
 712         uniform_beta_0_tex = glGetUniformLocation(equations_program, "beta_0_tex");
 713         uniform_smoothness_x_tex = glGetUniformLocation(equations_program, "smoothness_x_tex");
 714         uniform_smoothness_y_tex = glGetUniformLocation(equations_program, "smoothness_y_tex");
 715         uniform_gamma = glGetUniformLocation(equations_program, "gamma");
 716         uniform_delta = glGetUniformLocation(equations_program, "delta");
 717 }
 718
 719 void SetupEquations::exec(GLuint I_x_y_tex, GLuint I_t_tex, GLuint diff_flow_tex, GLuint base_flow_tex, GLuint beta_0_tex, GLuint smoothness_x_tex, GLuint smoothness_y_tex, GLuint equation_tex, int level_width, int level_height)
 720 {
 721         glUseProgram(equations_program);
 722
 723         bind_sampler(equations_program, uniform_I_x_y_tex, 0, I_x_y_tex, nearest_sampler);
 724         bind_sampler(equations_program, uniform_I_t_tex, 1, I_t_tex, nearest_sampler);
 725         bind_sampler(equations_program, uniform_diff_flow_tex, 2, diff_flow_tex, nearest_sampler);
 726         bind_sampler(equations_program, uniform_base_flow_tex, 3, base_flow_tex, nearest_sampler);
 727         bind_sampler(equations_program, uniform_beta_0_tex, 4, beta_0_tex, nearest_sampler);
 728         bind_sampler(equations_program, uniform_smoothness_x_tex, 5, smoothness_x_tex, smoothness_sampler);
 729         bind_sampler(equations_program, uniform_smoothness_y_tex, 6, smoothness_y_tex, smoothness_sampler);
 730         glProgramUniform1f(equations_program, uniform_delta, vr_delta);
 731         glProgramUniform1f(equations_program, uniform_gamma, vr_gamma);
 732
 733         glViewport(0, 0, level_width, level_height);
 734         glDisable(GL_BLEND);
 735         glBindVertexArray(equations_vao);
 736         fbos.render_to(equation_tex);
 737         glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
 738 }
 739
 740 // Actually solve the equation sets made by SetupEquations, by means of
 741 // successive over-relaxation (SOR).
 742 //
 743 // See variational_refinement.txt for more information.
 744 class SOR {
 745 public:
 746         SOR();
 747         void exec(GLuint diff_flow_tex, GLuint equation_tex, GLuint smoothness_x_tex, GLuint smoothness_y_tex, int level_width, int level_height, int num_iterations);
 748
 749 private:
 750         PersistentFBOSet<1> fbos;
 751
 752         GLuint sor_vs_obj;
 753         GLuint sor_fs_obj;
 754         GLuint sor_program;
 755         GLuint sor_vao;
 756
 757         GLuint uniform_diff_flow_tex;
 758         GLuint uniform_equation_tex;
 759         GLuint uniform_smoothness_x_tex, uniform_smoothness_y_tex;
 760 };
 761
 762 SOR::SOR()
 763 {
 764         sor_vs_obj = compile_shader(read_file("vs.vert"), GL_VERTEX_SHADER);
 765         sor_fs_obj = compile_shader(read_file("sor.frag"), GL_FRAGMENT_SHADER);
 766         sor_program = link_program(sor_vs_obj, sor_fs_obj);
 767
 768         // Set up the VAO containing all the required position/texcoord data.
 769         glCreateVertexArrays(1, &sor_vao);
 770         glBindVertexArray(sor_vao);
 771         glBindBuffer(GL_ARRAY_BUFFER, vertex_vbo);
 772
 773         GLint position_attrib = glGetAttribLocation(sor_program, "position");
 774         glEnableVertexArrayAttrib(sor_vao, position_attrib);
 775         glVertexAttribPointer(position_attrib, 2, GL_FLOAT, GL_FALSE, 0, BUFFER_OFFSET(0));
 776
 777         uniform_diff_flow_tex = glGetUniformLocation(sor_program, "diff_flow_tex");
 778         uniform_equation_tex = glGetUniformLocation(sor_program, "equation_tex");
 779         uniform_smoothness_x_tex = glGetUniformLocation(sor_program, "smoothness_x_tex");
 780         uniform_smoothness_y_tex = glGetUniformLocation(sor_program, "smoothness_y_tex");
 781 }
 782
 783 void SOR::exec(GLuint diff_flow_tex, GLuint equation_tex, GLuint smoothness_x_tex, GLuint smoothness_y_tex, int level_width, int level_height, int num_iterations)
 784 {
 785         glUseProgram(sor_program);
 786
 787         bind_sampler(sor_program, uniform_diff_flow_tex, 0, diff_flow_tex, nearest_sampler);
 788         bind_sampler(sor_program, uniform_smoothness_x_tex, 1, smoothness_x_tex, smoothness_sampler);
 789         bind_sampler(sor_program, uniform_smoothness_y_tex, 2, smoothness_y_tex, smoothness_sampler);
 790         bind_sampler(sor_program, uniform_equation_tex, 3, equation_tex, nearest_sampler);
 791
 792         glViewport(0, 0, level_width, level_height);
 793         glDisable(GL_BLEND);
 794         glBindVertexArray(sor_vao);
 795         fbos.render_to(diff_flow_tex);  // NOTE: Bind to same as we render from!
 796
 797         for (int i = 0; i < num_iterations; ++i) {
 798                 glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
 799                 if (i != num_iterations - 1) {
 800                         glTextureBarrier();
 801                 }
 802         }
 803 }
 804
 805 // Simply add the differential flow found by the variational refinement to the base flow.
 806 // The output is in base_flow_tex; we don't need to make a new texture.
 807 class AddBaseFlow {
 808 public:
 809         AddBaseFlow();
 810         void exec(GLuint base_flow_tex, GLuint diff_flow_tex, int level_width, int level_height);
 811
 812 private:
 813         PersistentFBOSet<1> fbos;
 814
 815         GLuint add_flow_vs_obj;
 816         GLuint add_flow_fs_obj;
 817         GLuint add_flow_program;
 818         GLuint add_flow_vao;
 819
 820         GLuint uniform_diff_flow_tex;
 821 };
 822
 823 AddBaseFlow::AddBaseFlow()
 824 {
 825         add_flow_vs_obj = compile_shader(read_file("vs.vert"), GL_VERTEX_SHADER);
 826         add_flow_fs_obj = compile_shader(read_file("add_base_flow.frag"), GL_FRAGMENT_SHADER);
 827         add_flow_program = link_program(add_flow_vs_obj, add_flow_fs_obj);
 828
 829         // Set up the VAO containing all the required position/texcoord data.
 830         glCreateVertexArrays(1, &add_flow_vao);
 831         glBindVertexArray(add_flow_vao);
 832         glBindBuffer(GL_ARRAY_BUFFER, vertex_vbo);
 833
 834         GLint position_attrib = glGetAttribLocation(add_flow_program, "position");
 835         glEnableVertexArrayAttrib(add_flow_vao, position_attrib);
 836         glVertexAttribPointer(position_attrib, 2, GL_FLOAT, GL_FALSE, 0, BUFFER_OFFSET(0));
 837
 838         uniform_diff_flow_tex = glGetUniformLocation(add_flow_program, "diff_flow_tex");
 839 }
 840
 841 void AddBaseFlow::exec(GLuint base_flow_tex, GLuint diff_flow_tex, int level_width, int level_height)
 842 {
 843         glUseProgram(add_flow_program);
 844
 845         bind_sampler(add_flow_program, uniform_diff_flow_tex, 0, diff_flow_tex, nearest_sampler);
 846
 847         glViewport(0, 0, level_width, level_height);
 848         glEnable(GL_BLEND);
 849         glBlendFunc(GL_ONE, GL_ONE);
 850         glBindVertexArray(add_flow_vao);
 851         fbos.render_to(base_flow_tex);
 852
 853         glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
 854 }
 855
 856 // Take a copy of the flow, bilinearly interpolated and scaled up.
 857 class ResizeFlow {
 858 public:
 859         ResizeFlow();
 860         void exec(GLuint in_tex, GLuint out_tex, int input_width, int input_height, int output_width, int output_height);
 861
 862 private:
 863         PersistentFBOSet<1> fbos;
 864
 865         GLuint resize_flow_vs_obj;
 866         GLuint resize_flow_fs_obj;
 867         GLuint resize_flow_program;
 868         GLuint resize_flow_vao;
 869
 870         GLuint uniform_flow_tex;
 871         GLuint uniform_scale_factor;
 872 };
 873
 874 ResizeFlow::ResizeFlow()
 875 {
 876         resize_flow_vs_obj = compile_shader(read_file("vs.vert"), GL_VERTEX_SHADER);
 877         resize_flow_fs_obj = compile_shader(read_file("resize_flow.frag"), GL_FRAGMENT_SHADER);
 878         resize_flow_program = link_program(resize_flow_vs_obj, resize_flow_fs_obj);
 879
 880         // Set up the VAO containing all the required position/texcoord data.
 881         glCreateVertexArrays(1, &resize_flow_vao);
 882         glBindVertexArray(resize_flow_vao);
 883         glBindBuffer(GL_ARRAY_BUFFER, vertex_vbo);
 884
 885         GLint position_attrib = glGetAttribLocation(resize_flow_program, "position");
 886         glEnableVertexArrayAttrib(resize_flow_vao, position_attrib);
 887         glVertexAttribPointer(position_attrib, 2, GL_FLOAT, GL_FALSE, 0, BUFFER_OFFSET(0));
 888
 889         uniform_flow_tex = glGetUniformLocation(resize_flow_program, "flow_tex");
 890         uniform_scale_factor = glGetUniformLocation(resize_flow_program, "scale_factor");
 891 }
 892
 893 void ResizeFlow::exec(GLuint flow_tex, GLuint out_tex, int input_width, int input_height, int output_width, int output_height)
 894 {
 895         glUseProgram(resize_flow_program);
 896
 897         bind_sampler(resize_flow_program, uniform_flow_tex, 0, flow_tex, nearest_sampler);
 898
 899         glProgramUniform2f(resize_flow_program, uniform_scale_factor, float(output_width) / input_width, float(output_height) / input_height);
 900
 901         glViewport(0, 0, output_width, output_height);
 902         glDisable(GL_BLEND);
 903         glBindVertexArray(resize_flow_vao);
 904         fbos.render_to(out_tex);
 905
 906         glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
 907 }
 908
 909 class GPUTimers {
 910 public:
 911         void print();
 912         pair<GLuint, GLuint> begin_timer(const string &name, int level);
 913
 914 private:
 915         struct Timer {
 916                 string name;
 917                 int level;
 918                 pair<GLuint, GLuint> query;
 919         };
 920         vector<Timer> timers;
 921 };
 922
 923 pair<GLuint, GLuint> GPUTimers::begin_timer(const string &name, int level)
 924 {
 925         GLuint queries[2];
 926         glGenQueries(2, queries);
 927         glQueryCounter(queries[0], GL_TIMESTAMP);
 928
 929         Timer timer;
 930         timer.name = name;
 931         timer.level = level;
 932         timer.query.first = queries[0];
 933         timer.query.second = queries[1];
 934         timers.push_back(timer);
 935         return timer.query;
 936 }
 937
 938 void GPUTimers::print()
 939 {
 940         for (const Timer &timer : timers) {
 941                 // NOTE: This makes the CPU wait for the GPU.
 942                 GLuint64 time_start, time_end;
 943                 glGetQueryObjectui64v(timer.query.first, GL_QUERY_RESULT, &time_start);
 944                 glGetQueryObjectui64v(timer.query.second, GL_QUERY_RESULT, &time_end);
 945                 //fprintf(stderr, "GPU time used = %.1f ms\n", time_elapsed / 1e6);
 946                 for (int i = 0; i < timer.level * 2; ++i) {
 947                         fprintf(stderr, " ");
 948                 }
 949                 fprintf(stderr, "%-30s %4.1f ms\n", timer.name.c_str(), GLint64(time_end - time_start) / 1e6);
 950         }
 951 }
 952
 953 // A simple RAII class for timing until the end of the scope.
 954 class ScopedTimer {
 955 public:
 956         ScopedTimer(const string &name, GPUTimers *timers)
 957                 : timers(timers), level(0)
 958         {
 959                 query = timers->begin_timer(name, level);
 960         }
 961
 962         ScopedTimer(const string &name, ScopedTimer *parent_timer)
 963                 : timers(parent_timer->timers),
 964                   level(parent_timer->level + 1)
 965         {
 966                 query = timers->begin_timer(name, level);
 967         }
 968
 969         ~ScopedTimer()
 970         {
 971                 end();
 972         }
 973
 974         void end()
 975         {
 976                 if (!ended) {
 977                         glQueryCounter(query.second, GL_TIMESTAMP);
 978                         ended = true;
 979                 }
 980         }
 981
 982 private:
 983         GPUTimers *timers;
 984         int level;
 985         pair<GLuint, GLuint> query;
 986         bool ended = false;
 987 };
 988
 989 int main(int argc, char **argv)
 990 {
 991         static const option long_options[] = {
 992                 { "alpha", required_argument, 0, 'a' },
 993                 { "delta", required_argument, 0, 'd' },
 994                 { "gamma", required_argument, 0, 'g' }
 995         };
 996
 997         for ( ;; ) {
 998                 int option_index = 0;
 999                 int c = getopt_long(argc, argv, "a:d:g:", long_options, &option_index);
1000
1001                 if (c == -1) {
1002                         break;
1003                 }
1004                 switch (c) {
1005                 case 'a':
1006                         vr_alpha = atof(optarg);
1007                         break;
1008                 case 'd':
1009                         vr_delta = atof(optarg);
1010                         break;
1011                 case 'g':
1012                         vr_gamma = atof(optarg);
1013                         break;
1014                 default:
1015                         fprintf(stderr, "Unknown option '%s'\n", argv[option_index]);
1016                         exit(1);
1017                 };
1018         }
1019
1020         if (SDL_Init(SDL_INIT_EVERYTHING) == -1) {
1021                 fprintf(stderr, "SDL_Init failed: %s\n", SDL_GetError());
1022                 exit(1);
1023         }
1024         SDL_GL_SetAttribute(SDL_GL_ALPHA_SIZE, 8);
1025         SDL_GL_SetAttribute(SDL_GL_DEPTH_SIZE, 0);
1026         SDL_GL_SetAttribute(SDL_GL_STENCIL_SIZE, 0);
1027         SDL_GL_SetAttribute(SDL_GL_DOUBLEBUFFER, 1);
1028
1029         SDL_GL_SetAttribute(SDL_GL_CONTEXT_PROFILE_MASK, SDL_GL_CONTEXT_PROFILE_CORE);
1030         SDL_GL_SetAttribute(SDL_GL_CONTEXT_MAJOR_VERSION, 4);
1031         SDL_GL_SetAttribute(SDL_GL_CONTEXT_MINOR_VERSION, 5);
1032         // SDL_GL_SetAttribute(SDL_GL_CONTEXT_FLAGS, SDL_GL_CONTEXT_DEBUG_FLAG);
1033         SDL_Window *window = SDL_CreateWindow("OpenGL window",
1034                         SDL_WINDOWPOS_UNDEFINED,
1035                         SDL_WINDOWPOS_UNDEFINED,
1036                         64, 64,
1037                         SDL_WINDOW_OPENGL);
1038         SDL_GLContext context = SDL_GL_CreateContext(window);
1039         assert(context != nullptr);
1040
1041         // Load pictures.
1042         unsigned width1, height1, width2, height2;
1043         GLuint tex0 = load_texture(argc >= (optind + 1) ? argv[optind] : "test1499.png", &width1, &height1);
1044         GLuint tex1 = load_texture(argc >= (optind + 2) ? argv[optind + 1] : "test1500.png", &width2, &height2);
1045
1046         if (width1 != width2 || height1 != height2) {
1047                 fprintf(stderr, "Image dimensions don't match (%dx%d versus %dx%d)\n",
1048                         width1, height1, width2, height2);
1049                 exit(1);
1050         }
1051
1052         // Make some samplers.
1053         glCreateSamplers(1, &nearest_sampler);
1054         glSamplerParameteri(nearest_sampler, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
1055         glSamplerParameteri(nearest_sampler, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
1056         glSamplerParameteri(nearest_sampler, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
1057         glSamplerParameteri(nearest_sampler, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
1058
1059         glCreateSamplers(1, &linear_sampler);
1060         glSamplerParameteri(linear_sampler, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
1061         glSamplerParameteri(linear_sampler, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
1062         glSamplerParameteri(linear_sampler, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
1063         glSamplerParameteri(linear_sampler, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
1064
1065         // The smoothness is sampled so that once we get to a smoothness involving
1066         // a value outside the border, the diffusivity between the two becomes zero.
1067         glCreateSamplers(1, &smoothness_sampler);
1068         glSamplerParameteri(smoothness_sampler, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
1069         glSamplerParameteri(smoothness_sampler, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
1070         glSamplerParameteri(smoothness_sampler, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_BORDER);
1071         glSamplerParameteri(smoothness_sampler, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_BORDER);
1072         float zero[] = { 0.0f, 0.0f, 0.0f, 0.0f };
1073         glSamplerParameterfv(smoothness_sampler, GL_TEXTURE_BORDER_COLOR, zero);
1074
1075         float vertices[] = {
1076                 0.0f, 1.0f,
1077                 0.0f, 0.0f,
1078                 1.0f, 1.0f,
1079                 1.0f, 0.0f,
1080         };
1081         glCreateBuffers(1, &vertex_vbo);
1082         glNamedBufferData(vertex_vbo, sizeof(vertices), vertices, GL_STATIC_DRAW);
1083         glBindBuffer(GL_ARRAY_BUFFER, vertex_vbo);
1084
1085         // Initial flow is zero, 1x1.
1086         GLuint initial_flow_tex;
1087         glCreateTextures(GL_TEXTURE_2D, 1, &initial_flow_tex);
1088         glTextureStorage2D(initial_flow_tex, 1, GL_RG16F, 1, 1);
1089         glClearTexImage(initial_flow_tex, 0, GL_RG, GL_FLOAT, nullptr);
1090         int prev_level_width = 1, prev_level_height = 1;
1091
1092         GLuint prev_level_flow_tex = initial_flow_tex;
1093
1094         Sobel sobel;
1095         MotionSearch motion_search;
1096         Densify densify;
1097         Prewarp prewarp;
1098         Derivatives derivatives;
1099         ComputeSmoothness compute_smoothness;
1100         SetupEquations setup_equations;
1101         SOR sor;
1102         AddBaseFlow add_base_flow;
1103         ResizeFlow resize_flow;
1104
1105         GLuint query;
1106         glGenQueries(1, &query);
1107         glBeginQuery(GL_TIME_ELAPSED, query);
1108
1109         GPUTimers timers;
1110
1111         ScopedTimer total_timer("Total", &timers);
1112         for (int level = coarsest_level; level >= int(finest_level); --level) {
1113                 char timer_name[256];
1114                 snprintf(timer_name, sizeof(timer_name), "Level %d", level);
1115                 ScopedTimer level_timer(timer_name, &total_timer);
1116
1117                 int level_width = width1 >> level;
1118                 int level_height = height1 >> level;
1119                 float patch_spacing_pixels = patch_size_pixels * (1.0f - patch_overlap_ratio);
1120                 int width_patches = 1 + lrintf((level_width - patch_size_pixels) / patch_spacing_pixels);
1121                 int height_patches = 1 + lrintf((level_height - patch_size_pixels) / patch_spacing_pixels);
1122
1123                 // Make sure we always read from the correct level; the chosen
1124                 // mipmapping could otherwise be rather unpredictable, especially
1125                 // during motion search.
1126                 // TODO: create these beforehand, and stop leaking them.
1127                 GLuint tex0_view, tex1_view;
1128                 glGenTextures(1, &tex0_view);
1129                 glTextureView(tex0_view, GL_TEXTURE_2D, tex0, GL_R8, level, 1, 0, 1);
1130                 glGenTextures(1, &tex1_view);
1131                 glTextureView(tex1_view, GL_TEXTURE_2D, tex1, GL_R8, level, 1, 0, 1);
1132
1133                 // Create a new texture; we could be fancy and render use a multi-level
1134                 // texture, but meh.
1135                 GLuint grad0_tex;
1136                 glCreateTextures(GL_TEXTURE_2D, 1, &grad0_tex);
1137                 glTextureStorage2D(grad0_tex, 1, GL_RG16F, level_width, level_height);
1138
1139                 // Find the derivative.
1140                 {
1141                         ScopedTimer timer("Sobel", &level_timer);
1142                         sobel.exec(tex0_view, grad0_tex, level_width, level_height);
1143                 }
1144
1145                 // Motion search to find the initial flow. We use the flow from the previous
1146                 // level (sampled bilinearly; no fancy tricks) as a guide, then search from there.
1147
1148                 // Create an output flow texture.
1149                 GLuint flow_out_tex;
1150                 glCreateTextures(GL_TEXTURE_2D, 1, &flow_out_tex);
1151                 glTextureStorage2D(flow_out_tex, 1, GL_RGB16F, width_patches, height_patches);
1152
1153                 // And draw.
1154                 {
1155                         ScopedTimer timer("Motion search", &level_timer);
1156                         motion_search.exec(tex0_view, tex1_view, grad0_tex, prev_level_flow_tex, flow_out_tex, level_width, level_height, prev_level_width, prev_level_height, width_patches, height_patches);
1157                 }
1158
1159                 // Densification.
1160
1161                 // Set up an output texture (initially zero).
1162                 GLuint dense_flow_tex;
1163                 glCreateTextures(GL_TEXTURE_2D, 1, &dense_flow_tex);
1164                 glTextureStorage2D(dense_flow_tex, 1, GL_RGB16F, level_width, level_height);
1165                 glClearTexImage(dense_flow_tex, 0, GL_RGB, GL_FLOAT, nullptr);
1166
1167                 // And draw.
1168                 {
1169                         ScopedTimer timer("Densification", &level_timer);
1170                         densify.exec(tex0_view, tex1_view, flow_out_tex, dense_flow_tex, level_width, level_height, width_patches, height_patches);
1171                 }
1172
1173                 // Everything below here in the loop belongs to variational refinement.
1174                 ScopedTimer varref_timer("Variational refinement", &level_timer);
1175
1176                 // Prewarping; create I and I_t, and a normalized base flow (so we don't
1177                 // have to normalize it over and over again, and also save some bandwidth).
1178                 //
1179                 // During the entire rest of the variational refinement, flow will be measured
1180                 // in pixels, not 0..1 normalized OpenGL texture coordinates.
1181                 // This is because variational refinement depends so heavily on derivatives,
1182                 // which are measured in intensity levels per pixel.
1183                 GLuint I_tex, I_t_tex, base_flow_tex;
1184                 glCreateTextures(GL_TEXTURE_2D, 1, &I_tex);
1185                 glCreateTextures(GL_TEXTURE_2D, 1, &I_t_tex);
1186                 glCreateTextures(GL_TEXTURE_2D, 1, &base_flow_tex);
1187                 glTextureStorage2D(I_tex, 1, GL_R16F, level_width, level_height);
1188                 glTextureStorage2D(I_t_tex, 1, GL_R16F, level_width, level_height);
1189                 glTextureStorage2D(base_flow_tex, 1, GL_RG16F, level_width, level_height);
1190                 {
1191                         ScopedTimer timer("Prewarping", &varref_timer);
1192                         prewarp.exec(tex0_view, tex1_view, dense_flow_tex, I_tex, I_t_tex, base_flow_tex, level_width, level_height);
1193                 }
1194
1195                 // Calculate I_x and I_y. We're only calculating first derivatives;
1196                 // the others will be taken on-the-fly in order to sample from fewer
1197                 // textures overall, since sampling from the L1 cache is cheap.
1198                 // (TODO: Verify that this is indeed faster than making separate
1199                 // double-derivative textures.)
1200                 GLuint I_x_y_tex, beta_0_tex;
1201                 glCreateTextures(GL_TEXTURE_2D, 1, &I_x_y_tex);
1202                 glCreateTextures(GL_TEXTURE_2D, 1, &beta_0_tex);
1203                 glTextureStorage2D(I_x_y_tex, 1, GL_RG16F, level_width, level_height);
1204                 glTextureStorage2D(beta_0_tex, 1, GL_R16F, level_width, level_height);
1205                 {
1206                         ScopedTimer timer("First derivatives", &varref_timer);
1207                         derivatives.exec(I_tex, I_x_y_tex, beta_0_tex, level_width, level_height);
1208                 }
1209
1210                 // We need somewhere to store du and dv (the flow increment, relative
1211                 // to the non-refined base flow u0 and v0). It starts at zero.
1212                 GLuint du_dv_tex;
1213                 glCreateTextures(GL_TEXTURE_2D, 1, &du_dv_tex);
1214                 glTextureStorage2D(du_dv_tex, 1, GL_RG16F, level_width, level_height);
1215                 glClearTexImage(du_dv_tex, 0, GL_RG, GL_FLOAT, nullptr);
1216
1217                 // And for smoothness.
1218                 GLuint smoothness_x_tex, smoothness_y_tex;
1219                 glCreateTextures(GL_TEXTURE_2D, 1, &smoothness_x_tex);
1220                 glCreateTextures(GL_TEXTURE_2D, 1, &smoothness_y_tex);
1221                 glTextureStorage2D(smoothness_x_tex, 1, GL_R16F, level_width, level_height);
1222                 glTextureStorage2D(smoothness_y_tex, 1, GL_R16F, level_width, level_height);
1223
1224                 // And finally for the equation set. See SetupEquations for
1225                 // the storage format.
1226                 GLuint equation_tex;
1227                 glCreateTextures(GL_TEXTURE_2D, 1, &equation_tex);
1228                 glTextureStorage2D(equation_tex, 1, GL_RGBA32UI, level_width, level_height);
1229
1230                 for (int outer_idx = 0; outer_idx < level + 1; ++outer_idx) {
1231                         // Calculate the smoothness terms between the neighboring pixels,
1232                         // both in x and y direction.
1233                         {
1234                                 ScopedTimer timer("Compute smoothness", &varref_timer);
1235                                 compute_smoothness.exec(base_flow_tex, du_dv_tex, smoothness_x_tex, smoothness_y_tex, level_width, level_height);
1236                         }
1237
1238                         // Set up the 2x2 equation system for each pixel.
1239                         {
1240                                 ScopedTimer timer("Set up equations", &varref_timer);
1241                                 setup_equations.exec(I_x_y_tex, I_t_tex, du_dv_tex, base_flow_tex, beta_0_tex, smoothness_x_tex, smoothness_y_tex, equation_tex, level_width, level_height);
1242                         }
1243
1244                         // Run a few SOR (or quasi-SOR, since we're not really Jacobi) iterations.
1245                         // Note that these are to/from the same texture.
1246                         {
1247                                 ScopedTimer timer("SOR", &varref_timer);
1248                                 sor.exec(du_dv_tex, equation_tex, smoothness_x_tex, smoothness_y_tex, level_width, level_height, 5);
1249                         }
1250                 }
1251
1252                 // Add the differential flow found by the variational refinement to the base flow,
1253                 // giving the final flow estimate for this level.
1254                 // The output is in diff_flow_tex; we don't need to make a new texture.
1255                 // You can comment out this part if you wish to test disabling of the variational refinement.
1256                 {
1257                         ScopedTimer timer("Add differential flow", &varref_timer);
1258                         add_base_flow.exec(base_flow_tex, du_dv_tex, level_width, level_height);
1259                 }
1260
1261                 prev_level_flow_tex = base_flow_tex;
1262                 prev_level_width = level_width;
1263                 prev_level_height = level_height;
1264         }
1265         total_timer.end();
1266
1267         timers.print();
1268
1269         // Scale up the flow to the final size (if needed).
1270         GLuint final_tex;
1271         if (finest_level == 0) {
1272                 final_tex = prev_level_flow_tex;
1273         } else {
1274                 glCreateTextures(GL_TEXTURE_2D, 1, &final_tex);
1275                 glTextureStorage2D(final_tex, 1, GL_RG16F, width1, height1);
1276                 resize_flow.exec(prev_level_flow_tex, final_tex, prev_level_width, prev_level_height, width1, height1);
1277         }
1278
1279         unique_ptr<float[]> dense_flow(new float[width1 * height1 * 2]);
1280         glGetTextureImage(final_tex, 0, GL_RG, GL_FLOAT, width1 * height1 * 2 * sizeof(float), dense_flow.get());
1281
1282         FILE *fp = fopen("flow.ppm", "wb");
1283         FILE *flowfp = fopen("flow.flo", "wb");
1284         fprintf(fp, "P6\n%d %d\n255\n", width1, height1);
1285         fprintf(flowfp, "FEIH");
1286         fwrite(&width1, 4, 1, flowfp);
1287         fwrite(&height1, 4, 1, flowfp);
1288         for (unsigned y = 0; y < unsigned(height1); ++y) {
1289                 int yy = height1 - y - 1;
1290                 for (unsigned x = 0; x < unsigned(width1); ++x) {
1291                         float du = dense_flow[(yy * width1 + x) * 2 + 0];
1292                         float dv = dense_flow[(yy * width1 + x) * 2 + 1];
1293
1294                         dv = -dv;
1295
1296                         fwrite(&du, 4, 1, flowfp);
1297                         fwrite(&dv, 4, 1, flowfp);
1298
1299                         uint8_t r, g, b;
1300                         flow2rgb(du, dv, &r, &g, &b);
1301                         putc(r, fp);
1302                         putc(g, fp);
1303                         putc(b, fp);
1304                 }
1305         }
1306         fclose(fp);
1307         fclose(flowfp);
1308
1309         fprintf(stderr, "err = %d\n", glGetError());
1310 }