Add support of arbitrary CCM and CFA in finishing. (#54)

* Add support of arbitrary CCM and CFA in finish. * Remove debug print.
timothybrooks · Mar 10, 2020 · 37431dd · 37431dd
1 parent 2c7e2ca
commit 37431dd
Show file tree

Hide file tree

Showing 9 changed files with 174 additions and 96 deletions.
diff --git a/bin/HDRPlus.cpp b/bin/HDRPlus.cpp
@@ -15,35 +15,38 @@
  * processes main stages of the pipeline.
  */
 class HDRPlus {
-    Halide::Runtime::Buffer<uint16_t> imgs;
+    const Burst& burst;
 public:
-    const int width;
-    const int height;
-
-    const BlackPoint bp;
-    const WhitePoint wp;
-    const WhiteBalance wb;
     const Compression c;
     const Gain g;
 
-    HDRPlus(Halide::Runtime::Buffer<uint16_t> imgs, BlackPoint bp, WhitePoint wp, WhiteBalance wb, Compression c, Gain g)
-        : imgs(imgs)
-        , width(imgs.width())
-        , height(imgs.height())
-        , bp(bp)
-        , wp(wp)
-        , wb(wb)
+    HDRPlus(Burst burst, const Compression  c, const Gain g)
+        : burst(burst)
         , c(c)
         , g(g)
     {
-        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.");
-        }
     }
 
     Halide::Runtime::Buffer<uint8_t> process() {
+        const int width = burst.GetWidth();
+        const int height = burst.GetHeight();
+
         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);
+
+        std::cerr << "Black point: " << burst.GetBlackLevel() << std::endl;
+        std::cerr << "White point: " << burst.GetWhiteLevel() << std::endl;
+
+        const WhiteBalance wb = burst.GetWhiteBalance();
+        std::cerr << "RGGB: " << wb.r << " " << wb.g0 << " " << wb.g1 << " " << wb.b << std::endl;
+
+        Halide::Runtime::Buffer<uint16_t> imgs = burst.ToBuffer();
+        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.");
+        }
+
+        const int cfa_pattern = static_cast<int>(burst.GetCfaPattern());
+        auto ccm = burst.GetColorCorrectionMatrix();
+        hdrplus_pipeline(imgs, burst.GetBlackLevel(), burst.GetWhiteLevel(), wb.r, wb.g0, wb.g1, wb.b, cfa_pattern, ccm, c, g, output_img);
 
         // transpose to account for interleaved layout
         output_img.transpose(0, 1);
@@ -53,11 +56,10 @@ class HDRPlus {
     }
 
     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)) {
+        const std::string img_path = dir_path + "/" + img_name;
+
+        const 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;
         }
