CodecJPEG.h

Go to the documentation of this file.

#pragma once
 
#include <algorithm>  // std::min
 
#include "AudioTools/CoreAudio/AudioBasic/Collections/Vector.h"
#include "AudioTools/Video/CodecVideo.h"
#include "TinyJPEGDecoder.h"  // https://github.com/pschatzmann/TinyJPEG
#include "TinyJPEGEncoder.h"

 
namespace audio_tools {

class MJPEGDecoder : public VideoDecoder, public VideoInfoSource {
 public:
  MJPEGDecoder() {
    decoder_.setCallback(onBlock);
    decoder_.setUserData(this);
    decoder_.setSwapBytes(byte_swap);
  }
  MJPEGDecoder(Print &out) : MJPEGDecoder() { setOutput(out); }
  MJPEGDecoder(VideoOutput &out) : MJPEGDecoder() { setOutput(out); }
 
  VideoFormat codecFormat() override { return VideoFormat::MJPEG; }

  bool isValid(const uint8_t *data, size_t len) override {
    return len >= 2 && data[0] == 0xFF && data[1] == 0xD8;
  }

  void setOutput(Print &out) override { p_out = &out; }
  void setOutput(VideoOutput &out) override { p_out_video = &out; }

  void setVideoFormat(VideoFormat format) override {
    if (format != VideoFormat::RGB565) {
      LOGW("MJPEGDecoder: unsupported VideoFormat %d - only RGB565",
           (int)format);
      return;
    }
  }

  void setByteSwap(bool active) {
    byte_swap = active;
    decoder_.setSwapBytes(active);
  }
  bool byteSwap() const { return byte_swap; }

  VideoInfo videoInfo() override {
    VideoInfo info;
    info.format = VideoFormat::RGB565;
    info.width = width_;
    info.height = height_;
    return info;
  }
 
  bool begin() override { return true; }
  void end() override {}

  size_t write(const uint8_t *data, size_t len) override {
    if (pos == 0) start_ms = millis();
    // prevent memory fragmentation, change size only if more memory is needed
    if (img_vector.size() < pos + len) {
      img_vector.resize(pos + len);
    }
    memcpy(&img_vector[pos], data, len);
    pos += len;
    return len;
  }

  void flush() override {
    if (pos == 0) return;
    if (p_out == nullptr && p_out_video == nullptr) {
      pos = 0;
      return;
    }
    uint16_t w = 0, h = 0;
    if (decoder_.getJpgSize(&w, &h, img_vector.data(), pos) != JDR_OK ||
        w == 0 || h == 0) {
      LOGE("MJPEGDecoder: could not determine JPEG size");
      pos = 0;
      return;
    }
    width_ = w;
    height_ = h;
    size_t needed = (size_t)width_ * height_ * 2;
    if (frame_buffer.size() < needed) {
      frame_buffer.resize(needed);
      if (frame_buffer.data() == nullptr) {
        LOGE("MJPEGDecoder: frame buffer allocation failed");
        pos = 0;
        return;
      }
    }
    if (decoder_.drawJpg(0, 0, img_vector.data(), pos) == JDR_OK) {
      writeToOutput(frame_buffer.data(), needed);
    }
    pos = 0;
  }

  bool isKeyFrame(const uint8_t *data, size_t len) override { return true; }
 
 protected:
  Vector<uint8_t> img_vector;
  size_t pos = 0;
  uint64_t start_ms = 0;
  tinyjpeg::TinyJPEGDecoder decoder_;
  Print *p_out = nullptr;
  VideoOutput *p_out_video = nullptr;
  uint16_t width_ = 0;
  uint16_t height_ = 0;
  bool byte_swap = true;  // see setByteSwap()
  Vector<uint8_t> frame_buffer;  // tightly packed RGB565, 2 bytes/pixel

  static bool onBlock(tinyjpeg::TinyJPEGDecoder &decoder, int16_t x,
                       int16_t y, uint16_t w, uint16_t h, uint16_t *data) {
    return static_cast<MJPEGDecoder *>(decoder.getUserData())
        ->writeBlock(x, y, w, h, data);
  }

  bool writeBlock(int16_t x, int16_t y, uint16_t w, uint16_t h,
                   uint16_t *data) {
    if (x < 0 || y < 0) return true;
    uint32_t max_x = width_;
    uint32_t max_y = height_;
    if ((uint32_t)x >= max_x || (uint32_t)y >= max_y) return true;
    uint32_t win_w = std::min((uint32_t)w, max_x - (uint32_t)x);
    uint32_t win_h = std::min((uint32_t)h, max_y - (uint32_t)y);
    uint16_t *dst = (uint16_t *)frame_buffer.data();
    for (uint32_t row = 0; row < win_h; row++) {
      memcpy(dst + ((uint32_t)y + row) * max_x + (uint32_t)x,
             data + row * w, win_w * sizeof(uint16_t));
    }
    return true;
  }
 
