JPEGRtpEncoder.h

Go to the documentation of this file.

/*
 * Author: Phil Schatzmann
 *
 * JPEG RTP Encoder for RTSP Video Streaming
 * Implements RFC 2435 compliant JPEG over RTP fragmentation
 */
 
#pragma once
 
#include "AudioTools/AudioCodecs/AudioCodecsBase.h"
#include "AudioTools/CoreAudio/AudioBasic/Collections/Vector.h"
#include "RTSPFormat.h"
#include "RTSPFragmentQueue.h"
#include "RTSPVideoEncoder.h"
 
namespace audio_tools {
 
class JPEGRtpEncoder : public RTSPVideoEncoder, protected RTSPFragmentQueue {
 public:
  JPEGRtpEncoder() = default;
 
  JPEGRtpEncoder(size_t maxFragmentSize) : m_maxFragmentSize(maxFragmentSize) {}
 
  void setMaxFragmentSize(size_t size) { m_maxFragmentSize = size; }
 
  void setMaxQueuedFrames(size_t frames) { m_maxQueuedFrames = frames; }
 
  void setFormat(RTSPFormat &format) override { p_format = &format; }
 
  // -- IMediaSource --------------------------------------------------
 
  RTSPFormat &getFormat() override {
    return p_format != nullptr ? *p_format : default_format;
  }
 
  void start() override { begin(); }
  void stop() override { end(); }
 
  int packetSize() override { return queuePacketSize(); }
 
  int readBytes(void *dest, int maxBytes) override {
    return queueReadBytes(dest, maxBytes);
  }
 
  // -- AudioEncoder ----------------------------------------------------
 
  bool begin() override {
    LOGI("Starting JPEG RTP Encoder");
    m_isStarted = true;
    m_currentFrame.clear();
    m_frameComplete = true;
    clearQueue();
    return true;
  }
 
  void end() override {
    LOGI("Stopping JPEG RTP Encoder");
    m_isStarted = false;
    m_currentFrame.clear();
    clearQueue();
  }
 
  void setAudioInfo(AudioInfo info) override {
    // Video encoder: dimensions/framerate come from the RTSPFormat, not
    // AudioInfo
    (void)info;
  }
 
  const char *mime() override { return "video/jpeg"; }
 
  size_t write(const uint8_t *data, size_t len) override {
    if (!m_isStarted || len == 0) {
      return 0;
    }
 
    if (m_frameComplete) {
      m_currentFrame.clear();
      m_frameComplete = false;
    }
 
    size_t oldSize = m_currentFrame.size();
    m_currentFrame.resize(oldSize + len);
    memcpy(m_currentFrame.data() + oldSize, data, len);
 
    if (isCompleteJpegFrame()) {
      processCompleteFrame();
      m_frameComplete = true;
    }
 
    return len;
  }
 
  operator bool() override { return m_isStarted; }
 
 protected:
  RTSPFormat *p_format = nullptr;
  RTSPFormatMJPEG default_format;  // fallback so getFormat() is never null
 
  bool m_isStarted = false;
  size_t m_maxFragmentSize = 1400;  // Typical MTU minus IP/UDP/RTP headers
 
  // Frame accumulation state
  Vector<uint8_t> m_currentFrame;
  bool m_frameComplete = true;
 
  // Reused scratch buffer for building one fragment at a time
  Vector<uint8_t> m_scratch;
 
  bool isCompleteJpegFrame() {
    size_t size = m_currentFrame.size();
    if (size < 4) return false;
    return (m_currentFrame[size - 2] == 0xFF && m_currentFrame[size - 1] == 0xD9);
  }
 
  size_t findJpegDataStart() {
    size_t size = m_currentFrame.size();
    if (size < 4) return 0;
    for (size_t i = 0; i + 4 <= size; i++) {
      if (m_currentFrame[i] == 0xFF && m_currentFrame[i + 1] == 0xDA) {
        size_t sosLength = (m_currentFrame[i + 2] << 8) | m_currentFrame[i + 3];
        size_t dataStart = i + 2 + sosLength;
        return dataStart <= size ? dataStart : 0;
      }
    }
    return 0;
  }
 
  int extractQuantTables(const uint8_t *data, size_t headerLen, uint8_t *out) {
    uint8_t table0[64];
    uint8_t table1[64];
    bool has0 = false, has1 = false;
 
    size_t i = 2;  // skip SOI (0xFFD8)
    while (i + 4 <= headerLen) {
      if (data[i] != 0xFF) {
        ++i;
        continue;
      }
      uint8_t marker = data[i + 1];
      if (marker == 0xD8 || marker == 0x01 || (marker >= 0xD0 && marker <= 0xD7)) {
        i += 2;
        continue;
      }
      size_t segLen = (data[i + 2] << 8) | data[i + 3];
      if (segLen < 2 || i + 2 + segLen > headerLen) break;
 
      if (marker == 0xDB) {  // DQT
        size_t p = i + 4;
        size_t segEnd = i + 2 + segLen;
        while (p < segEnd) {
          uint8_t pqtq = data[p++];
          uint8_t precision = pqtq >> 4;
          uint8_t id = pqtq & 0x0F;
          size_t tableBytes = precision == 0 ? 64 : 128;
          if (p + tableBytes > segEnd) break;
          if (precision == 0) {
            if (id == 0 && !has0) {
              memcpy(table0, data + p, 64);
              has0 = true;
            } else if (id == 1 && !has1) {
              memcpy(table1, data + p, 64);
              has1 = true;
            }
          }
          p += tableBytes;
        }
      } else if (marker == 0xDA) {
        break;  // reached scan data, no more tables before it
      }
      i += 2 + segLen;
    }
 
    if (!has0) return 0;
    memcpy(out, table0, 64);
    memcpy(out + 64, has1 ? table1 : table0, 64);
    return 128;
  }
 
