src/processingworker.cpp - master - ka_pie/PeakOff

src/processingworker.cpp
422 строки13 KB
Viktoria
final version of PeakOff
14 июл 2026, 11:13
14 июл 2026, 11:13a838933
#include "processingworker.h"
#include <QDebug>
#include <QThread>
#include <QFile>
#include <vector>

ProcessingWorker::ProcessingWorker(const QString &u2netPath,
                                   const QString &enhancePath,
                                   const QString &stylePath,
                                   QObject *parent)
    : QObject(parent)
    , m_u2netPath(u2netPath)
    , m_enhancePath(enhancePath)
    , m_stylePath(stylePath)
{
    try {
        if (QFile::exists(m_u2netPath)) {
            m_u2netNet = cv::dnn::readNetFromONNX(m_u2netPath.toStdString());
            m_u2netLoaded = !m_u2netNet.empty();
            if (m_u2netLoaded) {
                qDebug() << "[ProcessingWorker] U2-Net model loaded:" << m_u2netPath;
            } else {
                qWarning() << "[ProcessingWorker] U2-Net model empty:" << m_u2netPath;
            }
        } else {
            qWarning() << "[ProcessingWorker] Model file not found:" << m_u2netPath;
        }
    } catch (const cv::Exception &e) {
        qWarning() << "[ProcessingWorker] Model load error:" << e.what();
        m_u2netLoaded = false;
    }

    try {
        if (QFile::exists(m_enhancePath)) {
            m_enhanceNet = cv::dnn::readNetFromONNX(m_enhancePath.toStdString());
            m_enhanceLoaded = !m_enhanceNet.empty();
            if (m_enhanceLoaded) {
                qDebug() << "[ProcessingWorker] super-resolution-10 model loaded:" << m_enhancePath;
            } else {
                qWarning() << "[ProcessingWorker] super-resolution-10 model empty:" << m_enhancePath;
            }
        } else {
            qWarning() << "[ProcessingWorker] Model file not found:" << m_enhancePath;
        }
    } catch (const cv::Exception &e) {
        qWarning() << "[ProcessingWorker] Model load error:" << e.what();
        m_enhanceLoaded = false;
    }

    try {
        if (QFile::exists(m_stylePath)) {
            m_styleNet = cv::dnn::readNetFromONNX(m_stylePath.toStdString());
            m_styleLoaded = !m_styleNet.empty();
            if (m_styleLoaded) {
                qDebug() << "[ProcessingWorker] rain-princess-9 model loaded:" << m_stylePath;
            } else {
                qWarning() << "[ProcessingWorker] rain-princess-9 model empty:" << m_stylePath;
            }
        } else {
            qWarning() << "[ProcessingWorker] Model file not found:" << m_stylePath;
        }
    } catch (const cv::Exception &e) {
        qWarning() << "[ProcessingWorker] Model load error:" << e.what();
        m_styleLoaded = false;
    }
}


void ProcessingWorker::processEnhance(const QImage &img)
{
    QElapsedTimer timer;
        timer.start();
        ProcessingMetrics metrics;
        emit progressChanged(10);

        if (img.isNull()) {
            emit errorOccurred("Empty image");
            return;
        }

        if (!m_enhanceLoaded) {
            qWarning() << "[ProcessingWorker] Enhance model not loaded, using stub";
            QImage stubResult = stubEnhance(img);
            emit progressChanged(100);
            metrics.totalMs = timer.elapsed();
            emit metricsReady(metrics);
            emit enhanceFinished(stubResult);
            return;
        }

        try {
            QElapsedTimer stepTimer;

            stepTimer.start();
            cv::Mat input = qimageToMat(img);
            metrics.preprocessMs = stepTimer.elapsed();
            emit progressChanged(30);

            cv::Mat output = runGenericInference(m_enhanceNet, input, cv::Size(224, 224), metrics);
            emit progressChanged(80);


            stepTimer.start();
            QImage result = matToQImage(output);
            metrics.postprocessMs += stepTimer.elapsed();

            emit progressChanged(100);
            metrics.totalMs = timer.elapsed();

            emit metricsReady(metrics);
            emit enhanceFinished(result);
        } catch (const cv::Exception &e) {
            emit errorOccurred(QString("OpenCV Enhance error: %1").arg(e.what()));
        }
}

