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

src/segmentationengine.cpp
252 строки7 KB
Артемий
fix: model inputs/outputs
09 июл 2026, 10:50
09 июл 2026, 10:503be26b4
#include "segmentationengine.h"

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

#include <limits>

#ifdef CATEDIT_NCNN
#include "net.h"

namespace {

constexpr int kInputSize = 320;

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.loaded = false;
    state.failed = false;
    state.paramPath.clear();
    state.binPath.clear();

    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 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 ncnn bin:" << binPath << "code:" << modelResult;
        return false;
    }

    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;
}

} // namespace
#endif

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
{
#ifdef CATEDIT_NCNN
    return ensureLoaded();
#else
    return false;
#endif
}

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,
                                             qreal *backgroundPercent, qint64 *preprocessMs,
                                             qint64 *inferenceMs, qint64 *postprocessMs) const
{
#ifdef CATEDIT_NCNN
    if (source.isNull() || !ensureLoaded())
        return QImage();

    const int width = source.width();
    const int height = source.height();
    if (width <= 0 || height <= 0)
        return QImage();

    QElapsedTimer stageTimer;
    stageTimer.start();

    QImage rgbImage = source.convertToFormat(QImage::Format_RGB888);
    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));
        }
    }

    const QImage fullMask = smallMask.scaled(width, height, Qt::IgnoreAspectRatio,
                                             Qt::SmoothTransformation);

    const uchar replacement[3] = {uchar(color.red()), uchar(color.green()), uchar(color.blue())};
    const int total = width * height;
    int backgroundPixels = 0;

    for (int y = 0; y < height; ++y) {
        uchar *line = rgbImage.scanLine(y);
        const uchar *maskLine = fullMask.constScanLine(y);
        for (int x = 0; x < width; ++x) {
            if (maskLine[x] < 128) {
                line[x * 3 + 0] = replacement[0];
                line[x * 3 + 1] = replacement[1];
                line[x * 3 + 2] = replacement[2];
                ++backgroundPixels;
            }
        }
    }

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

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

    return rgbImage;
#else
    Q_UNUSED(source)
    Q_UNUSED(color)
    Q_UNUSED(backgroundPercent)
    Q_UNUSED(preprocessMs)
    Q_UNUSED(inferenceMs)
    Q_UNUSED(postprocessMs)
    return QImage();
#endif
}