Use generators (#51)

* Move binaries to separate directory * Make align accept Halide::Func * Make project compilable * Make hdrplus binary working * Add generator for stack_frames * Move binaries to separate directory * Make align accept Halide::Func * Make project compilable * Make hdrplus binary working * Add generator for stack_frames * Borrow changes from halide autoscheduler branch
timothybrooks · Mar 6, 2020 · 1b5fbbe · 1b5fbbe
1 parent 6c9f6d0
commit 1b5fbbe
Show file tree

Hide file tree

Showing 12 changed files with 243 additions and 219 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -17,10 +17,6 @@ find_package(JPEG REQUIRED)
 find_library(LIBRAW_LIBRARY NAMES raw raw_r)
 
 set(src_files
-    src/align.cpp
-    src/finish.cpp
-    src/merge.cpp
-    src/util.cpp
     src/InputSource.cpp
     src/Burst.cpp
     src/LibRaw2DngConverter.cpp)
@@ -31,10 +27,36 @@ set(header_files
     src/dngwriter.h
     src/LibRaw2DngConverter.h)
 
-include_directories(${HALIDE_DISTRIB_DIR}/include ${HALIDE_DISTRIB_DIR}/tools ${RAW2DNG_INCLUDE_DIRS})
-
-add_executable(hdrplus src/HDRPlus.cpp ${src_files})
-target_link_libraries(hdrplus Halide png ${LIBRAW_LIBRARY} ${TIFF_LIBRARIES} ${TIFFXX_LIBRARY})
-
-add_executable(stack_frames src/stack_frames.cpp ${src_files})
-target_link_libraries(stack_frames Halide ${LIBRAW_LIBRARY} ${PNG_LIBRARIES} ${JPEG_LIBRARIES} ${TIFF_LIBRARIES} ${TIFFXX_LIBRARY})
+include_directories(${CMAKE_CURRENT_SOURCE_DIR} ${HALIDE_DISTRIB_DIR}/include ${HALIDE_DISTRIB_DIR}/tools ${RAW2DNG_INCLUDE_DIRS})
+include_directories(${CMAKE_BINARY_DIR}/genfiles)
+
+set(HALIDE_TARGET "")
+set(HALIDE_TARGET_FEATURES "")
+halide_library(hdrplus_pipeline
+    SRCS src/hdrplus_pipeline_generator.cpp src/align.cpp src/merge.cpp src/finish.cpp src/util.cpp
+    # GENERATOR_DEPS  # We don't have any yet
+    GENERATOR_NAME hdrplus_pipeline
+    # GENERATOR_ARGS  # We don't have any yet
+    FUNCTION_NAME hdrplus_pipeline
+    # HALIDE_TARGET ${HALIDE_TARGET}  # TODO: add option with custom HALIDE_TARGET
+    # HALIDE_TARGET_FEATURES ${HALIDE_TARGET_FEATURES}  # TODO: add option with custom HALIDE_TARGET
+    # EXTRA_OUTPUTS "stmt;html;schedule") # uncomment for extra output
+)
+
+halide_library(align_and_merge
+    SRCS src/align_and_merge_generator.cpp src/align.cpp src/merge.cpp src/util.cpp
+    GENERATOR_NAME align_and_merge
+    FUNCTION_NAME align_and_merge
+    # HALIDE_TARGET ${HALIDE_TARGET}  # TODO: add option with custom HALIDE_TARGET
+    # HALIDE_TARGET_FEATURES ${HALIDE_TARGET_FEATURES}  # TODO: add option with custom HALIDE_TARGET
+    # EXTRA_OUTPUTS "stmt;html;schedule") # uncomment for extra output
+)
+
+
+add_executable(hdrplus bin/HDRPlus.cpp ${src_files})
+add_dependencies(hdrplus hdrplus_pipeline)
+target_link_libraries(hdrplus hdrplus_pipeline Halide png ${LIBRAW_LIBRARY} ${TIFF_LIBRARIES} ${TIFFXX_LIBRARY})
+
+add_executable(stack_frames bin/stack_frames.cpp ${src_files})
+add_dependencies(stack_frames align_and_merge)
+target_link_libraries(stack_frames Halide align_and_merge ${LIBRAW_LIBRARY} ${PNG_LIBRARIES} ${JPEG_LIBRARIES} ${TIFF_LIBRARIES} ${TIFFXX_LIBRARY})
diff --git a/src/HDRPlus.cpp → bin/HDRPlus.cpp b/src/HDRPlus.cpp → bin/HDRPlus.cpp
@@ -1,32 +1,22 @@
-#include "Halide.h"
-#include "halide_load_raw.h"
+#include <iostream>
+#include <fstream>
+#include <stdio.h>
 
-#include "Burst.h"
+#include <Halide.h>
 
 #define STB_IMAGE_WRITE_IMPLEMENTATION
-#include "../include/stb_image_write.h"
-#include <stdio.h>
-#include "align.h"
-#include "merge.h"
-#include "finish.h"
-#include <iostream>
-#include <fstream>
+#include <include/stb_image_write.h>
+
+#include <src/Burst.h>
+#include <hdrplus_pipeline/hdrplus_pipeline.h>
 
-using namespace Halide;
-using namespace std;
 /*
  * HDRPlus Class -- Houses file I/O, defines pipeline attributes and calls
  * processes main stages of the pipeline.
  */
 class HDRPlus {
-
-private:
-
-    const Halide::Runtime::Buffer<uint16_t> imgs;
-
+    Halide::Runtime::Buffer<uint16_t> imgs;
 public:
-
-
     const int width;
     const int height;
 
@@ -35,9 +25,6 @@ class HDRPlus {
     const WhiteBalance wb;
     const Compression c;
     const Gain g;
-    char *readfile;
-    std::string readfilestring;
-    void imagesize();
 
     HDRPlus(Halide::Runtime::Buffer<uint16_t> imgs, BlackPoint bp, WhitePoint wp, WhiteBalance wb, Compression c, Gain g)
         : imgs(imgs)
@@ -49,53 +36,31 @@ class HDRPlus {
         , c(c)
         , g(g)
     {
-        assert(imgs.dimensions() == 3);         // width * height * img_idx
-        assert(imgs.extent(2) >= 2);            // must have at least one alternate image
+        if (imgs.dimensions() != 3 || imgs.extent(2) < 2) {
+            throw std::invalid_argument("The input of HDRPlus must be a 3-dimensional buffer with at least two channels.");
+        }
     }
 
-    /*
-     * process -- Calls all of the main stages (align, merge, finish) of the pipeline.
-     */
-    Buffer<uint8_t> process() {
-        Halide::Buffer<uint16_t> imgsBuffer(*imgs.raw_buffer());
-
-        Func alignment = align(imgsBuffer);
-        Func merged = merge(imgsBuffer, alignment);
-        Func finished = finish(merged, width, height, bp, wp, wb, c, g);
-
-        ///////////////////////////////////////////////////////////////////////////
-        // realize image
-        ///////////////////////////////////////////////////////////////////////////
-
-        Buffer<uint8_t> output_img(3, width, height);
-
-        finished.realize(output_img);
+    Halide::Runtime::Buffer<uint8_t> process() {
+        Halide::Runtime::Buffer<uint8_t> output_img(3, width, height);
+        hdrplus_pipeline(imgs, bp, wp, wb.r, wb.g0, wb.g1, wb.b, c, g, output_img);
 
         // transpose to account for interleaved layout
-
         output_img.transpose(0, 1);
         output_img.transpose(1, 2);
 
         return output_img;
     }
 
-    /*
-     * save_png -- Writes an interleaved Halide image to an output file.
-     */
-    static bool save_png(std::string dir_path, std::string img_name, Buffer<uint8_t> &img) {
-
+    static bool save_png(const std::string& dir_path, const std::string& img_name, const Halide::Runtime::Buffer<uint8_t> &img) {
         std::string img_path = dir_path + "/" + img_name;
-
         std::remove(img_path.c_str());
-
         int stride_in_bytes = img.width() * img.channels();
 
         if(!stbi_write_png(img_path.c_str(), img.width(), img.height(), img.channels(), img.data(), stride_in_bytes)) {
-
             std::cerr << "Unable to write output image '" << img_name << "'" << std::endl;
             return false;
         }
-
         return true;
     }
 };
@@ -114,20 +79,15 @@ int main(int argc, char* argv[]) {
     int i = 1;
 
     while(argv[i][0] == '-') {
-
         if(argv[i][1] == 'c') {
-
             c = atof(argv[++i]);
             i++;
             continue;
-        }
-        else if(argv[i][1] == 'g') {
-
+        } else if(argv[i][1] == 'g') {
             g = atof(argv[++i]);
             i++;
             continue;
-        }
-        else {
+        } else {
             std::cerr << "Invalid flag '" << argv[i][1] << "'" << std::endl;
             return 1;
         }
@@ -142,23 +102,15 @@ int main(int argc, char* argv[]) {
     std::string out_name = argv[i++];
 
     std::vector<std::string> in_names;
-    while (i < argc) in_names.push_back(argv[i++]);
+    while (i < argc) {
+        in_names.emplace_back(argv[i++]);
+    }
 
     Burst burst(dir_path, in_names);
-    Halide::Runtime::Buffer<uint16_t> imgs = burst.ToBuffer();
-    if (imgs.channels() < 2) {
-        return EXIT_FAILURE;
-    }
 
-    HDRPlus hdr_plus(
-        imgs,
-        burst.GetBlackLevel(),
-        burst.GetWhiteLevel(),
-        burst.GetWhiteBalance(),
-        c,
-        g);
+    HDRPlus hdr_plus(burst.ToBuffer(), burst.GetBlackLevel(), burst.GetWhiteLevel(), burst.GetWhiteBalance(), c, g);
 
-    Buffer<uint8_t> output = hdr_plus.process();
+    Halide::Runtime::Buffer<uint8_t> output = hdr_plus.process();
 
     if (!HDRPlus::save_png(dir_path, out_name, output)) {
         return EXIT_FAILURE;

diff --git a/src/stack_frames.cpp → bin/stack_frames.cpp b/src/stack_frames.cpp → bin/stack_frames.cpp
@@ -5,27 +5,16 @@
 #include <Halide.h>
 #include <halide_image_io.h>
 
-#include "align.h"
-#include "Burst.h"
-#include "finish.h"
-#include "merge.h"
+#include <src/Burst.h>
 
-using namespace Halide;
-using namespace std;
+#include <align_and_merge/align_and_merge.h>
 
-Halide::Buffer<uint16_t> align_and_merge(const Halide::Runtime::Buffer<uint16_t>& burst) {
+Halide::Runtime::Buffer<uint16_t> align_and_merge(Halide::Runtime::Buffer<uint16_t> burst) {
     if (burst.channels() < 2) {
         return {};
     }
-
-    Halide::Buffer<uint16_t> imgsBuffer(*burst.raw_buffer());
-
-    Func alignment = align(imgsBuffer);
-    Func merged = merge(imgsBuffer, alignment);
-
-    Halide::Buffer<uint16_t> merged_buffer(burst.width(), burst.height());
-    merged.realize(merged_buffer);
-
+    Halide::Runtime::Buffer<uint16_t> merged_buffer(burst.width(), burst.height());
+    align_and_merge(burst, merged_buffer);
     return merged_buffer;
 }
 
@@ -49,12 +38,12 @@ int main(int argc, char* argv[]) {
 
     Burst burst(dir_path, in_names);
 
-    Halide::Buffer<uint16_t> merged = align_and_merge(burst.ToBuffer());
+    const auto merged = align_and_merge(burst.ToBuffer());
     std::cerr << "merged size: " << merged.width() << " " << merged.height() << std::endl;
 
     const RawImage& raw = burst.GetRaw(0);
     const std::string merged_filename = dir_path + "/" + out_name;
-    raw.WriteDng(merged_filename, *merged.get());
+    raw.WriteDng(merged_filename, merged);
 
     return EXIT_SUCCESS;
 }
diff --git a/src/align.cpp b/src/align.cpp
@@ -65,7 +65,7 @@ Func align_layer(Func layer, Func prev_alignment, Point prev_min, Point prev_max
  * by T_SIZE_2 in each dimension. align(imgs)(tile_x, tile_y, n) is a point representing the x and y offset
  * for a tile in layer n that most closely matches that tile in the reference (relative to the reference tile's location)
  */
-Func align(const Buffer<uint16_t> imgs) {
+Func align(const Halide::Func imgs, Halide::Expr width, Halide::Expr height) {
 
     Func alignment_3("layer_3_alignment");
     Func alignment("alignment");
@@ -74,7 +74,7 @@ Func align(const Buffer<uint16_t> imgs) {
 
     // mirror input with overlapping edges
 
-    Func imgs_mirror = BoundaryConditions::mirror_interior(imgs, 0, imgs.width(), 0, imgs.height());
+    Func imgs_mirror = BoundaryConditions::mirror_interior(imgs, 0, width, 0, height);
 
     // downsampled layers for alignment
 
@@ -107,8 +107,8 @@ Func align(const Buffer<uint16_t> imgs) {
 
     // number of tiles in the x and y dimensions
 
-    int num_tx = imgs.width() / T_SIZE_2 - 1;
-    int num_ty = imgs.height() / T_SIZE_2 - 1;
+    Expr num_tx = width / T_SIZE_2 - 1;
+    Expr num_ty = height / T_SIZE_2 - 1;
 
     // final alignment offsets for the original mosaic image; tiles outside of the bounds use the nearest alignment offset
 
@@ -118,3 +118,8 @@ Func align(const Buffer<uint16_t> imgs) {
 
     return alignment_repeat;
 }
+
+Halide::Func align(Halide::Buffer<uint16_t> imgs) {
+    Halide::Func imgs_function(imgs);
+    return align(imgs_function, imgs.width(), imgs.height());
+}
diff --git a/src/align.h b/src/align.h
@@ -55,6 +55,7 @@ inline Halide::Expr idx_layer(Halide::Expr t, Halide::Expr i) { return t * T_SIZ
  * by T_SIZE_2 in each dimension. align(imgs)(tile_x, tile_y, n) is a point representing the x and y offset
  * for a tile in layer n that most closely matches that tile in the reference (relative to the reference tile's location)
  */
-Halide::Func align(const Halide::Buffer<uint16_t> imgs);
+Halide::Func align(Halide::Buffer<uint16_t> imgs);
+Halide::Func align(const Halide::Func imgs, Halide::Expr width, Halide::Expr height);
 
 #endif
diff --git a/src/align_and_merge_generator.cpp b/src/align_and_merge_generator.cpp
@@ -0,0 +1,25 @@
+#include <Halide.h>
+
+#include "align.h"
+#include "merge.h"
+#include "finish.h"
+
+namespace {
+
+    class StackFrames : public Halide::Generator<StackFrames> {
+    public:
+        // 'inputs' is really a series of raw 2d frames; extent[2] specifies the count
+        Input<Halide::Buffer<uint16_t>> inputs{"inputs", 3};
+        // Merged buffer
+        Output<Halide::Buffer<uint16_t>> output{"output", 2};
+
+        void generate() {
+            Func alignment = align(inputs, inputs.width(), inputs.height());
+            Func merged = merge(inputs, inputs.width(), inputs.height(), inputs.dim(2).extent(), alignment);
+            output = merged;
+        }
+    };
+
+}  // namespace
+
+HALIDE_REGISTER_GENERATOR(StackFrames, align_and_merge)