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