Skip to content

Commit

Permalink
Add FountainHeader logic to CimbReader
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
sz3 committed Dec 20, 2023
1 parent 728d542 commit 47a221d
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 13 deletions.
15 changes: 15 additions & 0 deletions src/lib/cimb_translator/CimbReader.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 <opencv2/opencv.hpp>

using namespace cimbar;
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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();
Expand Down
4 changes: 4 additions & 0 deletions src/lib/cimb_translator/CimbReader.h
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
#include "PositionData.h"

#include "bit_file/bitbuffer.h"
#include "fountain/FountainMetadata.h"
#include <opencv2/opencv.hpp>

class CimbReader
Expand All @@ -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;
Expand Down
13 changes: 10 additions & 3 deletions src/lib/encoder/Decoder.h
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include "cimb_translator/Interleave.h"

#include <opencv2/opencv.hpp>
#include <functional>
#include <string>

class Decoder
Expand Down Expand Up @@ -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
Expand All @@ -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;
}

Expand Down Expand Up @@ -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);
}

Expand Down
29 changes: 21 additions & 8 deletions src/lib/encoder/aligned_stream.h
Original file line number Diff line number Diff line change
@@ -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 <functional>
#include <vector>

template <typename STREAM>
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<void(char*,size_t)>& on_flush=nullptr)
: _stream(stream)
, _buffer(align_increment, 0)
, _offset(0)
, _alignOffset(align_offset)
, _alignIncrement(align_increment)
, _onFlush(on_flush)
, _badChunk(false)
, _good(true)

{
}

Expand Down Expand Up @@ -64,6 +67,9 @@ class aligned_stream
{
_badChunk = false;
_offset = 0;
// notify callback w/ bad result
if (_onFlush)
_onFlush(nullptr, 0);
}
else
{
Expand Down Expand Up @@ -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;
}
Expand All @@ -112,6 +123,8 @@ class aligned_stream
unsigned _offset;
unsigned _alignOffset;
unsigned _alignIncrement;
std::function<void(char*,size_t)> _onFlush;

bool _badChunk;
bool _good;
size_t _totalCount = 0;
Expand Down
15 changes: 13 additions & 2 deletions src/lib/fountain/FountainMetadata.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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:
Expand Down Expand Up @@ -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();
Expand Down

0 comments on commit 47a221d

Please sign in to comment.