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

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

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

#ifdef CATEDIT_NCNN
#include "net.h"

namespace {

struct StyleNetState {
    ncnn::Net net;
    QString paramPath;
    QString binPath;
    bool loaded = false;
    bool failed = false;
};

StyleNetState &styleNetState()
{
    static thread_local StyleNetState state;
    return state;
}

bool loadStyleNet(const QString &paramPath, const QString &binPath)
{
    StyleNetState &state = styleNetState();
    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() << "Style 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 style model from:" << paramPath << binPath;

    const int paramResult = state.net.load_param_bin(localParam.constData());
    if (paramResult != 0) {
        state.failed = true;
        qWarning() << "Failed to load style 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 style 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() << "Style model has no input or output blobs."
                   << "inputs:" << int(inputIndexes.size())
                   << "outputs:" << int(outputIndexes.size());
        return false;
    }

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

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

void resetStyleNetState()
{
    StyleNetState &state = styleNetState();
    state.net.clear();
    state.paramPath.clear();
    state.binPath.clear();
    state.loaded = false;
    state.failed = false;
}

} // namespace
#endif

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

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

bool NeuralStyleEngine::isAvailable() const
{
#ifdef CATEDIT_NCNN
    return ensureLoaded();
#else
    return false;
#endif
}

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

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

QImage NeuralStyleEngine::applyStyle(const QImage &source, 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();

    const int maxSide = 512;
    int inferWidth = width;
    int inferHeight = height;
    if (qMax(width, height) > maxSide) {
        if (width >= height) {
            inferWidth = maxSide;
            inferHeight = qMax(1, height * maxSide / width);
        } else {
            inferHeight = maxSide;
            inferWidth = qMax(1, width * maxSide / height);
        }
    }

    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(),
                                                    inferWidth, inferHeight);

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

    stageTimer.restart();

    StyleNetState &state = styleNetState();
    ncnn::Extractor extractor = state.net.create_extractor();
    extractor.set_light_mode(true);
    const std::vector<int> &inputIndexes = state.net.input_indexes();
    const std::vector<int> &outputIndexes = state.net.output_indexes();
    if (inputIndexes.empty() || outputIndexes.empty()) {
        qWarning() << "Style model input/output indexes are empty";
        return QImage();
    }

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

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

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

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

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

    stageTimer.restart();

    QImage styled(output.w, output.h, QImage::Format_RGB888);
    output.to_pixels(styled.bits(), ncnn::Mat::PIXEL_RGB, styled.bytesPerLine());

    QImage result = styled;
    if (styled.width() != width || styled.height() != height)
        result = styled.scaled(width, height, Qt::IgnoreAspectRatio, Qt::SmoothTransformation);

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

    return result;
#else
    Q_UNUSED(source)
    Q_UNUSED(preprocessMs)
    Q_UNUSED(inferenceMs)
    Q_UNUSED(postprocessMs)
    return QImage();
#endif
}