src/segmentationengine.cpp - master - limb/catedit-aurora

src/segmentationengine.cpp
312 строк10 KB
Артемий
add: new filters
09 июл 2026, 22:25
09 июл 2026, 22:25882783f
#include "segmentationengine.h"

#include <QDebug>
#include <QElapsedTimer>
#include <QFile>
#include <QtGlobal>

#include <limits>

#include "opencv2/imgproc.hpp"

#ifdef CATEDIT_NCNN
#include "net.h"
#endif

namespace {

constexpr int kInputSize = 320;

cv::Mat qImageToRgbMat(QImage &image)
{
    return cv::Mat(image.height(), image.width(), CV_8UC3, image.bits(),
                   image.bytesPerLine());
}

QImage rgbMatToQImage(const cv::Mat &mat)
{
    if (mat.empty())
        return QImage();

    return QImage(mat.data, mat.cols, mat.rows, mat.step, QImage::Format_RGB888).copy();
}

#ifdef CATEDIT_NCNN
struct SegmentationNetState {
    ncnn::Net net;
    QString paramPath;
    QString binPath;
    bool loaded = false;
    bool failed = false;
};

SegmentationNetState &segmentationNetState()
{
    static thread_local SegmentationNetState state;
    return state;
}

bool loadSegmentationNet(const QString &paramPath, const QString &binPath)
{
    SegmentationNetState &state = segmentationNetState();
    if (state.loaded && state.paramPath == paramPath && state.binPath == binPath)
        return true;

    state.net.clear();
    state.paramPath.clear();
    state.binPath.clear();
    state.loaded = false;
    state.failed = false;

    if (paramPath.isEmpty() || binPath.isEmpty() || !QFile::exists(paramPath)
            || !QFile::exists(binPath)) {
        state.failed = true;
        qWarning() << "Segmentation model files are missing. param:" << paramPath
                   << "exists:" << QFile::exists(paramPath) << "bin:" << binPath
                   << "exists:" << QFile::exists(binPath);
        return false;
    }

    ncnn::Option opt;
    opt.num_threads = 2;
    opt.use_vulkan_compute = false;
    state.net.opt = opt;

    const QByteArray localParam = paramPath.toLocal8Bit();
    const QByteArray localBin = binPath.toLocal8Bit();
    qDebug() << "Loading segmentation model from:" << paramPath << binPath;

    const int paramResult = state.net.load_param(localParam.constData());
    if (paramResult != 0) {
        state.failed = true;
        qWarning() << "Failed to load segmentation ncnn param:" << paramPath
                   << "code:" << paramResult;
        return false;
    }

    const int modelResult = state.net.load_model(localBin.constData());
    if (modelResult != 0) {
        state.failed = true;
        qWarning() << "Failed to load segmentation ncnn bin:" << binPath
                   << "code:" << modelResult;
        return false;
    }

    const std::vector<int> &inputIndexes = state.net.input_indexes();
    const std::vector<int> &outputIndexes = state.net.output_indexes();
    if (inputIndexes.empty() || outputIndexes.empty()) {
        state.failed = true;
        qWarning() << "Segmentation model has no input or output blobs."
                   << "inputs:" << int(inputIndexes.size())
                   << "outputs:" << int(outputIndexes.size());
        return false;
    }

    qDebug() << "Segmentation model input index:" << inputIndexes.front()
             << "output index:" << outputIndexes.front();

    state.paramPath = paramPath;
    state.binPath = binPath;
    state.loaded = true;
    qDebug() << "Segmentation model loaded successfully";
    return true;
}

void resetSegmentationNetState()
{
    SegmentationNetState &state = segmentationNetState();
    state.net.clear();
    state.paramPath.clear();
    state.binPath.clear();
    state.loaded = false;
    state.failed = false;
}
#endif

} // namespace

SegmentationEngine &SegmentationEngine::instance()
{
    static SegmentationEngine engine;
    return engine;
}

void SegmentationEngine::setModelPaths(const QString &paramPath, const QString &binPath)
{
    m_paramPath = paramPath;
    m_binPath = binPath;
    m_loaded = false;
#ifdef CATEDIT_NCNN
    resetSegmentationNetState();
#endif
}

bool SegmentationEngine::isAvailable() const
{
    return ensureLoaded();
}

bool SegmentationEngine::ensureLoaded() const
{
#ifdef CATEDIT_NCNN
    if (m_loaded)
        return true;

    m_loaded = loadSegmentationNet(m_paramPath, m_binPath);
    return m_loaded;
#else
    return false;
#endif
}

