]> git.sesse.net Git - nageru/blobdiff - flow.cpp
Fix a typo.
[nageru] / flow.cpp
index b66642dbbc0c97cdae7bc769f9264d764dabbad7..5456d45121bdd8fafcf3c9e3e36890a0f9a12a87 100644 (file)
--- a/flow.cpp
+++ b/flow.cpp
@@ -1,8 +1,5 @@
 #define NO_SDL_GLEXT 1
 
-#define WIDTH 1280
-#define HEIGHT 720
-
 #include <epoxy/gl.h>
 
 #include <SDL2/SDL.h>
@@ -115,26 +112,41 @@ GLuint compile_shader(const string &shader_src, GLenum type)
        return obj;
 }
 
-
-GLuint load_texture(const char *filename, unsigned width, unsigned height)
+GLuint load_texture(const char *filename, unsigned *width_ret, unsigned *height_ret)
 {
-       FILE *fp = fopen(filename, "rb");
-       if (fp == nullptr) {
-               perror(filename);
+       SDL_Surface *surf = IMG_Load(filename);
+       if (surf == nullptr) {
+               fprintf(stderr, "IMG_Load(%s): %s\n", filename, IMG_GetError());
                exit(1);
        }
-       unique_ptr<uint8_t[]> pix(new uint8_t[width * height]);
-       if (fread(pix.get(), width * height, 1, fp) != 1) {
-               fprintf(stderr, "Short read from %s\n", filename);
+
+       // For whatever reason, SDL doesn't support converting to YUV surfaces
+       // nor grayscale, so we'll do it (slowly) ourselves.
+       SDL_Surface *rgb_surf = SDL_ConvertSurfaceFormat(surf, SDL_PIXELFORMAT_RGBA8888, /*flags=*/0);
+       if (rgb_surf == nullptr) {
+               fprintf(stderr, "SDL_ConvertSurfaceFormat(%s): %s\n", filename, SDL_GetError());
                exit(1);
        }
-       fclose(fp);
 
-       // Convert to bottom-left origin.
-       for (unsigned y = 0; y < height / 2; ++y) {
+       SDL_FreeSurface(surf);
+
+       unsigned width = rgb_surf->w, height = rgb_surf->h;
+       const uint8_t *sptr = (uint8_t *)rgb_surf->pixels;
+       unique_ptr<uint8_t[]> pix(new uint8_t[width * height]);
+
+       // Extract the Y component, and convert to bottom-left origin.
+       for (unsigned y = 0; y < height; ++y) {
                unsigned y2 = height - 1 - y;
-               swap_ranges(&pix[y * width], &pix[y * width + width], &pix[y2 * width]);
+               for (unsigned x = 0; x < width; ++x) {
+                       uint8_t r = sptr[(y2 * width + x) * 4 + 3];
+                       uint8_t g = sptr[(y2 * width + x) * 4 + 2];
+                       uint8_t b = sptr[(y2 * width + x) * 4 + 1];
+
+                       // Rec. 709.
+                       pix[y * width + x] = lrintf(r * 0.2126f + g * 0.7152f + b * 0.0722f);
+               }
        }
+       SDL_FreeSurface(rgb_surf);
 
        int levels = 1;
        for (int w = width, h = height; w > 1 || h > 1; ) {
@@ -149,6 +161,9 @@ GLuint load_texture(const char *filename, unsigned width, unsigned height)
        glTextureSubImage2D(tex, 0, 0, 0, width, height, GL_RED, GL_UNSIGNED_BYTE, pix.get());
        glGenerateTextureMipmap(tex);
 
+       *width_ret = width;
+       *height_ret = height;
+
        return tex;
 }
 
@@ -268,7 +283,7 @@ void Sobel::exec(GLint tex0_view, GLint grad0_tex, int level_width, int level_he
 class MotionSearch {
 public:
        MotionSearch();
-       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 width_patches, int height_patches);
+       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);
 
 private:
        GLuint motion_vs_obj;
@@ -276,7 +291,7 @@ private:
        GLuint motion_search_program;
        GLuint motion_search_vao;
 
-       GLuint uniform_image_size, uniform_inv_image_size;
+       GLuint uniform_image_size, uniform_inv_image_size, uniform_inv_prev_level_size;
        GLuint uniform_image0_tex, uniform_image1_tex, uniform_grad0_tex, uniform_flow_tex;
 };
 
@@ -297,13 +312,14 @@ MotionSearch::MotionSearch()
 
        uniform_image_size = glGetUniformLocation(motion_search_program, "image_size");
        uniform_inv_image_size = glGetUniformLocation(motion_search_program, "inv_image_size");
+       uniform_inv_prev_level_size = glGetUniformLocation(motion_search_program, "inv_prev_level_size");
        uniform_image0_tex = glGetUniformLocation(motion_search_program, "image0_tex");
        uniform_image1_tex = glGetUniformLocation(motion_search_program, "image1_tex");
        uniform_grad0_tex = glGetUniformLocation(motion_search_program, "grad0_tex");
        uniform_flow_tex = glGetUniformLocation(motion_search_program, "flow_tex");
 }
 
-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 width_patches, int height_patches)
+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)
 {
        glUseProgram(motion_search_program);
 
@@ -314,6 +330,7 @@ void MotionSearch::exec(GLuint tex0_view, GLuint tex1_view, GLuint grad0_tex, GL
 
        glProgramUniform2f(motion_search_program, uniform_image_size, level_width, level_height);
        glProgramUniform2f(motion_search_program, uniform_inv_image_size, 1.0f / level_width, 1.0f / level_height);
+       glProgramUniform2f(motion_search_program, uniform_inv_prev_level_size, 1.0f / prev_level_width, 1.0f / prev_level_height);
 
        GLuint flow_fbo;  // TODO: cleanup
        glCreateFramebuffers(1, &flow_fbo);
@@ -409,11 +426,14 @@ void Densify::exec(GLuint tex0_view, GLuint tex1_view, GLuint flow_tex, GLuint d
 // I_0 and I_w. The prewarping is what enables us to solve the variational
 // flow for du,dv instead of u,v.
 //
+// Also calculates the normalized flow, ie. divides by z (this is needed because
+// Densify works by additive blending) and multiplies by the image size.
+//
 // See variational_refinement.txt for more information.
 class Prewarp {
 public:
        Prewarp();
-       void exec(GLuint tex0_view, GLuint tex1_view, GLuint flow_tex, GLuint I_tex, GLuint I_t_tex, int level_width, int level_height);
+       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);
 
 private:
        GLuint prewarp_vs_obj;
@@ -422,6 +442,7 @@ private:
        GLuint prewarp_vao;
 
        GLuint uniform_image0_tex, uniform_image1_tex, uniform_flow_tex;
+       GLuint uniform_image_size;
 };
 
 Prewarp::Prewarp()
@@ -442,9 +463,11 @@ Prewarp::Prewarp()
        uniform_image0_tex = glGetUniformLocation(prewarp_program, "image0_tex");
        uniform_image1_tex = glGetUniformLocation(prewarp_program, "image1_tex");
        uniform_flow_tex = glGetUniformLocation(prewarp_program, "flow_tex");
+
+       uniform_image_size = glGetUniformLocation(prewarp_program, "image_size");
 }
 
-void Prewarp::exec(GLuint tex0_view, GLuint tex1_view, GLuint flow_tex, GLuint I_tex, GLuint I_t_tex, int level_width, int level_height)
+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)
 {
        glUseProgram(prewarp_program);
 
@@ -452,12 +475,15 @@ void Prewarp::exec(GLuint tex0_view, GLuint tex1_view, GLuint flow_tex, GLuint I
        bind_sampler(prewarp_program, uniform_image1_tex, 1, tex1_view, linear_sampler);
        bind_sampler(prewarp_program, uniform_flow_tex, 2, flow_tex, nearest_sampler);
 
+       glProgramUniform2f(prewarp_program, uniform_image_size, level_width, level_height);
+
        GLuint prewarp_fbo;  // TODO: cleanup
        glCreateFramebuffers(1, &prewarp_fbo);
-       GLenum bufs[] = { GL_COLOR_ATTACHMENT0, GL_COLOR_ATTACHMENT1 };
-       glNamedFramebufferDrawBuffers(prewarp_fbo, 2, bufs);
+       GLenum bufs[] = { GL_COLOR_ATTACHMENT0, GL_COLOR_ATTACHMENT1, GL_COLOR_ATTACHMENT2 };
+       glNamedFramebufferDrawBuffers(prewarp_fbo, 3, bufs);
        glNamedFramebufferTexture(prewarp_fbo, GL_COLOR_ATTACHMENT0, I_tex, 0);
        glNamedFramebufferTexture(prewarp_fbo, GL_COLOR_ATTACHMENT1, I_t_tex, 0);
+       glNamedFramebufferTexture(prewarp_fbo, GL_COLOR_ATTACHMENT2, normalized_flow_tex, 0);
 
        glViewport(0, 0, level_width, level_height);
        glDisable(GL_BLEND);
@@ -582,19 +608,17 @@ void ComputeSmoothness::exec(GLuint flow_tex, GLuint diff_flow_tex, GLuint smoot
 
        glViewport(0, 0, level_width, level_height);
 
-       // Make sure the smoothness on the right and upper borders is zero.
-       // We could have done this by making a (W-1)x(H-1) texture instead
-       // (we're sampling smoothness with all-zero border color), but we'd
-       // have to adjust the sampling coordinates, which is annoying.
-       glScissor(0, 0, level_width - 1, level_height - 1);
-       glEnable(GL_SCISSOR_TEST);
-
        glDisable(GL_BLEND);
        glBindVertexArray(smoothness_vao);
        glBindFramebuffer(GL_FRAMEBUFFER, smoothness_fbo);
        glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
 
-       glDisable(GL_SCISSOR_TEST);
+       // Make sure the smoothness on the right and upper borders is zero.
+       // We could have done this by making (W-1)xH and Wx(H-1) textures instead
+       // (we're sampling smoothness with all-zero border color), but we'd
+       // have to adjust the sampling coordinates, which is annoying.
+       glClearTexSubImage(smoothness_x_tex, 0,  level_width - 1, 0, 0,   1, level_height, 1,  GL_RED, GL_FLOAT, nullptr);
+       glClearTexSubImage(smoothness_y_tex, 0,  0, level_height - 1, 0,  level_width, 1, 1,   GL_RED, GL_FLOAT, nullptr);
 }
 
 // Set up the equations set (two equations in two unknowns, per pixel).
@@ -619,10 +643,9 @@ private:
        GLuint equations_vao;
 
        GLuint uniform_I_x_y_tex, uniform_I_t_tex;
-       GLuint uniform_diff_flow_tex, uniform_flow_tex;
+       GLuint uniform_diff_flow_tex, uniform_base_flow_tex;
        GLuint uniform_beta_0_tex;
        GLuint uniform_smoothness_x_tex, uniform_smoothness_y_tex;
-
 };
 
 SetupEquations::SetupEquations()
@@ -643,20 +666,20 @@ SetupEquations::SetupEquations()
        uniform_I_x_y_tex = glGetUniformLocation(equations_program, "I_x_y_tex");
        uniform_I_t_tex = glGetUniformLocation(equations_program, "I_t_tex");
        uniform_diff_flow_tex = glGetUniformLocation(equations_program, "diff_flow_tex");
-       uniform_flow_tex = glGetUniformLocation(equations_program, "flow_tex");
+       uniform_base_flow_tex = glGetUniformLocation(equations_program, "base_flow_tex");
        uniform_beta_0_tex = glGetUniformLocation(equations_program, "beta_0_tex");
        uniform_smoothness_x_tex = glGetUniformLocation(equations_program, "smoothness_x_tex");
        uniform_smoothness_y_tex = glGetUniformLocation(equations_program, "smoothness_y_tex");
 }
 
