From 47a221ddc7a9466b08c3a5b4a5578c848840e4ca Mon Sep 17 00:00:00 2001 From: sz Date: Tue, 19 Dec 2023 22:06:31 -0600 Subject: [PATCH] Add FountainHeader logic to CimbReader The premise here is to use the symbol decode to compute the expected fountain headers for the upcoming color decode -- and then use those known values to compute the average real-world color values for each relevant tile in the image. We can then use those averages to compute a more accurate color correction matrix. --- src/lib/cimb_translator/CimbReader.cpp | 15 +++++++++++++ src/lib/cimb_translator/CimbReader.h | 4 ++++ src/lib/encoder/Decoder.h | 13 +++++++++--- src/lib/encoder/aligned_stream.h | 29 +++++++++++++++++++------- src/lib/fountain/FountainMetadata.h | 15 +++++++++++-- 5 files changed, 63 insertions(+), 13 deletions(-) diff --git a/src/lib/cimb_translator/CimbReader.cpp b/src/lib/cimb_translator/CimbReader.cpp index 45d1d195..b78ef9c6 100644 --- a/src/lib/cimb_translator/CimbReader.cpp +++ b/src/lib/cimb_translator/CimbReader.cpp @@ -7,6 +7,8 @@ #include "bit_file/bitmatrix.h" #include "chromatic_adaptation/adaptation_transform.h" #include "chromatic_adaptation/color_correction.h" +#include "serialize/format.h" + #include using namespace cimbar; @@ -90,6 +92,7 @@ namespace { CimbReader::CimbReader(const cv::Mat& img, CimbDecoder& decoder, bool needs_sharpen, bool color_correction) : _image(img) + , _fountainColorHeader(0U) , _cellSize(Config::cell_size() + 2) , _positions(Config::cell_spacing(), Config::cells_per_col(), Config::cell_offset(), Config::corner_padding()) , _decoder(decoder) @@ -141,6 +144,18 @@ bool CimbReader::done() const return !_good or _positions.done(); } +void CimbReader::update_metadata(char* buff, unsigned len) +{ + if (len == 0 and _fountainColorHeader.id() == 0) + return; + + if (_fountainColorHeader.id() == 0) + _fountainColorHeader = FountainMetadata(buff, len); + + std::cout << fmt::format("FountainMd {}, {}, {}", _fountainColorHeader.encode_id(), _fountainColorHeader.file_size(), _fountainColorHeader.block_id()) << std::endl; + _fountainColorHeader.increment_block_id(); // we always want to be +1 +} + unsigned CimbReader::num_reads() const { return _positions.size(); diff --git a/src/lib/cimb_translator/CimbReader.h b/src/lib/cimb_translator/CimbReader.h index 64790c6a..c5136e28 100644 --- a/src/lib/cimb_translator/CimbReader.h +++ b/src/lib/cimb_translator/CimbReader.h @@ -6,6 +6,7 @@ #include "PositionData.h" #include "bit_file/bitbuffer.h" +#include "fountain/FountainMetadata.h" #include class CimbReader @@ -18,11 +19,14 @@ class CimbReader unsigned read_color(const PositionData& pos); bool done() const; + void update_metadata(char* buff, unsigned len); + unsigned num_reads() const; protected: cv::Mat _image; bitbuffer _grayscale; + FountainMetadata _fountainColorHeader; unsigned _cellSize; FloodDecodePositions _positions; diff --git a/src/lib/encoder/Decoder.h b/src/lib/encoder/Decoder.h index 14d21965..29cc3ecc 100644 --- a/src/lib/encoder/Decoder.h +++ b/src/lib/encoder/Decoder.h @@ -9,6 +9,7 @@ #include "cimb_translator/Interleave.h" #include +#include #include class Decoder @@ -93,7 +94,7 @@ inline unsigned Decoder::do_decode(CimbReader& reader, STREAM& ostream) // flush symbols reed_solomon_stream rss(ostream, _eccBytes, _eccBlockSize); - bytesDecoded += symbolBits.flush(rss); + symbolBits.flush(rss); // TODO: get fountain headers out of ostream somehow. // after we call symbolBits.flush(), the underlying aligned stream (ostream) may have knowledge about @@ -109,7 +110,7 @@ inline unsigned Decoder::do_decode(CimbReader& reader, STREAM& ostream) } reed_solomon_stream rss(ostream, _eccBytes, _eccBlockSize); - bytesDecoded += colorBits.flush(rss); + bytesDecoded += colorBits.flush(rss); // will return the pos after this flush(), which includes the previous flush()... return bytesDecoded; } @@ -169,7 +170,13 @@ inline unsigned Decoder::decode_fountain(const MAT& img, FOUNTAINSTREAM& ostream // reader takes cimbar::Config::color_mode() ? CimbReader reader(img, _decoder, should_preprocess, color_correction); - aligned_stream aligner(ostream, ostream.chunk_size()); + // maybe give aligned_stream a callback function that can poke the CimbReader on flush()? + // and then in do_decode(), after the first flush, call CimbReader::try_to_make_new_ccm(), + // ... maybe a replacement for the existing updateColorCorrection()? (called in constructor, would now happen later)... + // which reads the values from the image (using what the aligned_stream gave it -- if anything, and maybe the cellpositions from FloodFillDecode?) + // iff we got enough data from aligned_stream to sample enough colors, try_to_make_new_ccm() pushes the updated ccm to CimbDecoder's threadlocal + // if not, we use whatever we previously had in CimbDecoder's threadlocal...? + aligned_stream aligner(ostream, ostream.chunk_size(), 0, std::bind(&CimbReader::update_metadata, &reader, std::placeholders::_1, std::placeholders::_2)); return do_decode(reader, aligner); } diff --git a/src/lib/encoder/aligned_stream.h b/src/lib/encoder/aligned_stream.h index df039461..b1bca176 100644 --- a/src/lib/encoder/aligned_stream.h +++ b/src/lib/encoder/aligned_stream.h @@ -1,20 +1,23 @@ /* This code is subject to the terms of the Mozilla Public License, v.2.0. http://mozilla.org/MPL/2.0/. */ #pragma once +#include #include template class aligned_stream { public: - aligned_stream(STREAM& stream, unsigned align_increment, unsigned align_offset=0) - : _stream(stream) - , _buffer(align_increment, 0) - , _offset(0) - , _alignOffset(align_offset) - , _alignIncrement(align_increment) - , _badChunk(false) - , _good(true) + aligned_stream(STREAM& stream, unsigned align_increment, unsigned align_offset=0, const std::function& on_flush=nullptr) + : _stream(stream) + , _buffer(align_increment, 0) + , _offset(0) + , _alignOffset(align_offset) + , _alignIncrement(align_increment) + , _onFlush(on_flush) + , _badChunk(false) + , _good(true) + { } @@ -64,6 +67,9 @@ class aligned_stream { _badChunk = false; _offset = 0; + // notify callback w/ bad result + if (_onFlush) + _onFlush(nullptr, 0); } else { @@ -100,7 +106,12 @@ class aligned_stream void flush() { if (_offset > 0) + { _stream.write(_buffer.data(), _offset); + if (_onFlush) + _onFlush(_buffer.data(), _offset); + // notify callback w/ header bytes! + } _totalCount += _offset; _offset = 0; } @@ -112,6 +123,8 @@ class aligned_stream unsigned _offset; unsigned _alignOffset; unsigned _alignIncrement; + std::function _onFlush; + bool _badChunk; bool _good; size_t _totalCount = 0; diff --git a/src/lib/fountain/FountainMetadata.h b/src/lib/fountain/FountainMetadata.h index 583a9169..3896773e 100644 --- a/src/lib/fountain/FountainMetadata.h +++ b/src/lib/fountain/FountainMetadata.h @@ -5,6 +5,13 @@ class FountainMetadata { +protected: + static void update_block_id_internal(uint16_t block_id, uint8_t* arr) + { + arr[0] = (block_id >> 8) & 0xFF; + arr[1] = block_id & 0xFF; + } + public: static const unsigned md_size = 6; @@ -14,8 +21,7 @@ class FountainMetadata arr[1] = (size >> 16) & 0xFF; arr[2] = (size >> 8) & 0xFF; arr[3] = size & 0xFF; - arr[4] = (block_id >> 8) & 0xFF; - arr[5] = block_id & 0xFF; + update_block_id_internal(block_id, arr+4); } public: @@ -59,6 +65,11 @@ class FountainMetadata return res; } + void increment_block_id() + { + update_block_id_internal(block_id()+1, data()+4); + } + unsigned file_size() const { const uint8_t* d = data();