  uint8_t detectJpegType(const uint8_t *data, size_t headerLen) {
    size_t i = 2;
    while (i + 4 <= headerLen) {
      if (data[i] != 0xFF) {
        ++i;
        continue;
      }
      uint8_t marker = data[i + 1];
      if (marker == 0xD8 || marker == 0x01 || (marker >= 0xD0 && marker <= 0xD7)) {
        i += 2;
        continue;
      }
      size_t segLen = (data[i + 2] << 8) | data[i + 3];
      if (segLen < 2 || i + 2 + segLen > headerLen) break;
 
      if (marker == 0xC0 || marker == 0xC1) {  // SOF0 / SOF1 (baseline)
        size_t p = i + 4;
        size_t segEnd = i + 2 + segLen;
        // layout from p: precision(1) Y(2) X(2) Nf(1) [Cid(1) HV(1) Tq(1)]...
        if (p + 8 <= segEnd) {
          uint8_t hv = data[p + 7];
          uint8_t h = hv >> 4;
          uint8_t v = hv & 0x0F;
          if (h == 2 && v == 1) return 0;  // 4:2:2
          if (h == 2 && v == 2) return 1;  // 4:2:0
          LOGW(
              "JPEGRtpEncoder: unsupported chroma subsampling H=%d V=%d; "
              "assuming 4:2:2",
              h, v);
        }
        return 0;
      }
      if (marker == 0xDA) break;
      i += 2 + segLen;
    }
    LOGW("JPEGRtpEncoder: SOF0 marker not found; assuming 4:2:2 subsampling");
    return 0;
  }
 
  void buildJpegRtpHeader(uint8_t *buffer, size_t fragmentOffset, int width,
                          int height, uint8_t type, uint8_t qValue) {
    buffer[0] = 0;  // Type-specific field (0 for baseline JPEG)
    buffer[1] = (fragmentOffset >> 16) & 0xFF;
    buffer[2] = (fragmentOffset >> 8) & 0xFF;
    buffer[3] = fragmentOffset & 0xFF;
    buffer[4] = type;
    buffer[5] = qValue;
    buffer[6] = (uint8_t)((width / 8) & 0xFF);
    buffer[7] = (uint8_t)((height / 8) & 0xFF);
  }
 
  void writeQuantTableHeader(uint8_t *buffer, const uint8_t *qTable, int qTableLen) {
    buffer[0] = 0;  // MBZ
    buffer[1] = 0;  // Precision: 0 = 8-bit for all tables
    buffer[2] = (qTableLen >> 8) & 0xFF;
    buffer[3] = qTableLen & 0xFF;
    memcpy(buffer + 4, qTable, qTableLen);
  }
 
  void processCompleteFrame() {
    if (p_format == nullptr) {
      LOGE("JPEGRtpEncoder: no RTSPFormat configured (call setFormat()); dropping frame");
      return;
    }
 
    size_t headerSize = findJpegDataStart();
    if (headerSize == 0 || headerSize >= m_currentFrame.size()) {
      LOGW("JPEGRtpEncoder: could not locate JPEG scan data (SOS marker); dropping frame");
      return;
    }
 
    uint8_t qTable[128];
    int qTableLen = extractQuantTables(m_currentFrame.data(), headerSize, qTable);
    uint8_t qValue = qTableLen > 0 ? 255 : 127;
    if (qTableLen == 0) {
      LOGW(
          "JPEGRtpEncoder: could not extract quantization tables from frame; "
          "using approximate standard tables (Q=%d)",
          qValue);
    }
    uint8_t type = detectJpegType(m_currentFrame.data(), headerSize);
 
    RTSPVideoInfo vi = p_format->videoInfo();
    int width = vi.width;
    int height = vi.height;
 
    const uint8_t *frameData = m_currentFrame.data() + headerSize;
    size_t frameSize = m_currentFrame.size() - headerSize;
    size_t processedOffset = 0;
    int fragmentCount = 0;
 
    while (processedOffset < frameSize) {
      bool isFirst = (processedOffset == 0);
      size_t qLen = isFirst ? (size_t)qTableLen : 0;
      size_t qHeaderLen = qLen > 0 ? 4 + qLen : 0;
 
      if (m_maxFragmentSize <= 8 + qHeaderLen) {
        LOGE("JPEGRtpEncoder: maxFragmentSize too small to fit RTP/JPEG headers");
        return;
      }
      size_t maxData = m_maxFragmentSize - 8 - qHeaderLen;
      size_t chunk = frameSize - processedOffset;
      if (chunk > maxData) chunk = maxData;
      bool isLast = (processedOffset + chunk >= frameSize);
 
      m_scratch.resize(8 + qHeaderLen + chunk);
      buildJpegRtpHeader(m_scratch.data(), processedOffset, width, height, type, qValue);
      if (qHeaderLen > 0) {
        writeQuantTableHeader(m_scratch.data() + 8, qTable, qTableLen);
      }
      memcpy(m_scratch.data() + 8 + qHeaderLen, frameData + processedOffset, chunk);
 
      appendFragment(m_scratch.data(), m_scratch.size(), isLast);
 
      processedOffset += chunk;
      ++fragmentCount;
    }
 
    LOGI("JPEGRtpEncoder: queued %d RTP fragments for %u byte frame",
         fragmentCount, (unsigned)frameSize);
  }
};
 
}  // namespace audio_tools