segplugin.h

#pragma once
 
#include <ModelManager.h>
#include <NeuralNetFactory.h>
 
#include <QCoreApplication>
#include <QObject>
#include <QTranslator>
#include <memory>
#include <optional>
 
#include "ibase.h"
#include "imask.h"
#include "ipreview.h"
 
class SegmentationSettingsWidget;
 
class SegmentationPlugin : public PLUG::IBase, public PLUG::IMask, public PLUG::IPreview {
    Q_OBJECT
    Q_PLUGIN_METADATA(IID "iVS3D.IMask")
    Q_PLUGIN_METADATA(IID "iVS3D.IPreview")
    Q_PLUGIN_METADATA(IID "iVS3D.IBase")
    Q_INTERFACES(PLUG::IMask PLUG::IBase PLUG::IPreview)
 
   public:
    using PLUG::IBase::IBase;
 
    SegmentationPlugin();
 
    // IBase interface
    QString getName() const override { return tr("Segmentation Masks"); }
    PLUG::SettingsWidgetResult getSettingsWidget() override;
    QMap<QString, QVariant> getSettings() const override;
    QString getSettingsString() const override;
    PLUG::ApplySettingsResult applySettings(
        const QMap<QString, QVariant>& settings) override;
    void onCudaChanged(bool enabled) override;
    void deactivate() override;
 
    // IMask interface
    PLUG::MaskResult generateMask(const PLUG::MaskData& data) override;
 
    // IPreview interface
    VIS::VisualizationResult generatePreview(const PLUG::PreviewData& data) override;
 
   signals:
    void syncSettingsWidget(QString selectedModelName, float overlayAlpha);
 
   private slots:
    void onModelChanged(const QString& modelName, bool isValid);
    void onClassSelectionChanged(const QVector<uint>& selectedClassIds);
    void onOverlayAlphaChanged(float value);
 
   private:
    std::optional<PLUG::Error> ensureModelReady();
 
    tl::expected<NN::Tensor, NN::NeuralError> runInference(
        const cv::Mat& image);
    tl::expected<cv::Mat, NN::NeuralError> runColorization(
        const NN::Tensor& inferenceTensor);
    tl::expected<cv::Mat, NN::NeuralError> runMasking(
        const NN::Tensor& inferenceTensor);
 
    MCFG::ModelManager m_modelManager{QCoreApplication::applicationDirPath() +
                                "/plugins/resources/neural_network_models"};
 
    bool m_useCuda = false;
    float m_overlayAlpha = 0.5f;
 
    QString m_currentModelName;
    NN::NeuralNetPtr m_currentModel;
    std::unique_ptr<QTranslator> m_segmentationTranslator;
 
    // cache for the last inference result (if any and valid)
    struct SegmentationCache {
        uint idx;
        cv::Mat image;
        NN::Tensor inferenceTensor;
        cv::Mat colorizedImage =
            cv::Mat();                  // optional: can be empty if invalid
        cv::Mat maskImage = cv::Mat();  // optional: can be empty if invalid
        long inferenceDurationMs = 0;   // Time taken for inference
    };
    std::optional<SegmentationCache> m_cache;
};