  size_t writeToOutput(const uint8_t *data, size_t len) {
    size_t result = 0;
    if (p_out != nullptr) result = p_out->write(data, len);
    if (p_out_video != nullptr) result += p_out_video->write(data, len);
    return result;
  }
};

class MJPEGEncoder : public VideoEncoder {
 public:
  static constexpr size_t kDefaultBufferSize = 16384;
 
  MJPEGEncoder(int width = 0, int height = 0) { setSize(width, height); }
  MJPEGEncoder(Print &out) { setOutput(out); }
  MJPEGEncoder(VideoOutput &out) { setOutput(out); }

  void setOutput(Print &out) override { p_out = &out; }
  void setOutput(VideoOutput &out) { p_out_video = &out; }

  void setVideoFormat(VideoFormat format) override {
    switch (format) {
      case VideoFormat::RGB565:
      case VideoFormat::RGB888:
      case VideoFormat::RGB666:
        input_format = format;
        break;
      default:
        LOGW("MJPEGEncoder: unsupported VideoFormat %d", (int)format);
        break;
    }
  }

  VideoInfo videoInfo() override {
    VideoInfo info;
    info.format = VideoFormat::MJPEG;
    info.width = (uint16_t)width_;
    info.height = (uint16_t)height_;
    return info;
  }

  void setSize(int width, int height) {
    width_ = width;
    height_ = height;
  }

  void setQuality(uint8_t quality) { encoder_.setQuality(quality); }
  void setSubsample(bool enable) { encoder_.setSubsample(enable); }

  void setBufferSize(size_t size) {
    if (size > jpeg.size()) jpeg.resize(size);
  }
 
  bool begin() override {
    if (width_ <= 0 || height_ <= 0) {
      LOGE("MJPEGEncoder: setSize() was not called");
      return false;
    }
    if (jpeg.size() == 0) jpeg.resize(kDefaultBufferSize);
    return jpeg.data() != nullptr;
  }
 
  void end() override {}

  size_t write(const uint8_t *data, size_t len) override {
    if (p_out == nullptr && p_out_video == nullptr) return 0;
    size_t expected = videoFrameSizeBytes(input_format, width_, height_);
    if (expected == 0 || len < expected) {
      LOGE("MJPEGEncoder: %u bytes, expected %u", (unsigned)len,
           (unsigned)expected);
      return 0;
    }
    if (jpeg.size() == 0) jpeg.resize(kDefaultBufferSize);
    size_t size = 0;
    JERESULT rc = encode(data, size);
    // the image didn't fit: grow the buffer and try again
    for (int attempt = 0; rc == JER_INTR && attempt < kMaxGrowRetries;
         ++attempt) {
      jpeg.resize(jpeg.size() * 2);
      if (jpeg.data() == nullptr) {
        LOGE("MJPEGEncoder: buffer allocation failed");
        return 0;
      }
      rc = encode(data, size);
    }
    if (rc != JER_OK) {
      LOGE("MJPEGEncoder: encoding failed (%d)", (int)rc);
      return 0;
    }
    if (p_out != nullptr) p_out->write(jpeg.data(), size);
    if (p_out_video != nullptr) p_out_video->write(jpeg.data(), size);
    return size;
  }

  tinyjpeg::TinyJPEGEncoder &driver() { return encoder_; }
 
 protected:
  static constexpr int kMaxGrowRetries = 4;
 
  tinyjpeg::TinyJPEGEncoder encoder_;
  Print *p_out = nullptr;
  VideoOutput *p_out_video = nullptr;
  Vector<uint8_t> jpeg;
  VideoFormat input_format = VideoFormat::RGB565;
  int width_ = 0;
  int height_ = 0;
 
  JERESULT encode(const uint8_t *data, size_t &size) {
    JEPixelFormat format = input_format == VideoFormat::RGB565 ? JE_FMT_RGB565
                           : input_format == VideoFormat::RGB666
                               ? JE_FMT_RGB666
                               : JE_FMT_RGB888;
    return encoder_.encodeJpg(data, width_, height_, format, jpeg.data(),
                              jpeg.size(), size);
  }
};
 
}  // namespace audio_tools