AntiClipSettings/DataCollection.cpp

#include "DataCollection.h"
#include "BoostLog.h"
#include <QFile>
#include <QNetworkAccessManager>
#include <QNetworkReply>
#include <QTimer>
#include <boost/json/object.hpp>
#include <boost/json/parse.hpp>
#include <boost/json/serialize.hpp>
#include <turbojpeg.h>

DataCollection::DataCollection(QObject *parent) : QObject{parent} {
    m_manager = new QNetworkAccessManager(this);
}

void DataCollection::start(const QString &address) {
    m_address = address;
    m_enabled = true;
    emit enabledChanged();
    start();
}

void DataCollection::stop() {
    m_enabled = false;
    emit enabledChanged();
}

void DataCollection::setImageFormat(ImageFormat format, int quality) {
    if (m_format != format) {
        m_format = format;
    }
    if (m_quality != quality) {
        m_quality = quality;
    }
}

QString DataCollection::path() const {
    return m_path;
}

void DataCollection::setPath(const QString &path) {
    // file:///E:/Downloads/logs
    auto p = path;
    if (p.startsWith("file:///")) {
        p.remove(0, 8);
    }
    if (m_path != p) {
        m_path = p;
        emit pathChanged();
        LOG(info) << "set data path: " << m_path.toStdString();
    }
}

bool DataCollection::enabled() const {
    return m_enabled;
}

void DataCollection::start() {
    auto url = QString("http://%1:8080/capture.do").arg(m_address);
    QNetworkRequest request;
    request.setHeader(QNetworkRequest::ContentTypeHeader, "application/json");
    request.setUrl(url);

    boost::json::object data;
    data["devid"] = 1;
    auto content = boost::json::serialize(data);
    QNetworkReply *reply = m_manager->post(request, QString::fromStdString(content).toLocal8Bit());
    connect(reply, &QNetworkReply::finished, this, &DataCollection::onCaptureFinished);
}

void DataCollection::onCaptureFinished() {
    auto reply = dynamic_cast<QNetworkReply *>(sender());
    reply->deleteLater();

    auto text = reply->readAll().toStdString();

    auto replyVale = boost::json::parse(text);
    auto root = replyVale.as_object();

    m_filename = QString::fromStdString(std::string(root.at("path").as_string()));
    auto url = QString("http://%1:8080/%2").arg(m_address, m_filename);
    QNetworkRequest request;
    request.setUrl(url);
    QNetworkReply *dataReply = m_manager->get(request);
    connect(dataReply, &QNetworkReply::finished, this, &DataCollection::onDataGetFinished);
}

void DataCollection::onDataGetFinished() {
    auto reply = dynamic_cast<QNetworkReply *>(sender());
    auto data = reply->readAll();
    if (m_format == ImageFormat::Jpeg) {
        m_filename.replace("yuv", "jpg");
    }
    LOG(info) << "filename: " << m_filename.toStdString() << ", error: " << reply->error()
              << ", capture data size: " << data.size();
    QFile file(m_path + "/" + m_filename);
    file.open(QIODevice::WriteOnly);
    if (m_format == ImageFormat::Jpeg) {
        auto jpegBuffer =
            encodeNv21ToJpeg(reinterpret_cast<uint8_t *>(data.data()), ImageWidth, ImageHeight, m_quality);
        file.write(reinterpret_cast<const char *>(jpegBuffer.data()), jpegBuffer.size());
    } else {
        file.write(data);
    }
    file.close();
    if (m_enabled) {
        QTimer::singleShot(0, this, [this]() { start(); });
    }
}

static void convertNV21ToYUV420(uint8_t *nv21, uint8_t *yuv420, int width, int height) {
    int frameSize = width * height;
    int chromaSize = frameSize / 4;

    memcpy(yuv420, nv21, frameSize);

    uint8_t *u = yuv420 + frameSize;
    uint8_t *v = u + chromaSize;

    uint8_t *uv = nv21 + frameSize;

    for (int i = 0; i < chromaSize; i++) {
        u[i] = uv[2 * i + 1];
        v[i] = uv[2 * i];
    }
}

std::vector<uint8_t> DataCollection::encodeNv21ToJpeg(uint8_t *nv21, int width, int height, int quality) {
    tjhandle handle = tjInitCompress();
    if (handle == nullptr) {
        fprintf(stderr, "Failed to initialize turbojpeg compressor\n");
        return {};
    }

    unsigned long jpeg_size = 0;
    uint8_t *jpeg_buffer = nullptr;

    int yuv_size = tjBufSizeYUV2(width, 4, height, TJSAMP_420);
    uint8_t *yuv_buffer = (uint8_t *)malloc(yuv_size);
    if (yuv_buffer == nullptr) {
        fprintf(stderr, "Failed to allocate YUV buffer\n");
        tjDestroy(handle);
        return {};
    }

    // Convert NV21 to YUV420
    convertNV21ToYUV420(nv21, yuv_buffer, width, height);

    // Compress YUV to JPEG
    int retval = tjCompressFromYUV(handle, yuv_buffer, width, 4, height, TJSAMP_420, &jpeg_buffer, &jpeg_size, quality,
                                   TJFLAG_FASTDCT);
    if (retval != 0) {
        fprintf(stderr, "Failed to compress to JPEG: %s\n", tjGetErrorStr());
        free(yuv_buffer);
        tjDestroy(handle);
        return {};
    }

    // Copy the JPEG buffer into a std::vector
    std::vector<uint8_t> jpeg_vector(jpeg_buffer, jpeg_buffer + jpeg_size);

    // Free the JPEG buffer allocated by libjpeg-turbo
    tjFree(jpeg_buffer);
    free(yuv_buffer);
    tjDestroy(handle);

    return jpeg_vector;
}