Bayer.hpp

#pragma once
#include <Gfx/Graph/decoders/CaptureAdjustGLSL.hpp>
#include <Gfx/Graph/decoders/GPUVideoDecoder.hpp>
 
#include <QString>
 
extern "C" {
#include <libavformat/avformat.h>
}
 
namespace score::gfx
{
struct BayerDecoder : GPUVideoDecoder
{
  enum class Phase
  {
    RGGB, 
    GRBG, 
    GBRG, 
    BGGR, 
  };
 
  // %1 = user filter, %2/%3 = red-site offset, %4 = sample scale
#define SCORE_GFX_BAYER_DEMOSAIC R"_(
layout(binding=3) uniform sampler2D u_tex;
 
layout(location = 0) in vec2 v_texcoord;
layout(location = 0) out vec4 fragColor;
 
vec4 processTexture(vec4 tex) {
  vec4 processed = tex;
  { %1 }
  return processed;
}
 
float smp(ivec2 q, ivec2 lim) {
  return texelFetch(u_tex, clamp(q, ivec2(0), lim), 0).r;
}
 
void main() {
  ivec2 sz = textureSize(u_tex, 0);
  ivec2 lim = sz - ivec2(1);
  ivec2 ip = clamp(ivec2(floor(v_texcoord * vec2(sz))), ivec2(0), lim);
 
  // (0,0) marks the red site; the other three cells follow from its parity.
  ivec2 c = (ip + ivec2(%2, %3)) & 1;
 
  float ctr = smp(ip, lim);
  float n  = smp(ip + ivec2( 0,-1), lim);
  float s  = smp(ip + ivec2( 0, 1), lim);
  float w  = smp(ip + ivec2(-1, 0), lim);
  float e  = smp(ip + ivec2( 1, 0), lim);
  float nw = smp(ip + ivec2(-1,-1), lim);
  float ne = smp(ip + ivec2( 1,-1), lim);
  float sw = smp(ip + ivec2(-1, 1), lim);
  float se = smp(ip + ivec2( 1, 1), lim);
 
  float cross4 = (n + s + w + e) * 0.25;
  float diag4  = (nw + ne + sw + se) * 0.25;
  float horz2  = (w + e) * 0.5;
  float vert2  = (n + s) * 0.5;
 
  vec3 rgb;
  if(c.x == 0 && c.y == 0)
    rgb = vec3(ctr, cross4, diag4);          // red site
  else if(c.x == 1 && c.y == 1)
    rgb = vec3(diag4, cross4, ctr);          // blue site
  else if(c.x == 1 && c.y == 0)
    rgb = vec3(horz2, ctr, vert2);           // green, red neighbours across
  else
    rgb = vec3(vert2, ctr, horz2);           // green, blue neighbours across
 
  fragColor = processTexture(vec4(adjustCapture(rgb * %4), 1.0));
}
)_"
 
  static const constexpr auto frag = R"_(#version 450
 
)_" SCORE_GFX_CAPTURE_UNIFORMS SCORE_GFX_CAPTURE_ADJUST_FN SCORE_GFX_BAYER_DEMOSAIC;
 
  static const constexpr auto fileFrag = R"_(#version 450
 
)_" SCORE_GFX_VIDEO_UNIFORMS R"_(
vec3 adjustCapture(vec3 c) { return c; }
)_" SCORE_GFX_BAYER_DEMOSAIC;
#undef SCORE_GFX_BAYER_DEMOSAIC
 
  enum class Source
  {
    Capture,
    File
  };
 
  BayerDecoder(
      QRhiTexture::Format fmt, int bytesPerSample, Video::ImageFormat& d,
      Phase phase, double sampleScale = 1.0, Source source = Source::Capture,
      QString f = "")
      : format{fmt}
      , bytesPerSample{bytesPerSample}
      , decoder{d}
      , phase{phase}
      , sampleScale{sampleScale}
      , source{source}
      , filter{std::move(f)}
  {
  }
 
  QRhiTexture::Format format;
  int bytesPerSample{};
  Video::ImageFormat& decoder;
  Phase phase{};
  double sampleScale{1.0};
  Source source{Source::Capture};
  QString filter;
 
  std::pair<QShader, QShader> init(RenderList& r) override
  {
    auto& rhi = *r.state.rhi;
    const auto w = decoder.width, h = decoder.height;
 
    {
      const bool supported = rhi.isTextureFormatSupported(format);
      auto tex = rhi.newTexture(format, QSize{w, h}, 1, QRhiTexture::Flag{});
      const bool created = tex->create();
      if(!supported || !created)
        qWarning() << "BayerDecoder: the backend will not give us a" << int(format)
                   << "texture at" << w << "x" << h
                   << "(supported:" << supported << "created:" << created
                   << ") -- the mosaic has nowhere to land";
      else
        qDebug() << "BayerDecoder:" << w << "x" << h << "format" << int(format)
                 << "ok";
 
      // The mosaic is reconstructed per texel: linear filtering would blend
      // neighbouring colour sites together before the demosaic can separate
      // them, which is the same hazard the byte-reassembling decoders have.
      auto sampler = rhi.newSampler(
          QRhiSampler::Nearest, QRhiSampler::Nearest, QRhiSampler::None,
          QRhiSampler::ClampToEdge, QRhiSampler::ClampToEdge);
      sampler->create();
 
      samplers.push_back({sampler, tex});
    }
 
    int px = 0, py = 0;
    switch(phase)
    {
      case Phase::RGGB: px = 0; py = 0; break;
      case Phase::GRBG: px = 1; py = 0; break;
      case Phase::GBRG: px = 0; py = 1; break;
      case Phase::BGGR: px = 1; py = 1; break;
    }
 
    const bool file = source == Source::File;
    return score::gfx::makeShaders(
        r.state, file ? vertexShader() : QString(score::gfx::captureVertexShader),
        QString(file ? fileFrag : frag)
            .arg(filter)
            .arg(px)
            .arg(py)
            .arg(QString::number(sampleScale, 'f', 6)));
  }
 
  void exec(RenderList&, QRhiResourceUpdateBatch& res, AVFrame& frame) override
  {
    auto tex = samplers[0].texture;
    QRhiTextureUploadEntry entry{
        0, 0,
        createTextureUpload(
            frame.data[0], decoder.width, decoder.height, bytesPerSample,
            frame.linesize[0])};
    QRhiTextureUploadDescription desc{entry};
    res.uploadTexture(tex, desc);
  }
};
 
} // namespace score::gfx