-void SetupEquations::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)
+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)
 {
        glUseProgram(equations_program);
 
        bind_sampler(equations_program, uniform_I_x_y_tex, 0, I_x_y_tex, nearest_sampler);
        bind_sampler(equations_program, uniform_I_t_tex, 1, I_t_tex, nearest_sampler);
        bind_sampler(equations_program, uniform_diff_flow_tex, 2, diff_flow_tex, nearest_sampler);
-       bind_sampler(equations_program, uniform_flow_tex, 3, flow_tex, nearest_sampler);
+       bind_sampler(equations_program, uniform_base_flow_tex, 3, base_flow_tex, nearest_sampler);
        bind_sampler(equations_program, uniform_beta_0_tex, 4, beta_0_tex, nearest_sampler);
        bind_sampler(equations_program, uniform_smoothness_x_tex, 5, smoothness_x_tex, smoothness_sampler);
        bind_sampler(equations_program, uniform_smoothness_y_tex, 6, smoothness_y_tex, smoothness_sampler);
@@ -743,7 +766,7 @@ void SOR::exec(GLuint diff_flow_tex, GLuint equation_tex, GLuint smoothness_x_te
 }
 
 // Simply add the differential flow found by the variational refinement to the base flow.
-// The output is in diff_flow_tex; we don't need to make a new texture.
+// The output is in base_flow_tex; we don't need to make a new texture.
 class AddBaseFlow {
 public:
        AddBaseFlow();
@@ -755,7 +778,7 @@ private:
        GLuint add_flow_program;
        GLuint add_flow_vao;
 
-       GLuint uniform_base_flow_tex;
+       GLuint uniform_diff_flow_tex;
 };
 
 AddBaseFlow::AddBaseFlow()
@@ -773,18 +796,18 @@ AddBaseFlow::AddBaseFlow()
        glEnableVertexArrayAttrib(add_flow_vao, position_attrib);
        glVertexAttribPointer(position_attrib, 2, GL_FLOAT, GL_FALSE, 0, BUFFER_OFFSET(0));
 
-       uniform_base_flow_tex = glGetUniformLocation(add_flow_program, "base_flow_tex");
+       uniform_diff_flow_tex = glGetUniformLocation(add_flow_program, "diff_flow_tex");
 }
 
 void AddBaseFlow::exec(GLuint base_flow_tex, GLuint diff_flow_tex, int level_width, int level_height)
 {
        glUseProgram(add_flow_program);
 
-       bind_sampler(add_flow_program, uniform_base_flow_tex, 0, base_flow_tex, nearest_sampler);
+       bind_sampler(add_flow_program, uniform_diff_flow_tex, 0, diff_flow_tex, nearest_sampler);
 
        GLuint add_flow_fbo;  // TODO: cleanup
        glCreateFramebuffers(1, &add_flow_fbo);
-       glNamedFramebufferTexture(add_flow_fbo, GL_COLOR_ATTACHMENT0, diff_flow_tex, 0);
+       glNamedFramebufferTexture(add_flow_fbo, GL_COLOR_ATTACHMENT0, base_flow_tex, 0);
 
        glViewport(0, 0, level_width, level_height);
        glEnable(GL_BLEND);
@@ -875,7 +898,7 @@ private:
        bool ended = false;
 };
 
