Add horizontal carving

src/seam-carving.cpp

@@ -18,6 +18,11 @@ bool test_energy;
 
 #define min(a, b) { (a < b ? a : b) }
 
+// Get index for any table indexed by [width*(i : height) + (j : width)], but a : dim_long, b : dim_large
+#define im_index(a, b) \
+  (vertical ? (width*a + b) : (width*b + a))
+
+
 bool nearly_equal(float a, float b) {
   return std::nextafter(a, std::numeric_limits<float>::lowest()) <= b
     && std::nextafter(a, std::numeric_limits<float>::max()) >= b;
@@ -66,61 +71,65 @@ std::vector<float> energy_e1(std::vector<unsigned char> source, int width, int h
   return energy;
 }
 
-/** Given the energy value, returns the optimal vertical seam */
+/** Given the energy value, returns the optimal seam */
-std::vector<int> optimal_vertical_seam(std::vector<float> energy, int width, int height) {
+std::vector<int> optimal_seam(std::vector<float> energy, int width, int height, bool vertical) {
+  // dyn_energy is indexed by [dim_large*(i : dim_long) + (j : dim_large)]
   std::vector<float> dyn_energy(width*height);
 
+  int dim_large = vertical ? width : height;
+  int dim_long  = vertical ? height : width; // Number of elements in the seam
+
   //* Find an end of the minimal connected vertical/horizontal seam
-  for (auto i=0; i < width; i++) {
+  for (auto i=0; i < dim_large; i++) {
     dyn_energy[i] = energy[i];
   }
 
-  for (auto i=1; i < height; i++) {
+  for (auto i=1; i < dim_long; i++) { // Propagate dyn_energy
-    for (auto j=0; j < width; j++) {
+    for (auto j=0; j < dim_large; j++) {
-      float bot_center = dyn_energy[width*(i-1) + j];
+      float bot_center = dyn_energy[dim_large*(i-1) + j];
-      float bot_left  = (j > 0) ? dyn_energy[width*(i-1) + (j-1)] : __FLT_MAX__;
+      float bot_left  = (j > 0) ? dyn_energy[dim_large*(i-1) + (j-1)] : __FLT_MAX__;
-      float bot_right = (j+1 < width) ? dyn_energy[width*(i-1) + (j+1)] : __FLT_MAX__;
+      float bot_right = (j+1 < dim_large) ? dyn_energy[dim_large*(i-1) + (j+1)] : __FLT_MAX__;
 
-      dyn_energy[width*i+j] = fmin(
+      dyn_energy[dim_large*i+j] = fmin(
         bot_center,
         fmin(
           bot_left,
           bot_right
         )
-      ) + energy[width*i + j];
+      ) + energy[im_index(i, j)];
     }
   }
 
-  std::vector<int> result(height);
+  std::vector<int> result(dim_long);
   // Find the seam end
   int min_idx = -1;
   float min_val = __FLT_MAX__;
-  for (auto j=0; j < width; j++) {
+  for (auto j=0; j < dim_large; j++) {
-    if (min_val > dyn_energy[width*(height-1) + j]) {
+    if (min_val > dyn_energy[dim_large*(dim_long-1) + j]) {
       min_idx = j;
-      min_val = dyn_energy[width*(height-1) + j];
+      min_val = dyn_energy[dim_large*(dim_long-1) + j];
     }
   }
-  result[height-1] = min_idx;
+  result[dim_long-1] = min_idx;
 
   //* Backtracking to find the path
-  for (auto i=height-1; i > 0; i--) {
+  for (auto i=dim_long-1; i > 0; i--) {
     // We want to find either (bot_l, bot_c, bot_r) with dyn_energy[.] = min_val - energy[cur]
     
     //Idea : float next_energy = min_val - energy[width*i + min_idx];
     //! With floats, we don't always have x + y - y == x, so we check is x+y == x+y
     // This define is a bit ugly but 200x faster than using a lambda function
     #define is_next_idx(idx) \
-        (dyn_energy[(i-1)*width + idx]+energy[width*i + min_idx] == min_val)
+        (dyn_energy[(i-1)*dim_large + idx]+energy[im_index(i, min_idx)] == min_val)
 
     if (is_next_idx(min_idx)) {
       // min_idx does not change
-      min_val = dyn_energy[(i-1)*width + min_idx];
+      min_val = dyn_energy[(i-1)*dim_large + min_idx];
     } else if (min_idx > 0 && is_next_idx(min_idx-1)) {
-      min_val = dyn_energy[(i-1)*width + (min_idx-1)];
+      min_val = dyn_energy[(i-1)*dim_large + (min_idx-1)];
       min_idx = min_idx - 1;
-    } else if (min_idx+1 < width && is_next_idx(min_idx+1)) {
+    } else if (min_idx+1 < dim_large && is_next_idx(min_idx+1)) {
-      min_val = dyn_energy[(i-1)*width + (min_idx+1)];
+      min_val = dyn_energy[(i-1)*dim_large + (min_idx+1)];
       min_idx = min_idx + 1;
     } else {
       std::cerr << "Unable to backtrack path !" << std::endl;
@@ -132,27 +141,31 @@ std::vector<int> optimal_vertical_seam(std::vector<float> energy, int width, int
   return result;
 }
 
-/** Carves an image by one vertical seam. Returns the optimal seam used */
+/** Carves an image by one seam. Returns the optimal seam used */
-std::vector<int> carving_step_vertical(const std::vector<unsigned char> source, std::vector<unsigned char> &output, int width, int height, int nbChannels) {
+std::vector<int> carving_step(const std::vector<unsigned char> source, std::vector<unsigned char> &output, int width, int height, int nbChannels, bool vertical) {
-  
   std::vector<float> energy = energy_e1(source, width, height, nbChannels);
-  std::vector<int> opt_seam = optimal_vertical_seam(energy, width, height);
+  std::vector<int> opt_seam = optimal_seam(energy, width, height, vertical);
   std::vector<bool> blacklist(width*height);
 
+  int dim_large = vertical ? width : height;
+  int dim_long  = vertical ? height : width;
+
   for (auto k=0; k < width*height; k++) { blacklist[k] = false; }
-  for (auto i=0; i < height; i++) {
+  for (auto i=0; i < dim_long; i++) {
-    blacklist[i*width+opt_seam[i]] = true;
+    int index = vertical ? opt_seam[i]+i*width : i+width*opt_seam[i];
+    blacklist[index] = true;
   }
 
-  int outWidth = width-1;
+  for (auto i=0; i < dim_long; i++) {
-
-  for (auto i=0; i < height; i++) {
     int cur_j = 0;
-    for (auto j=0; cur_j < outWidth && j < width; j++) {
+    for (auto j=0; cur_j < dim_large-1 && j < dim_large; j++) {
-      if (!blacklist[i*width+j]) {
+      if (!blacklist[im_index(i, j)]) {
-        output[3*(i*(width-1)+cur_j)] =   source[3*(i*width+j)];
+        int out_pixelIndex = 3*(vertical ? ((width-1)*i + cur_j) : (width*cur_j + i));
-        output[3*(i*(width-1)+cur_j)+1] = source[3*(i*width+j)+1];
+        int src_pixelIndex = 3*im_index(i, j);
-        output[3*(i*(width-1)+cur_j)+2] = source[3*(i*width+j)+2];
+
+        output[ out_pixelIndex ] = source[ src_pixelIndex ];
+        output[out_pixelIndex+1] = source[src_pixelIndex+1];
+        output[out_pixelIndex+2] = source[src_pixelIndex+2];
         cur_j++;
       }
     }
@@ -161,6 +174,72 @@ std::vector<int> carving_step_vertical(const std::vector<unsigned char> source,
 }
 
 
+void seam_carving(unsigned char* source, int width, int height, int nbChannels, const char* out_filename, int nbSeams, bool vertical, bool test_energy=false) {
+  int curWidth = width;
+  int curHeight = height;
+
+  int dim_large = vertical ? width : height;
+  int dim_long  = vertical ? height : width;
+
+  std::vector<unsigned char> carve_output(width*height*3); // Receives at each step the newly carved image
+  std::vector<unsigned char> source_img(width*height*nbChannels); // Contains at each step the carved image
+  std::vector<bool> complete_blacklist(width*height); // Contains all removed pixels, for "test_energy"
+  std::vector<float> ini_energy; // Contains the initial energy, only for "test_energy"
+  std::vector<unsigned char> test_energy_output(width*height*3); // Final output for "test_energy"
+
+
+  for (auto i=0; i < width*height*nbChannels; i++) { source_img[i] = source[i]; }
+
+  if (test_energy) {
+    ini_energy = energy_e1(source_img, width, height, nbChannels);
+    for (auto k=0; k < width*height; k++) { complete_blacklist[k] = false; }
+
+    //* Prepare final output
+    for (auto k=0; k < width*height; k++) {
+      //output[3*k] = source_img[3*k]/3; //* Uncomment if you prefer to see darkened source image
+      //output[3*k+1] = source_img[3*k+1]/3;
+      //output[3*k+2] = source_img[3*k+2]/3;
+      test_energy_output[3*k] = ini_energy[k]*255;
+      test_energy_output[3*k+1] = ini_energy[k]*255;
+      test_energy_output[3*k+2] = ini_energy[k]*255;
+    }
+  }
+
+  SimpleProgressBar::ProgressBar bar(nbSeams);
+  for (auto seam=0; seam < nbSeams; seam++) {
+    std::vector<int> opt_seam = carving_step(source_img, carve_output, curWidth, curHeight, nbChannels, vertical);
+    std::copy(carve_output.begin(), carve_output.end(), source_img.begin());
+
+    if (vertical) // We just reduced the dimension
+      curWidth--;
+    else
+      curHeight--;
+
+    if (test_energy) { // Update blacklist
+      for (auto i=0; i < dim_long; i++) {
+        int j, cur_j = 0; // cur_j is the index relative to the current carved image. j is absolute in the source image
+        for (j=0; j < dim_large && cur_j < opt_seam[i]; j++) {
+          if (!complete_blacklist[im_index(i, j)]) { cur_j++; }
+        }
+        assert(cur_j == opt_seam[i]); // Else, j == width and cur_j is not in the source image..
+  
+        complete_blacklist[im_index(i, j)] = true;
+        test_energy_output[3*im_index(i, j)] = 255; // Set carved pixel to red
+      }
+    }
+    bar.increment();
+    bar.print();
+  }
+  std::cout << std::endl;
+
+  if (test_energy) {
+    export_image(out_filename, test_energy_output.data(), width, height, nbChannels);
+  } else {
+    export_image(out_filename, source_img.data(), curWidth, curHeight, nbChannels);
+  }
+}
+
+
 int main(int argc, char **argv) {
   CLI::App app{"seam-carving"};
   std::string sourceImage;
@@ -169,6 +248,8 @@ int main(int argc, char **argv) {
   app.add_option("-o,--output", outputImage, "Output image")->required();
   int nbSeams = 1;
   app.add_option("-n,--nb-seams", nbSeams, "Number of seams");
+  bool vertical = false;
+  app.add_flag("--vertical", vertical, "Vertical carving");
   silent = false;
   app.add_flag("--silent", silent, "No verbose messages");
   test_energy = false;
@@ -185,70 +266,7 @@ int main(int argc, char **argv) {
     exit(1);
   }
 
-  if (test_energy) { // Vertical seam carving. Prepare
+  seam_carving(source, width, height, nbChannels, outputImage.c_str(), nbSeams, vertical, test_energy=test_energy);
-    std::vector<unsigned char> output(width*height*3); // Final output (initial energy function + where we set removed pixels red)
-    std::vector<unsigned char> carve_output(width*height*3); // Receives at each step the newly carved image
-    std::vector<unsigned char> source_img(width*height*nbChannels); // Contains at each step the carved image
-    std::vector<bool> complete_blacklist(width*height); // Contains all removed pixels
-
-    for (auto k=0; k < width*height; k++) { complete_blacklist[k] = false; }
-    for (auto i=0; i < width*height*nbChannels; i++) { source_img[i] = source[i]; }
-
-    std::vector<float> ini_energy = energy_e1(source_img, width, height, nbChannels);
-
-    //* Prepare final output
-    for (auto k=0; k < width*height; k++) {
-      //output[3*k] = source_img[3*k]/3; //* Uncomment if you prefer to see darkened source image
-      //output[3*k+1] = source_img[3*k+1]/3;
-      //output[3*k+2] = source_img[3*k+2]/3;
-      output[3*k] = ini_energy[k]*255;
-      output[3*k+1] = ini_energy[k]*255;
-      output[3*k+2] = ini_energy[k]*255;
-    }
-
-    int curWidth = width;
-    SimpleProgressBar::ProgressBar bar(nbSeams);
-    for (auto seam=0; seam < nbSeams; seam++) {
-      std::vector<int> opt_seam = carving_step_vertical(source_img, carve_output, curWidth, height, nbChannels); // Perform a carving step
-      std::copy(carve_output.begin(), carve_output.end(), source_img.begin()); // Copy output to next input
-      curWidth--;
-
-      for (auto i=0; i < height; i++) {
-        int cur_j = 0; // cur_j is the index relative to the current carved image. j is absolute in the source image
-        int j;
-        for (j=0; j < width && cur_j < opt_seam[i]; j++) {
-          if (!complete_blacklist[i*width+j]) { cur_j++; }
-        }
-        assert(cur_j == opt_seam[i]); // Else, j == width and cur_j is not in the source image..
-
-        complete_blacklist[i*width+j] = true;
-        output[3*(i*width+j)] = 255;
-      }
-      bar.increment();
-      bar.print();
-    }
-    std::cout << std::endl;
-
-    export_image(outputImage.c_str(), output.data(), width, height, nbChannels);
-
-  } else { // Standard vertical seam carving
-    int outWidth = width;
-    std::vector<unsigned char> output(width*height*3);
-    std::vector<unsigned char> vect_source(width*height*nbChannels);
-
-    for (auto i=0; i < width*height*nbChannels; i++) { vect_source[i] = source[i]; }
-
-    SimpleProgressBar::ProgressBar bar(nbSeams);
-    for (auto seam=0; seam < nbSeams; seam++) {
-      carving_step_vertical(vect_source, output, outWidth, height, nbChannels);
-      std::copy(output.begin(), output.end(), vect_source.begin());
-      outWidth--;
-      bar.increment();
-      bar.print();
-    }
-    std::cout << std::endl;
-    export_image(outputImage.c_str(), vect_source.data(), outWidth, height, nbChannels);
-  }
 
   stbi_image_free(source);
   exit(0);