QImage SegmentationEngine::replaceBackground(const QImage &source, const QColor &color,
                                             BackgroundMode mode, qreal *backgroundPercent,
                                             qint64 *preprocessMs, qint64 *inferenceMs,
                                             qint64 *postprocessMs) const
{
    if (backgroundPercent)
        *backgroundPercent = 0.0;
    if (preprocessMs)
        *preprocessMs = 0;
    if (inferenceMs)
        *inferenceMs = 0;
    if (postprocessMs)
        *postprocessMs = 0;

#ifndef CATEDIT_NCNN
    Q_UNUSED(source)
    Q_UNUSED(color)
    Q_UNUSED(mode)
    return QImage();
#else
    if (source.isNull() || !ensureLoaded())
        return QImage();

    QElapsedTimer stageTimer;
    stageTimer.start();

    QImage rgbImage = source.convertToFormat(QImage::Format_RGB888);
    const int width = rgbImage.width();
    const int height = rgbImage.height();
    const int total = width * height;
    if (total <= 0)
        return QImage();

    ncnn::Mat input = ncnn::Mat::from_pixels_resize(rgbImage.constBits(), ncnn::Mat::PIXEL_RGB,
                                                    width, height, rgbImage.bytesPerLine(),
                                                    kInputSize, kInputSize);

    const float meanVals[3] = {0.485f * 255.f, 0.456f * 255.f, 0.406f * 255.f};
    const float normVals[3] = {1.f / (0.229f * 255.f), 1.f / (0.224f * 255.f),
                               1.f / (0.225f * 255.f)};
    input.substract_mean_normalize(meanVals, normVals);

    if (preprocessMs)
        *preprocessMs = stageTimer.elapsed();

    stageTimer.restart();

    SegmentationNetState &state = segmentationNetState();
    ncnn::Extractor extractor = state.net.create_extractor();
    extractor.set_light_mode(true);

    const int inputResult = extractor.input("in0", input);
    if (inputResult != 0) {
        qWarning() << "Failed to set segmentation input. code:" << inputResult;
        return QImage();
    }

    ncnn::Mat maskOutput;
    const int extractResult = extractor.extract("out0", maskOutput);
    if (extractResult != 0) {
        qWarning() << "Failed to extract segmentation output. code:" << extractResult;
        return QImage();
    }

    if (maskOutput.w <= 0 || maskOutput.h <= 0) {
        qWarning() << "Segmentation output is empty. w:" << maskOutput.w
                   << "h:" << maskOutput.h << "c:" << maskOutput.c;
        return QImage();
    }

    qDebug() << "Segmentation output shape:"
             << "w:" << maskOutput.w << "h:" << maskOutput.h << "c:" << maskOutput.c;

    if (inferenceMs)
        *inferenceMs = stageTimer.elapsed();

    stageTimer.restart();

    float minMask = std::numeric_limits<float>::max();
    float maxMask = std::numeric_limits<float>::lowest();
    for (int y = 0; y < maskOutput.h; ++y) {
        const float *row = maskOutput.row(y);
        for (int x = 0; x < maskOutput.w; ++x) {
            minMask = qMin(minMask, row[x]);
            maxMask = qMax(maxMask, row[x]);
        }
    }
    qDebug() << "Segmentation mask range:" << minMask << maxMask;

    const float maskRange = maxMask - minMask;
    QImage smallMask(maskOutput.w, maskOutput.h, QImage::Format_Grayscale8);
    for (int y = 0; y < maskOutput.h; ++y) {
        uchar *line = smallMask.scanLine(y);
        const float *row = maskOutput.row(y);
        for (int x = 0; x < maskOutput.w; ++x) {
            const float normalized = maskRange > 0.00001f ? (row[x] - minMask) / maskRange : row[x];
            line[x] = uchar(qBound(0, int(normalized * 255.f + 0.5f), 255));
        }
    }

    cv::Mat smallMaskMat(smallMask.height(), smallMask.width(), CV_8UC1, smallMask.bits(),
                         smallMask.bytesPerLine());
    cv::Mat refinedMask = smallMaskMat.clone();
    cv::medianBlur(refinedMask, refinedMask, 3);

    const cv::Mat kernel = cv::getStructuringElement(cv::MORPH_ELLIPSE, cv::Size(3, 3));
    cv::morphologyEx(refinedMask, refinedMask, cv::MORPH_OPEN, kernel);
    cv::morphologyEx(refinedMask, refinedMask, cv::MORPH_CLOSE, kernel);
    cv::GaussianBlur(refinedMask, refinedMask, cv::Size(5, 5), 0.0);

    cv::Mat fullMask;
    cv::resize(refinedMask, fullMask, cv::Size(width, height), 0, 0, cv::INTER_LINEAR);

    cv::Mat fullImage = qImageToRgbMat(rgbImage);
    cv::Mat backgroundImage;
    if (mode == BackgroundMode::BlurOriginal) {
        cv::GaussianBlur(fullImage, backgroundImage, cv::Size(0, 0), 18.0);
    }

    const cv::Vec3b replacement(color.red(), color.green(), color.blue());
    int backgroundPixels = 0;

    for (int y = 0; y < height; ++y) {
        const uchar *maskLine = fullMask.ptr<uchar>(y);
        cv::Vec3b *imageLine = fullImage.ptr<cv::Vec3b>(y);
        for (int x = 0; x < width; ++x) {
            const float alpha = qBound(0.0f, maskLine[x] / 255.0f, 1.0f);
            if (alpha < 0.5f)
                ++backgroundPixels;

            cv::Vec3b &pixel = imageLine[x];
            const cv::Vec3b backgroundPixel = mode == BackgroundMode::BlurOriginal
                    ? backgroundImage.ptr<cv::Vec3b>(y)[x] : replacement;
            for (int channel = 0; channel < 3; ++channel) {
                pixel[channel] = uchar(qBound(0, int(pixel[channel] * alpha
                                                     + backgroundPixel[channel] * (1.0f - alpha)
                                                     + 0.5f),
                                              255));
            }
        }
    }

    if (backgroundPercent)
        *backgroundPercent = qreal(backgroundPixels) * 100.0 / qreal(total);

    if (postprocessMs)
        *postprocessMs = stageTimer.elapsed();

    return rgbMatToQImage(fullImage);
#endif
}