image_pipeline.cpp

#include "image_pipeline.h"
#include "config.h"
#include "palette.h"

#include <TJpg_Decoder.h>
#include <esp_heap_caps.h>
#include <math.h>

namespace {

// --- Decoded-image state shared with the (capture-less) TJpgDec callback ---
uint8_t* g_dec = nullptr;   // RGB888, g_dw x g_dh
int g_dw = 0, g_dh = 0;

void* psAlloc(size_t n) {
  void* p = heap_caps_malloc(n, MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
  return p ? p : malloc(n);
}

// Ordered-dither threshold map (0..63).
const uint8_t kBayer8[64] = {
   0, 48, 12, 60,  3, 51, 15, 63,
  32, 16, 44, 28, 35, 19, 47, 31,
   8, 56,  4, 52, 11, 59,  7, 55,
  40, 24, 36, 20, 43, 27, 39, 23,
   2, 50, 14, 62,  1, 49, 13, 61,
  34, 18, 46, 30, 33, 17, 45, 29,
  10, 58,  6, 54,  9, 57,  5, 53,
  42, 26, 38, 22, 41, 25, 37, 21,
};

// TJpgDec streams 16x16 (or scaled) blocks of RGB565; expand into g_dec.
bool onBlock(int16_t x, int16_t y, uint16_t w, uint16_t h, uint16_t* bmp) {
  for (int j = 0; j < h; ++j) {
    int py = y + j;
    if (py < 0 || py >= g_dh) continue;
    for (int i = 0; i < w; ++i) {
      int px = x + i;
      if (px < 0 || px >= g_dw) continue;
      uint16_t c = bmp[j * w + i];
      uint8_t r = (c >> 11) & 0x1f, g = (c >> 5) & 0x3f, b = c & 0x1f;
      uint8_t* d = g_dec + ((size_t)py * g_dw + px) * 3;
      d[0] = (r << 3) | (r >> 2);
      d[1] = (g << 2) | (g >> 4);
      d[2] = (b << 3) | (b >> 2);
    }
  }
  return true;
}

inline void sampleBilinear(float fx, float fy, uint8_t* out) {
  int x0 = (int)floorf(fx), y0 = (int)floorf(fy);
  float ax = fx - x0, ay = fy - y0;
  int x1 = x0 + 1 < g_dw ? x0 + 1 : g_dw - 1;
  int y1 = y0 + 1 < g_dh ? y0 + 1 : g_dh - 1;
  if (x0 < 0) x0 = 0;
  if (y0 < 0) y0 = 0;
  const uint8_t* p00 = g_dec + ((size_t)y0 * g_dw + x0) * 3;
  const uint8_t* p10 = g_dec + ((size_t)y0 * g_dw + x1) * 3;
  const uint8_t* p01 = g_dec + ((size_t)y1 * g_dw + x0) * 3;
  const uint8_t* p11 = g_dec + ((size_t)y1 * g_dw + x1) * 3;
  for (int c = 0; c < 3; ++c) {
    float top = p00[c] * (1 - ax) + p10[c] * ax;
    float bot = p01[c] * (1 - ax) + p11[c] * ax;
    out[c] = (uint8_t)(top * (1 - ay) + bot * ay + 0.5f);
  }
}

}  // namespace

bool imgpipe::decode(const uint8_t* jpg, size_t len) {
  uint16_t w = 0, h = 0;
  TJpgDec.setJpgScale(1);
  TJpgDec.setSwapBytes(false);
  TJpgDec.setCallback(onBlock);
  if (TJpgDec.getJpgSize(&w, &h, jpg, len) != 0 || w == 0 || h == 0) { Serial.println("[pipe] getJpgSize FAIL (not a baseline JPEG?)"); return false; }

  int scale = 1;
  while (scale < 8 && (max(w, h) / scale) > MAX_DECODE_LONG_SIDE) scale <<= 1;
  // Allocate the decoded-RGB buffer; if PSRAM is tight, drop to a coarser scale.
  for (;;) {
    g_dw = (w + scale - 1) / scale;
    g_dh = (h + scale - 1) / scale;
    g_dec = (uint8_t*)psAlloc((size_t)g_dw * g_dh * 3);
    if (g_dec) break;
    if (scale >= 8) { Serial.println("[pipe] g_dec alloc FAIL"); return false; }
    scale <<= 1;
  }
  TJpgDec.setJpgScale(scale);
  memset(g_dec, 255, (size_t)g_dw * g_dh * 3);
  if (TJpgDec.drawJpg(0, 0, jpg, len) != 0) { Serial.println("[pipe] drawJpg FAIL"); free(g_dec); g_dec = nullptr; return false; }
  return true;
}

void imgpipe::freeDecode() { if (g_dec) { free(g_dec); g_dec = nullptr; } }

bool imgpipe::render(uint8_t* out, int fit, int dither) {
  if (!g_dec) return false;
  // Fit params (COVER = fill+crop, CONTAIN = letterbox).
  float sCover = fmaxf((float)PANEL_W / g_dw, (float)PANEL_H / g_dh);
  float sContain = fminf((float)PANEL_W / g_dw, (float)PANEL_H / g_dh);
  float s = (fit == 0) ? sCover : sContain;
  float ox = (PANEL_W - g_dw * s) * 0.5f;
  float oy = (PANEL_H - g_dh * s) * 0.5f;

  const int W = PANEL_W;
  // Resize + dither one output row at a time, sampling the decoded image on the
  // fly. Avoids a full-frame PANEL_W*PANEL_H*3 RGB buffer (5.76 MB at 1200x1600)
  // that would blow the 8 MB PSRAM budget alongside g_dec and the index frame.
  uint8_t* row = (uint8_t*)malloc((size_t)W * 3);
  if (!row) { free(g_dec); g_dec = nullptr; return false; }
  auto fillRow = [&](int dy) {
    for (int dx = 0; dx < W; ++dx) {
      uint8_t* d = row + (size_t)dx * 3;
      float fx = (dx - ox) / s, fy = (dy - oy) / s;
      if (fx >= 0 && fx <= g_dw - 1 && fy >= 0 && fy <= g_dh - 1) sampleBilinear(fx, fy, d);
      else { d[0] = d[1] = d[2] = 255; }  // outside source -> white border
    }
  };

  if (dither == 0) {
    int16_t* curErr = (int16_t*)psAlloc(sizeof(int16_t) * W * 3);
    int16_t* nextErr = (int16_t*)psAlloc(sizeof(int16_t) * W * 3);
    if (!curErr || !nextErr) { free(row); free(g_dec); g_dec = nullptr; free(curErr); free(nextErr); return false; }
    memset(curErr, 0, sizeof(int16_t) * W * 3);
    for (int y = 0; y < PANEL_H; ++y) {
      fillRow(y);
      memset(nextErr, 0, sizeof(int16_t) * W * 3);
      for (int x = 0; x < W; ++x) {
        uint8_t* px = row + (size_t)x * 3;
        int v[3];
        for (int c = 0; c < 3; ++c) {
          int t = px[c] + curErr[x * 3 + c];
          v[c] = t < 0 ? 0 : (t > 255 ? 255 : t);
        }
        uint8_t idx = e6_nearest(v[0], v[1], v[2]);
        out[(size_t)y * W + x] = idx;
        for (int c = 0; c < 3; ++c) {
          int err = v[c] - (c == 0 ? kE6[idx].r : c == 1 ? kE6[idx].g : kE6[idx].b);
          if (x + 1 < W) curErr[(x + 1) * 3 + c] += err * 7 / 16;
          if (x > 0)     nextErr[(x - 1) * 3 + c] += err * 3 / 16;
          nextErr[x * 3 + c] += err * 5 / 16;
          if (x + 1 < W) nextErr[(x + 1) * 3 + c] += err * 1 / 16;
        }
      }
      int16_t* tmp = curErr; curErr = nextErr; nextErr = tmp;
    }
    free(curErr);
    free(nextErr);
  } else if (dither == 1) {
    for (int y = 0; y < PANEL_H; ++y) {
      fillRow(y);
      for (int x = 0; x < W; ++x) {
        uint8_t* px = row + (size_t)x * 3;
        int bias = (kBayer8[(y & 7) * 8 + (x & 7)] - 32) * 2;
        int v[3];
        for (int c = 0; c < 3; ++c) {
          int t = px[c] + bias;
          v[c] = t < 0 ? 0 : (t > 255 ? 255 : t);
        }
        out[(size_t)y * W + x] = e6_nearest(v[0], v[1], v[2]);
      }
    }
  } else {
    for (int y = 0; y < PANEL_H; ++y) {
      fillRow(y);
      for (int x = 0; x < W; ++x) {
        uint8_t* px = row + (size_t)x * 3;
        out[(size_t)y * W + x] = e6_nearest(px[0], px[1], px[2]);
      }
    }
  }

  free(row);
  free(g_dec);
  g_dec = nullptr;
  return true;
}