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