void ProcessingWorker::processRemoveBackground(const QImage &img)
{
    QElapsedTimer timer;
    timer.start();
    emit progressChanged(10);

    if (img.isNull()) {
        emit errorOccurred("Empty image");
        return;
    }

    cv::Mat input = qimageToMat(img);
    if (input.empty()) {
        emit errorOccurred("Failed to convert image");
        return;
    }

    if (!m_u2netLoaded) {
        qWarning() << "[ProcessingWorker] Model not loaded, using stub";
        QImage stubResult = stubEnhance(img);
        emit progressChanged(100);
        ProcessingMetrics metrics;
        metrics.totalMs = timer.elapsed();
        emit metricsReady(metrics);
        emit removeBackgroundFinished(stubResult);
        return;
    }

    try {
        emit progressChanged(30);
        cv::Mat binaryMask = runInference(input);
        emit progressChanged(70);

        cv::Mat rgba;
        cv::cvtColor(input, rgba, cv::COLOR_RGB2RGBA);
        std::vector<cv::Mat> channels;
        cv::split(rgba, channels);
        channels[3] = binaryMask;
        cv::merge(channels, rgba);

        QImage result = matToQImage(rgba);
        emit progressChanged(100);

        ProcessingMetrics metrics;
        metrics.totalMs = timer.elapsed();
        emit metricsReady(metrics);
        emit removeBackgroundFinished(result);
    } catch (const cv::Exception &e) {
        emit errorOccurred(QString("OpenCV error: %1").arg(e.what()));
    }
}

void ProcessingWorker::processStyle(const QImage &img, int styleIndex)
{
     QElapsedTimer timer;
     timer.start();
     ProcessingMetrics metrics;
     emit progressChanged(10);

     if (img.isNull()) {
         emit errorOccurred("Empty image");
         return;
     }

     if (!m_styleLoaded) {
         qWarning() << "[ProcessingWorker] Style model not loaded, using stub";
         QImage stubResult = stubStyle(img, styleIndex);
         emit progressChanged(100);
         metrics.totalMs = timer.elapsed();
         emit metricsReady(metrics);
         emit styleFinished(stubResult, styleIndex);
         return;
     }

     try {
         QElapsedTimer stepTimer;

         stepTimer.start();
         cv::Mat input = qimageToMat(img);
         metrics.preprocessMs = stepTimer.elapsed();
         emit progressChanged(30);

         cv::Mat output = runStyleInference(m_styleNet, input, cv::Size(640, 640), metrics);
         emit progressChanged(80);

         stepTimer.start();
         QImage result = matToQImage(output);
         metrics.postprocessMs += stepTimer.elapsed();

         emit progressChanged(100);
         metrics.totalMs = timer.elapsed();

         emit metricsReady(metrics);
         emit styleFinished(result, styleIndex);
     } catch (const cv::Exception &e) {
         emit errorOccurred(QString("OpenCV Style error: %1").arg(e.what()));
     }
}