-int main(void)
+int main(int argc, char **argv)
 {
        if (SDL_Init(SDL_INIT_EVERYTHING) == -1) {
                fprintf(stderr, "SDL_Init failed: %s\n", SDL_GetError());
@@ -899,8 +922,15 @@ int main(void)
        assert(context != nullptr);
 
        // Load pictures.
-       GLuint tex0 = load_texture("test1499.pgm", WIDTH, HEIGHT);
-       GLuint tex1 = load_texture("test1500.pgm", WIDTH, HEIGHT);
+       unsigned width1, height1, width2, height2;
+       GLuint tex0 = load_texture(argc >= 2 ? argv[1] : "test1499.png", &width1, &height1);
+       GLuint tex1 = load_texture(argc >= 3 ? argv[2] : "test1500.png", &width2, &height2);
+
+       if (width1 != width2 || height1 != height2) {
+               fprintf(stderr, "Image dimensions don't match (%dx%d versus %dx%d)\n",
+                       width1, height1, width2, height2);
+               exit(1);
+       }
 
        // Make some samplers.
        glCreateSamplers(1, &nearest_sampler);
@@ -939,6 +969,7 @@ int main(void)
        GLuint initial_flow_tex;
        glCreateTextures(GL_TEXTURE_2D, 1, &initial_flow_tex);
        glTextureStorage2D(initial_flow_tex, 1, GL_RG16F, 1, 1);
+       int prev_level_width = 1, prev_level_height = 1;
 
        GLuint prev_level_flow_tex = initial_flow_tex;
 
@@ -964,8 +995,8 @@ int main(void)
                snprintf(timer_name, sizeof(timer_name), "Level %d", level);
                ScopedTimer level_timer(timer_name, &total_timer);
 
-               int level_width = WIDTH >> level;
-               int level_height = HEIGHT >> level;
+               int level_width = width1 >> level;
+               int level_height = height1 >> level;
                float patch_spacing_pixels = patch_size_pixels * (1.0f - patch_overlap_ratio);
                int width_patches = 1 + lrintf((level_width - patch_size_pixels) / patch_spacing_pixels);
                int height_patches = 1 + lrintf((level_height - patch_size_pixels) / patch_spacing_pixels);
@@ -1003,7 +1034,7 @@ int main(void)
                // And draw.
                {
                        ScopedTimer timer("Motion search", &level_timer);
-                       motion_search.exec(tex0_view, tex1_view, grad0_tex, prev_level_flow_tex, flow_out_tex, level_width, level_height, width_patches, height_patches);
+                       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);
                }
 
                // Densification.
@@ -1012,6 +1043,7 @@ int main(void)
                GLuint dense_flow_tex;
                glCreateTextures(GL_TEXTURE_2D, 1, &dense_flow_tex);
                glTextureStorage2D(dense_flow_tex, 1, GL_RGB16F, level_width, level_height);
+               glClearTexImage(dense_flow_tex, 0, GL_RGB, GL_FLOAT, nullptr);
 
                // And draw.
                {
@@ -1022,15 +1054,23 @@ int main(void)
                // Everything below here in the loop belongs to variational refinement.
                ScopedTimer varref_timer("Variational refinement", &level_timer);
 
-               // Prewarping; create I and I_t.
-               GLuint I_tex, I_t_tex;
+               // Prewarping; create I and I_t, and a normalized base flow (so we don't
+               // have to normalize it over and over again, and also save some bandwidth).
+               //
+               // During the entire rest of the variational refinement, flow will be measured
+               // in pixels, not 0..1 normalized OpenGL texture coordinates.
+               // This is because variational refinement depends so heavily on derivatives,
+               // which are measured in intensity levels per pixel.
+               GLuint I_tex, I_t_tex, base_flow_tex;
                glCreateTextures(GL_TEXTURE_2D, 1, &I_tex);
                glCreateTextures(GL_TEXTURE_2D, 1, &I_t_tex);
+               glCreateTextures(GL_TEXTURE_2D, 1, &base_flow_tex);
                glTextureStorage2D(I_tex, 1, GL_R16F, level_width, level_height);
                glTextureStorage2D(I_t_tex, 1, GL_R16F, level_width, level_height);
+               glTextureStorage2D(base_flow_tex, 1, GL_RG16F, level_width, level_height);
                {
                        ScopedTimer timer("Prewarping", &varref_timer);
-                       prewarp.exec(tex0_view, tex1_view, dense_flow_tex, I_tex, I_t_tex, level_width, level_height);
+                       prewarp.exec(tex0_view, tex1_view, dense_flow_tex, I_tex, I_t_tex, base_flow_tex, level_width, level_height);
                }
 
                // Calculate I_x and I_y. We're only calculating first derivatives;
@@ -1053,6 +1093,7 @@ int main(void)
                GLuint du_dv_tex;
                glCreateTextures(GL_TEXTURE_2D, 1, &du_dv_tex);
                glTextureStorage2D(du_dv_tex, 1, GL_RG16F, level_width, level_height);
+               glClearTexImage(du_dv_tex, 0, GL_RG, GL_FLOAT, nullptr);
 
                // And for smoothness.
                GLuint smoothness_x_tex, smoothness_y_tex;
@@ -1072,13 +1113,13 @@ int main(void)
                        // both in x and y direction.
                        {
                                ScopedTimer timer("Compute smoothness", &varref_timer);
-                               compute_smoothness.exec(dense_flow_tex, du_dv_tex, smoothness_x_tex, smoothness_y_tex, level_width, level_height);
+                               compute_smoothness.exec(base_flow_tex, du_dv_tex, smoothness_x_tex, smoothness_y_tex, level_width, level_height);
                        }
 
                        // Set up the 2x2 equation system for each pixel.
                        {
                                ScopedTimer timer("Set up equations", &varref_timer);
-                               setup_equations.exec(I_x_y_tex, I_t_tex, du_dv_tex, dense_flow_tex, beta_0_tex, smoothness_x_tex, smoothness_y_tex, equation_tex, level_width, level_height);
+                               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);
                        }
 
                        // Run a few SOR (or quasi-SOR, since we're not really Jacobi) iterations.
@@ -1092,19 +1133,22 @@ int main(void)
                // Add the differential flow found by the variational refinement to the base flow,
                // giving the final flow estimate for this level.
                // The output is in diff_flow_tex; we don't need to make a new texture.
+               // You can comment out this part if you wish to test disabling of the variational refinement.
                {
                        ScopedTimer timer("Add differential flow", &varref_timer);
-                       add_base_flow.exec(dense_flow_tex, du_dv_tex, level_width, level_height);
+                       add_base_flow.exec(base_flow_tex, du_dv_tex, level_width, level_height);
                }
 
-               prev_level_flow_tex = du_dv_tex;
+               prev_level_flow_tex = base_flow_tex;
+               prev_level_width = level_width;
+               prev_level_height = level_height;
        }
        total_timer.end();
 
        timers.print();
 
-       int level_width = WIDTH >> finest_level;
-       int level_height = HEIGHT >> finest_level;
+       int level_width = width1 >> finest_level;
+       int level_height = height1 >> finest_level;
        unique_ptr<float[]> dense_flow(new float[level_width * level_height * 2]);
        glGetTextureImage(prev_level_flow_tex, 0, GL_RG, GL_FLOAT, level_width * level_height * 2 * sizeof(float), dense_flow.get());
 
@@ -1120,8 +1164,7 @@ int main(void)
                        float du = dense_flow[(yy * level_width + x) * 2 + 0];
                        float dv = dense_flow[(yy * level_width + x) * 2 + 1];
 
-                       du = du * level_width;
-                       dv = -dv * level_height;
+                       dv = -dv;
 
                        fwrite(&du, 4, 1, flowfp);
                        fwrite(&dv, 4, 1, flowfp);