@@ -108,7 +110,7 @@ int main(int argc, char* argv[]) {
 
     Burst burst(dir_path, in_names);
 
-    HDRPlus hdr_plus(burst.ToBuffer(), burst.GetBlackLevel(), burst.GetWhiteLevel(), burst.GetWhiteBalance(), c, g);
+    HDRPlus hdr_plus(burst, c, g);
 
     Halide::Runtime::Buffer<uint8_t> output = hdr_plus.process();
 

diff --git a/src/Burst.h b/src/Burst.h
@@ -6,7 +6,6 @@
 
 #include <string>
 #include <vector>
-#include <Halide.h>
 
 class Burst {
 public:
@@ -22,12 +21,16 @@ class Burst {
 
     int GetHeight() const { return Raws.empty() ? -1 : Raws[0].GetHeight(); }
 
-    int GetBlackLevel() const { return Raws.empty() ? -1 : Raws[0].GetBlackLevel(); }
+    int GetBlackLevel() const { return Raws.empty() ? -1 : Raws[0].GetScalarBlackLevel(); }
 
     int GetWhiteLevel() const { return Raws.empty() ? -1 : Raws[0].GetWhiteLevel(); }
 
     WhiteBalance GetWhiteBalance() const { return Raws.empty() ? WhiteBalance{-1, -1, -1, -1} : Raws[0].GetWhiteBalance(); }
 
+    CfaPattern GetCfaPattern() const { return Raws.empty() ? CfaPattern::CFA_UNKNOWN : Raws[0].GetCfaPattern(); }
+
+    Halide::Runtime::Buffer<float> GetColorCorrectionMatrix() const { return Raws.empty() ? Halide::Runtime::Buffer<float>() : Raws[0].GetColorCorrectionMatrix(); }
+
     Halide::Runtime::Buffer<uint16_t> ToBuffer() const;
 
     void CopyToBuffer(Halide::Runtime::Buffer<uint16_t>& buffer) const;

diff --git a/src/InputSource.cpp b/src/InputSource.cpp
@@ -1,4 +1,8 @@
 #include "InputSource.h"
+
+#include <algorithm>
+#include <unordered_map>
+
 #include "LibRaw2DngConverter.h"
 
 RawImage::RawImage(const std::string &path)
@@ -46,7 +50,82 @@ void RawImage::CopyToBuffer(Halide::Runtime::Buffer<uint16_t> &buffer) const {
 
 void RawImage::WriteDng(const std::string &output_path, const Halide::Runtime::Buffer<uint16_t> &buffer) const
 {
-    LibRaw2DngConverter converter(*RawProcessor);
+    LibRaw2DngConverter converter(*this);
     converter.SetBuffer(buffer);
     converter.Write(output_path);
 }
+
+std::array<float, 4> RawImage::GetBlackLevel() const {
+    // See https://www.libraw.org/node/2471
+    const auto raw_color = RawProcessor->imgdata.color;
+    const auto base_black_level = static_cast<float>(raw_color.black);
+
+    std::array<float, 4> black_level = {
+        base_black_level + static_cast<float>(raw_color.cblack[0]),
+        base_black_level + static_cast<float>(raw_color.cblack[1]),
+        base_black_level + static_cast<float>(raw_color.cblack[2]),
+        base_black_level + static_cast<float>(raw_color.cblack[3])
+    };
+
+    if (raw_color.cblack[4] == 2 && raw_color.cblack[5] == 2) {
+        for (int x = 0; x < raw_color.cblack[4]; ++x) {
+            for (int y = 0; y < raw_color.cblack[5]; ++y) {
+                const auto index = y * 2 + x;
+                black_level[index] = raw_color.cblack[6 + index];
+            }
+        }
+    }
+
+    return black_level;
+}
+
+int RawImage::GetScalarBlackLevel() const {
+    const auto black_level = GetBlackLevel();
+    return static_cast<int>(*std::min_element(black_level.begin(), black_level.end()));
+}
+
+std::string RawImage::GetCfaPatternString() const {
+    static const std::unordered_map<char, char> CDESC_TO_CFA = {
+        {'R', 0},
+        {'G', 1},
+        {'B', 2},
+        {'r', 0},
+        {'g', 1},
+        {'b', 2}
+    };
+    const auto &cdesc = RawProcessor->imgdata.idata.cdesc;
+    return {
+        CDESC_TO_CFA.at(cdesc[RawProcessor->COLOR(0, 0)]),
+        CDESC_TO_CFA.at(cdesc[RawProcessor->COLOR(0, 1)]),
+        CDESC_TO_CFA.at(cdesc[RawProcessor->COLOR(1, 0)]),
+        CDESC_TO_CFA.at(cdesc[RawProcessor->COLOR(1, 1)])
+    };
+}
+
+CfaPattern RawImage::GetCfaPattern() const {
+    const auto cfa_pattern = GetCfaPatternString();
+    if (cfa_pattern == std::string{0, 1, 1, 2}) {
+        return CfaPattern::CFA_RGGB;
+    } else if (cfa_pattern == std::string{1, 0, 2, 1}) {
+        return CfaPattern::CFA_GRBG;
+    } else if (cfa_pattern == std::string{2, 1, 1, 0}) {
+        return CfaPattern::CFA_BGGR;
+    } else if (cfa_pattern == std::string{1, 2, 0, 1}) {
+        return CfaPattern::CFA_GBRG;
+    }
+    throw std::invalid_argument("Unsupported CFA pattern: " + cfa_pattern);
+    return CfaPattern::CFA_UNKNOWN;
+}
+
+Halide::Runtime::Buffer<float> RawImage::GetColorCorrectionMatrix() const {
+    const auto raw_color = RawProcessor->imgdata.color;
+    Halide::Runtime::Buffer<float> ccm(3, 3);
+    for (int i = 0; i < 3; ++i) {
+        for (int j = 0; j < 3; ++j) {
+            ccm(i, j) = raw_color.rgb_cam[j][i];
+        }
+    }
+    return ccm;
+}
+
+
diff --git a/src/InputSource.h b/src/InputSource.h
@@ -17,17 +17,25 @@ class RawImage {
 
     int GetHeight() const { return RawProcessor->imgdata.rawdata.sizes.height; }
 
-    int GetBlackLevel() const { return RawProcessor->imgdata.color.black; }
+    int GetScalarBlackLevel() const;
+
+    std::array<float, 4> GetBlackLevel() const;
 
     int GetWhiteLevel() const { return RawProcessor->imgdata.color.maximum; }
 
     WhiteBalance GetWhiteBalance() const;
 
+    std::string GetCfaPatternString() const;
+    CfaPattern GetCfaPattern() const;
+
+    Halide::Runtime::Buffer<float> GetColorCorrectionMatrix() const;
+
     void CopyToBuffer(Halide::Runtime::Buffer<uint16_t>& buffer) const;
 
     // Writes current RawImage as DNG. If buffer was provided, then use it instead of internal buffer.
     void WriteDng(const std::string& path, const Halide::Runtime::Buffer<uint16_t>& buffer = {}) const;
 
+    std::shared_ptr<LibRaw> GetRawProcessor() const { return RawProcessor; }
 private:
     std::string Path;
     std::shared_ptr<LibRaw> RawProcessor;

diff --git a/src/LibRaw2DngConverter.cpp b/src/LibRaw2DngConverter.cpp
@@ -2,57 +2,19 @@
 
 #include <unordered_map>
 
-namespace {
-    std::array<float, 4> GetBlackLevel(const libraw_colordata_t &raw_color) {
-        // See https://www.libraw.org/node/2471
-        const auto base_black_level = static_cast<float>(raw_color.black);
-
-        std::array<float, 4> black_level = {
-            base_black_level + static_cast<float>(raw_color.cblack[0]),
-            base_black_level + static_cast<float>(raw_color.cblack[1]),
-            base_black_level + static_cast<float>(raw_color.cblack[2]),
-            base_black_level + static_cast<float>(raw_color.cblack[3])
-        };
-
-        if (raw_color.cblack[4] == 2 && raw_color.cblack[5] == 2) {
-            for (int x = 0; x < raw_color.cblack[4]; ++x) {
-                for (int y = 0; y < raw_color.cblack[5]; ++y) {
-                    const auto index = y * 2 + x;
-                    black_level[index] = raw_color.cblack[6 + index];
-                }
-            }
-        }
-        return black_level;
-    }
+#include <libraw/libraw.h>
 
-    std::string GetCfaPattern(const LibRaw &raw) {
-        static const std::unordered_map<char, char> CDESC_TO_CFA = {
-            {'R', 0},
-            {'G', 1},
-            {'B', 2},
-            {'r', 0},
-            {'g', 1},
-            {'b', 2}
-        };
-        const auto &cdesc = raw.imgdata.idata.cdesc;
-        auto &mutable_raw = const_cast<LibRaw &>(raw); // LibRaw is not const correct :-(
-        return {
-            CDESC_TO_CFA.at(cdesc[mutable_raw.COLOR(0, 0)]),
-            CDESC_TO_CFA.at(cdesc[mutable_raw.COLOR(0, 1)]),
-            CDESC_TO_CFA.at(cdesc[mutable_raw.COLOR(1, 0)]),
-            CDESC_TO_CFA.at(cdesc[mutable_raw.COLOR(1, 1)])
-        };
-    }
-}
+#include "InputSource.h"
 
-LibRaw2DngConverter::LibRaw2DngConverter(const LibRaw &raw)
+LibRaw2DngConverter::LibRaw2DngConverter(const RawImage& raw)
     : OutputStream()
     , Raw(raw)
     , Tiff(SetTiffFields(TiffPtr(TIFFStreamOpen("", &OutputStream), TIFFClose)))
 {}
 
 LibRaw2DngConverter::TiffPtr LibRaw2DngConverter::SetTiffFields(LibRaw2DngConverter::TiffPtr tiff_ptr) {
-    const auto raw_color = Raw.imgdata.color;
+    const auto RawProcessor = Raw.GetRawProcessor();
+    const auto raw_color = RawProcessor->imgdata.color;
 
     const uint16_t bayer_pattern_dimensions[] = {2, 2};
 
@@ -70,7 +32,7 @@ LibRaw2DngConverter::TiffPtr LibRaw2DngConverter::SetTiffFields(LibRaw2DngConver
     TIFFSetField(tiff, TIFFTAG_SAMPLEFORMAT, SAMPLEFORMAT_UINT);
     TIFFSetField(tiff, TIFFTAG_CFAREPEATPATTERNDIM, &bayer_pattern_dimensions);
 
-    const std::string cfa = GetCfaPattern(Raw);
+    const std::string cfa = Raw.GetCfaPatternString();
     TIFFSetField(tiff, TIFFTAG_CFAPATTERN, cfa.c_str());
 
     TIFFSetField(tiff, TIFFTAG_MAKE, "hdr-plus");
@@ -94,15 +56,15 @@ LibRaw2DngConverter::TiffPtr LibRaw2DngConverter::SetTiffFields(LibRaw2DngConver
     TIFFSetField(tiff, TIFFTAG_CFALAYOUT, 1); // Rectangular (or square) layout
     TIFFSetField(tiff, TIFFTAG_CFAPLANECOLOR, 3, "\00\01\02"); // RGB https://www.awaresystems.be/imaging/tiff/tifftags/cfaplanecolor.html
 
-    const std::array<float, 4> black_level = GetBlackLevel(raw_color);
+    const std::array<float, 4> black_level = Raw.GetBlackLevel();
     TIFFSetField(tiff, TIFFTAG_BLACKLEVEL, 4, &black_level);
 
     static const uint32_t white_level = raw_color.maximum;
     TIFFSetField(tiff, TIFFTAG_WHITELEVEL, 1, &white_level);
 
-    if (Raw.imgdata.sizes.flip > 0) {
+    if (RawProcessor->imgdata.sizes.flip > 0) {
         // Seems that LibRaw uses LibTIFF notation.
-        TIFFSetField(tiff, TIFFTAG_ORIENTATION, Raw.imgdata.sizes.flip);
+        TIFFSetField(tiff, TIFFTAG_ORIENTATION, RawProcessor->imgdata.sizes.flip);
     } else {
         TIFFSetField(tiff, TIFFTAG_ORIENTATION, ORIENTATION_TOPLEFT);
     }

diff --git a/src/LibRaw2DngConverter.h b/src/LibRaw2DngConverter.h
@@ -5,20 +5,21 @@
 #include <tiffio.hxx>
 
 #include <HalideBuffer.h>
-#include <libraw/libraw.h>
+
+class RawImage;
 
 class LibRaw2DngConverter {
     using TiffPtr = std::shared_ptr<TIFF>;
     TiffPtr SetTiffFields(TiffPtr tiff_ptr);
 public:
-    explicit LibRaw2DngConverter(const LibRaw& raw);
+    explicit LibRaw2DngConverter(const RawImage& raw);
 
     void SetBuffer(const Halide::Runtime::Buffer<uint16_t>& buffer) const;
 
     void Write(const std::string& path) const;
 
 private:
     std::ostringstream OutputStream;
-    const LibRaw& Raw;
+    const RawImage& Raw;
     std::shared_ptr<TIFF> Tiff;
 };