cv::Mat ProcessingWorker::runGenericInference(cv::dnn::Net &net, const cv::Mat &input, const cv::Size &inputSize, ProcessingMetrics &metrics)
{
    QElapsedTimer timer;

    timer.start();
    
    cv::Mat ycrcb;
    cv::cvtColor(input, ycrcb, cv::COLOR_RGB2YCrCb);
    
    std::vector<cv::Mat> channels;
    cv::split(ycrcb, channels);
    
    cv::Mat y_resized;
    cv::resize(channels[0], y_resized, inputSize, 0, 0, cv::INTER_CUBIC);
    
    cv::Mat blob = cv::dnn::blobFromImage(y_resized, 1.0 / 255.0, inputSize, cv::Scalar(0), false, false);
    metrics.preprocessMs += timer.elapsed();

    timer.start();
    net.setInput(blob);
    cv::Mat outputBlob = net.forward();
    metrics.inferenceMs = timer.elapsed();

    timer.start();
    std::vector<cv::Mat> images;
    cv::dnn::imagesFromBlob(outputBlob, images);
    cv::Mat outY = images[0];

    if (outY.depth() == CV_32F) {
        outY.convertTo(outY, CV_8U, 255.0);
    }

    cv::Mat outY_rescaled;
    cv::resize(outY, outY_rescaled, input.size(), 0, 0, cv::INTER_CUBIC);

    std::vector<cv::Mat> updatedChannels = { outY_rescaled, channels[1], channels[2] };
    cv::Mat finalYCrCb;
    cv::merge(updatedChannels, finalYCrCb);

    cv::Mat finalOutput;
    cv::cvtColor(finalYCrCb, finalOutput, cv::COLOR_YCrCb2BGR);

    metrics.postprocessMs = timer.elapsed();

    return finalOutput;
}


cv::Mat ProcessingWorker::runStyleInference(cv::dnn::Net &net, const cv::Mat &input, const cv::Size &inputSize, ProcessingMetrics &metrics)
{
    QElapsedTimer timer;

    timer.start();
    cv::Mat blob = cv::dnn::blobFromImage(input, 1.0, inputSize, cv::Scalar(0, 0, 0), true, false);
    metrics.preprocessMs += timer.elapsed();

    timer.start();
    net.setInput(blob);
    cv::Mat outputBlob = net.forward();
    metrics.inferenceMs = timer.elapsed();

    timer.start();
    std::vector<cv::Mat> images;
    cv::dnn::imagesFromBlob(outputBlob, images);
    cv::Mat outImg = images[0];

    if (outImg.depth() == CV_32F) {
        double minVal, maxVal;
        cv::minMaxLoc(outImg, &minVal, &maxVal);
        if (maxVal <= 1.01) {
            outImg.convertTo(outImg, CV_8U, 255.0);
        } else {
            outImg.convertTo(outImg, CV_8U);
        }
    }

    cv::Mat resizedOutput;
    cv::resize(outImg, resizedOutput, input.size(), 0, 0, cv::INTER_LINEAR);
    metrics.postprocessMs = timer.elapsed();

    return resizedOutput;
}

cv::Mat ProcessingWorker::runInference(const cv::Mat &input)
{
    cv::Mat blob = preprocess(input);
    m_u2netNet.setInput(blob);
    std::vector<cv::Mat> outputs;
    m_u2netNet.forward(outputs, m_u2netNet.getUnconnectedOutLayersNames());
    cv::Mat rawMask(320, 320, CV_32F, outputs[0].ptr<float>());
    return postprocess(rawMask, input.size());
}

cv::Mat ProcessingWorker::preprocess(const cv::Mat &input)
{
    cv::Mat blob = cv::dnn::blobFromImage(input, 1.0/255.0, cv::Size(320, 320),
                                          cv::Scalar(0.485*255, 0.456*255, 0.406*255),
                                          true, false);
    blob.convertTo(blob, CV_32F);
    cv::divide(blob, cv::Scalar(0.229, 0.224, 0.225), blob);
    return blob;
}

cv::Mat ProcessingWorker::postprocess(const cv::Mat &rawMask, const cv::Size &originalSize)
{
    double minVal, maxVal;
    cv::minMaxLoc(rawMask, &minVal, &maxVal);

    if (maxVal - minVal < 1e-5) {
        qWarning() << "U2Net: Объект не найден, маска пуста!";
        return cv::Mat::zeros(originalSize, CV_8UC1);
    }

    cv::Mat normalized;
    rawMask.convertTo(normalized, CV_8U, 255.0 / (maxVal - minVal), -minVal * 255.0 / (maxVal - minVal));
    cv::Mat resized;
    cv::resize(normalized, resized, originalSize, 0, 0, cv::INTER_LINEAR);
    cv::Mat binaryMask;
    cv::threshold(resized, binaryMask, 128, 255, cv::THRESH_BINARY);
    return binaryMask;
}

cv::Mat ProcessingWorker::qimageToMat(const QImage &img)
{
    QImage rgb = img.convertToFormat(QImage::Format_RGB888);
    return cv::Mat(rgb.height(), rgb.width(), CV_8UC3,
                   const_cast<uchar*>(rgb.bits()),
                   static_cast<size_t>(rgb.bytesPerLine())).clone();
}

QImage ProcessingWorker::matToQImage(const cv::Mat &mat)
{
    if (mat.channels() == 4) {
        return QImage(mat.data, mat.cols, mat.rows, mat.step,
                      QImage::Format_RGBA8888).copy();
    } else {
        cv::Mat rgb;
        cv::cvtColor(mat, rgb, cv::COLOR_BGR2RGB);
        return QImage(rgb.data, rgb.cols, rgb.rows, rgb.step,
                      QImage::Format_RGB888).copy();
    }
}

QImage ProcessingWorker::stubEnhance(const QImage &img)
{
    QImage result = img.convertToFormat(QImage::Format_ARGB32);
    for (int y = 0; y < result.height(); ++y) {
        QRgb *line = reinterpret_cast<QRgb*>(result.scanLine(y));
        for (int x = 0; x < result.width(); ++x) {
            int a = qAlpha(line[x]);
            int r = qBound(0, qRed(line[x]) + 20, 255);
            int g = qBound(0, qGreen(line[x]) + 20, 255);
            int b = qBound(0, qBlue(line[x]) + 20, 255);
            line[x] = qRgba(r, g, b, a);
        }
    }
    return result;
}

QImage ProcessingWorker::stubStyle(const QImage &img, int styleIndex)
{
    Q_UNUSED(styleIndex);
    QImage result = img.convertToFormat(QImage::Format_ARGB32);
    QImage blurred = result.scaled(result.width()/2, result.height()/2,
                                   Qt::IgnoreAspectRatio, Qt::SmoothTransformation)
                         .scaled(result.width(), result.height(),
                                 Qt::IgnoreAspectRatio, Qt::SmoothTransformation);
    for (int y = 0; y < result.height(); ++y) {
        QRgb *line = reinterpret_cast<QRgb*>(result.scanLine(y));
        const QRgb *blurLine = reinterpret_cast<const QRgb*>(blurred.constScanLine(y));
        for (int x = 0; x < result.width(); ++x) {
            int a = qAlpha(line[x]);
            int r = (qRed(line[x]) + qRed(blurLine[x])) / 2;
            int g = (qGreen(line[x]) + qGreen(blurLine[x])) / 2;
            int b = (qBlue(line[x]) + qBlue(blurLine[x])) / 2;
            line[x] = qRgba(r, g, b, a);
        }
    }
    return result;
}

void ProcessingWorker::processCorrection(const QImage &img, int brightness, float contrast)
{
    if (img.isNull()) {
        emit errorOccurred("Empty image for correction");
        return;
    }

    try {
        cv::Mat input = qimageToMat(img);
        cv::Mat output;

        input.convertTo(output, -1, contrast, brightness);

        cv::Mat rgbOutput;
        if (output.channels() == 3) {
            cv::cvtColor(output, rgbOutput, cv::COLOR_BGR2RGB);
        } else {
            rgbOutput = output;
        }

        QImage result = matToQImage(rgbOutput);
        emit correctionFinished(result);
    } catch (const cv::Exception &e) {
        emit errorOccurred(QString("OpenCV Correction error: %1").arg(e.what()));
    }
}