Merge pull request #186 from saeugetier/182-yolo-image-segmentation

saeugetier · web-flow · commit d196df45e06d · 2025-05-06T23:33:11.000+02:00
182 yolo image segmentation
diff --git a/.gitignore b/.gitignore
@@ -1,4 +1,4 @@
 build
 software/build*
-
+libs
 *.user
diff --git a/models/coco.names b/models/coco.names
@@ -0,0 +1,80 @@
+person
+bicycle
+car
+motorbike
+aeroplane
+bus
+train
+truck
+boat
+traffic light
+fire hydrant
+stop sign
+parking meter
+bench
+bird
+cat
+dog
+horse
+sheep
+cow
+elephant
+bear
+zebra
+giraffe
+backpack
+umbrella
+handbag
+tie
+suitcase
+frisbee
+skis
+snowboard
+sports ball
+kite
+baseball bat
+baseball glove
+skateboard
+surfboard
+tennis racket
+bottle
+wine glass
+cup
+fork
+knife
+spoon
+bowl
+banana
+apple
+sandwich
+orange
+broccoli
+carrot
+hot dog
+pizza
+donut
+cake
+chair
+sofa
+pottedplant
+bed
+diningtable
+toilet
+tvmonitor
+laptop
+mouse
+remote
+keyboard
+cell phone
+microwave
+oven
+toaster
+sink
+refrigerator
+book
+clock
+vase
+scissors
+teddy bear
+hair drier
+toothbrush
diff --git a/models/yolo11n-seg.onnx b/models/yolo11n-seg.onnx
diff --git a/prepare_onnx.sh b/prepare_onnx.sh
@@ -0,0 +1,33 @@
+#!/bin/bash
+
+VERSION="1.21.0"
+TARGET_DIR="libs/onnxruntime"
+
+mkdir -p "$TARGET_DIR"
+
+FILE_NAME="onnxruntime-linux-x64-${VERSION}.tgz"
+DOWNLOAD_URL="https://github.com/microsoft/onnxruntime/releases/download/v${VERSION}/${FILE_NAME}"
+
+echo "Downloading ONNX Runtime version $VERSION from $DOWNLOAD_URL..."
+
+curl -L -o "$FILE_NAME" "$DOWNLOAD_URL"
+
+if [ $? -ne 0 ]; then
+    echo "Error: Failed to download ONNX Runtime version $VERSION."
+    exit 1
+fi
+
+echo "Extracting $FILE_NAME to $TARGET_DIR..."
+tar -xzf "$FILE_NAME" -C "$TARGET_DIR" --strip-components=1
+
+if [ $? -ne 0 ]; then
+    echo "Error: Extraction failed."
+    rm -f "$FILE_NAME"
+    exit 1
+fi
+
+echo "Cleaning up..."
+rm -f "$FILE_NAME"
+
+echo "ONNX Runtime version $VERSION successfully downloaded and extracted to $TARGET_DIR."
+
diff --git a/qml.qrc b/qml.qrc
@@ -94,5 +94,7 @@
         <file>shaders/fisheye.frag.qsb</file>
         <file>shaders/previewshader.frag</file>
         <file>shaders/previewshader.frag.qsb</file>
+        <file>models/yolo11n-seg.onnx</file>
+        <file>models/coco.names</file>
     </qresource>
 </RCC>
diff --git a/qml/SnapshotMenu.qml b/qml/SnapshotMenu.qml
@@ -123,8 +123,8 @@ SnapshotMenuForm {
         failureAnimation.start()
     }
 
-    cameraRenderer.backgroundFilter.method: snapshotSettings.chromaKeyEnabled ? "Chroma" : "None"
-    cameraRenderer.backgroundFilterEnabled: snapshotSettings.chromaKeyEnabled
+    cameraRenderer.backgroundFilter.method: snapshotSettings.chromaKeyEnabled ? "Chroma" : "Neural"
+    cameraRenderer.backgroundFilterEnabled: snapshotSettings.backgroundFilterEnabled
     cameraRenderer.backgroundFilter.keyColor: snapshotSettings.chromaKeyColor
 
     SequentialAnimation
diff --git a/qml/SnapshotSettings.qml b/qml/SnapshotSettings.qml
@@ -6,8 +6,8 @@ SnapshotSettingsForm {
     property alias viewFinderBrightness: settings.viewFinderBrightness
     property alias flashBrightness: settings.flashBrightness
     property alias countdown: settings.countdown
+    property alias backgroundFilterEnabled: settings.backgroundFilterEnabled
     property alias chromaKeyEnabled: settings.chromaKeyEnabled
-    //property alias chromaKeyStrength: settings.chromaKeyEnabled
     property alias chromaKeyColor: settings.chromaKeyColor
 
     Settings
@@ -18,8 +18,8 @@ SnapshotSettingsForm {
         property real viewFinderBrightness: 0.1
         property real flashBrightness: 1.0
         property int countdown: 1
+        property bool backgroundFilterEnabled: false
         property bool chromaKeyEnabled: false
-        //property real chromaKeyStrength: 0.5
         property real chromaKeyColor: 0.5
     }
 
@@ -67,18 +67,18 @@ SnapshotSettingsForm {
         settings.countdown = spinBoxCountdownTime.value
     }
 
+    switchBackgroundFilterEnable.checked: settings.backgroundFilterEnabled
+    switchBackgroundFilterEnable.onCheckedChanged:
+    {
+        settings.backgroundFilterEnabled = switchBackgroundFilterEnable.checked
+    }
+
     switchChromaKeyEnable.checked: settings.chromaKeyEnabled
     switchChromaKeyEnable.onCheckedChanged:
     {
         settings.chromaKeyEnabled = switchChromaKeyEnable.checked
     }
 
-   /* sliderChromaKeyFilterStrength.value: settings.chromaKeyStrength
-    sliderChromaKeyFilterStrength.onValueChanged:
-    {
-        settings.chromaKeyStrength = sliderChromaKeyFilterStrength.value
-    }*/
-
     sliderChromaKeyFilterColor.value: settings.chromaKeyColor
     sliderChromaKeyFilterColor.onValueChanged:
     {
diff --git a/qml/SnapshotSettingsForm.ui.qml b/qml/SnapshotSettingsForm.ui.qml
@@ -12,8 +12,8 @@ Pane {
     property alias sliderViewfinderBrightness: sliderViewfinderBrightness
     property alias sliderFlashBrightness: sliderFlashBrightness
     property alias switchFlashEnable: switchFlashEnable
+    property alias switchBackgroundFilterEnable: switchBackgroundFilterEnable
     property alias switchChromaKeyEnable: switchChromaKeyEnable
-    //property alias sliderChromaKeyFilterStrength: sliderChromaKeyFilterStrength
     property alias sliderChromaKeyFilterColor: sliderChromaKeyFilterColor
     opacity: 0.5
     spacing: 5
@@ -233,7 +233,7 @@ Pane {
 
             RowLayout
             {
-                id: chromaKey
+                id: backgroundFilterLabel
 
                 spacing: 15
 
@@ -245,63 +245,56 @@ Pane {
                 }
 
                 Label {
-                    id: labelChromaKey
+                    id: labelBackgroundFilter
                     font.pixelSize: 32
-                    text: qsTr("Chroma Key")
+                    text: qsTr("Background Filter")
                     font.family: "DejaVu Serif"
                 }
             }
 
             RowLayout
             {
                 Label {
-                    id: labelChromaKeyFilterEnable
+                    id: labelBackgroundFilterEnable
                     font.pixelSize: 32
                     text: "Filter: "
                 }
 
                 Switch {
-                    id: switchChromaKeyEnable
+                    id: switchBackgroundFilterEnable
                     text: checked ? qsTr("enabled") : qsTr("disabled")
                     font.pixelSize: 32
                     font.family: "DejaVu Serif"
                     width: 250
                 }
-
             }
 
-
-           /* RowLayout {
-                id: rowChromaKeyFilterStrength
-                spacing: 20
+            RowLayout
+            {
+                visible: switchBackgroundFilterEnable.checked
 
                 Label {
-                    id: labelChromaKeyFilterStrengthLow
+                    id: labelChromaKeyFilterEnable
                     font.pixelSize: 32
-                    text: "\uE814"
-                    font.family: "fontello"
+                    text: "Type: "
                 }
 
-                Slider {
-                    id: sliderChromaKeyFilterStrength
-                    value: 0.5
-                    from: 0.5
-                    to: 10.0
+                Switch {
+                    id: switchChromaKeyEnable
+                    text: checked ? qsTr("chroma") : qsTr("neural")
+                    font.pixelSize: 32
+                    font.family: "DejaVu Serif"
                     width: 250
                 }
 
-                Label {
-                    id: labelChromaKeyFilterStrengthHigh
-                    font.pixelSize: 32
-                    text: "\uE813"
-                    font.family: "fontello"
-                }
-            } */
+            }
 
             RowLayout {
                 id: rowChromaKeyFilterColor
                 spacing: 20
 
+                visible: switchChromaKeyEnable.checked & switchBackgroundFilterEnable.checked
+
                 Rectangle
                 {
                     color: "Green"
diff --git a/qtbooth.pro b/qtbooth.pro
@@ -24,7 +24,8 @@ SOURCES += src/collageiconmodel.cpp \
     src/selphyprinter.cpp \
     src/standardprinter.cpp \
     src/system.cpp \
-    src/translationhelper.cpp
+    src/translationhelper.cpp \
+    src/yolo11seg.cpp
 
 RESOURCES += qml.qrc
 
@@ -47,7 +48,8 @@ DISTFILES += \
     android/gradlew.bat \
     android/res/values/libs.xml
 
-INCLUDEPATH += src/
+INCLUDEPATH += src/ \
+    libs/onnxruntime/include/ \
 
 HEADERS += \
     src/abstractprinter.h \
@@ -66,11 +68,14 @@ HEADERS += \
     src/selphyprinter.h \
     src/standardprinter.h \
     src/system.h \
-    src/translationhelper.h
+    src/translationhelper.h \
+    src/yolo11seg.h
 
 contains(ANDROID_TARGET_ARCH,x86) {
     ANDROID_PACKAGE_SOURCE_DIR = \
         $$PWD/android
 }
 
 DEFINES += GIT_CURRENT_SHA1="$(shell git -C \""$$_PRO_FILE_PWD_"\" describe)"
+
+LIBS += -L"$$PWD/libs/onnxruntime/lib" -lonnxruntime
diff --git a/src/replacebackgroundvideofilter.cpp b/src/replacebackgroundvideofilter.cpp
@@ -36,7 +36,7 @@ void ReplaceBackgroundVideoFilter::setMethod(QString method)
     {
         mFilterMethod = FilterMethod::CHROMA;
     }
-    else if(method.contains("Neural Net"))
+    else if(method.contains("Neural"))
     {
         mFilterMethod = FilterMethod::NEURAL;
     }
@@ -66,7 +66,7 @@ QString ReplaceBackgroundVideoFilter::getMethod() const
     }
     else if(mFilterMethod == FilterMethod::NEURAL)
     {
-        return QString("Neural Net");
+        return QString("Neural");
     }
     else
     {
@@ -100,7 +100,7 @@ void ReplaceBackgroundVideoFilter::onProcessingFinished(const QImage& maskImage)
     mProcessing = false;
 }
 
-ReplaceBackgroundFilterRunable::ReplaceBackgroundFilterRunable(ReplaceBackgroundVideoFilter* filter) : mFilter(filter)
+ReplaceBackgroundFilterRunable::ReplaceBackgroundFilterRunable(ReplaceBackgroundVideoFilter* filter) : mFilter(filter), mYoloSegmentor(":/models/yolo11n-seg.onnx", ":/models/coco.names" , false)
 {
 }
 
@@ -156,7 +156,22 @@ void ReplaceBackgroundFilterRunable::run(const QVideoFrame& input)
     break;
 
     case ReplaceBackgroundVideoFilter::FilterMethod::NEURAL:
-        break;
+    {
+        std::vector<Segmentation> results = mYoloSegmentor.segment(mMat, 0.2f, 0.45f);
+
+        cv::Mat mask = cv::Mat::zeros(mMat.size(), CV_8UC1);
+
+        std::vector<int> objectFilter = {0};
+
+        mYoloSegmentor.drawSegmentationMask(mask, results, objectFilter);
+
+        std::vector<cv::Mat> channels;
+        split(mMat, channels);
+        channels.push_back(mask);
+        cv::merge(channels, mMat);
+    }
+
+    break;
 
     case ReplaceBackgroundVideoFilter::FilterMethod::NONE:
         break;
diff --git a/src/replacebackgroundvideofilter.h b/src/replacebackgroundvideofilter.h
@@ -7,6 +7,7 @@
 #include <QVideoFrameInput>
 #include <QThread>
 #include <opencv2/core/core.hpp>
+#include "yolo11seg.h"
 
 class ReplaceBackgroundFilterRunable;
 
@@ -129,6 +130,8 @@ public slots:
     bool mYuv;
     ReplaceBackgroundVideoFilter* mFilter;
 
+    YOLOv11SegDetector mYoloSegmentor;
+
 signals:
     void processingFinished(const QImage& maskImage);
 };
diff --git a/src/yolo11seg.cpp b/src/yolo11seg.cpp
diff --git a/src/yolo11seg.h b/src/yolo11seg.h

-Original file line number
+Diff line change
@@ @@ -1,4 +1,4 @@ @@
 build
 software/build*
+-
 +libs
 *.user
Original file line number	Diff line number	Diff line change
`@@ -123,8 +123,8 @@ SnapshotMenuForm {`
`123`	`123`	`failureAnimation.start()`
`124`	`124`	`}`
`125`	`125`
`126`		`- cameraRenderer.backgroundFilter.method: snapshotSettings.chromaKeyEnabled ? "Chroma" : "None"`
`127`		`- cameraRenderer.backgroundFilterEnabled: snapshotSettings.chromaKeyEnabled`
	`126`	`+ cameraRenderer.backgroundFilter.method: snapshotSettings.chromaKeyEnabled ? "Chroma" : "Neural"`
	`127`	`+ cameraRenderer.backgroundFilterEnabled: snapshotSettings.backgroundFilterEnabled`
`128`	`128`	`cameraRenderer.backgroundFilter.keyColor: snapshotSettings.chromaKeyColor`
`129`	`129`
`130`	`130`	`SequentialAnimation`
Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,7 @@ void ReplaceBackgroundVideoFilter::setMethod(QString method)`
`36`	`36`	`{`
`37`	`37`	`mFilterMethod = FilterMethod::CHROMA;`
`38`	`38`	`}`
`39`		`- else if(method.contains("Neural Net"))`
	`39`	`+ else if(method.contains("Neural"))`
`40`	`40`	`{`
`41`	`41`	`mFilterMethod = FilterMethod::NEURAL;`
`42`	`42`	`}`
`@@ -66,7 +66,7 @@ QString ReplaceBackgroundVideoFilter::getMethod() const`
`66`	`66`	`}`
`67`	`67`	`else if(mFilterMethod == FilterMethod::NEURAL)`
`68`	`68`	`{`
`69`		`- return QString("Neural Net");`
	`69`	`+ return QString("Neural");`
`70`	`70`	`}`
`71`	`71`	`else`
`72`	`72`	`{`
`@@ -100,7 +100,7 @@ void ReplaceBackgroundVideoFilter::onProcessingFinished(const QImage& maskImage)`
`100`	`100`	`mProcessing = false;`
`101`	`101`	`}`
`102`	`102`
`103`		`-ReplaceBackgroundFilterRunable::ReplaceBackgroundFilterRunable(ReplaceBackgroundVideoFilter* filter) : mFilter(filter)`
	`103`	`+ReplaceBackgroundFilterRunable::ReplaceBackgroundFilterRunable(ReplaceBackgroundVideoFilter* filter) : mFilter(filter), mYoloSegmentor(":/models/yolo11n-seg.onnx", ":/models/coco.names" , false)`
`104`	`104`	`{`
`105`	`105`	`}`
`106`	`106`
`@@ -156,7 +156,22 @@ void ReplaceBackgroundFilterRunable::run(const QVideoFrame& input)`
`156`	`156`	`break;`
`157`	`157`
`158`	`158`	`case ReplaceBackgroundVideoFilter::FilterMethod::NEURAL:`
`159`		`- break;`
	`159`	`+ {`
	`160`	`+ std::vector<Segmentation> results = mYoloSegmentor.segment(mMat, 0.2f, 0.45f);`
	`161`	`+`
	`162`	`+ cv::Mat mask = cv::Mat::zeros(mMat.size(), CV_8UC1);`
	`163`	`+`
	`164`	`+ std::vector<int> objectFilter = {0};`
	`165`	`+`
	`166`	`+ mYoloSegmentor.drawSegmentationMask(mask, results, objectFilter);`
	`167`	`+`
	`168`	`+ std::vector<cv::Mat> channels;`
	`169`	`+ split(mMat, channels);`
	`170`	`+ channels.push_back(mask);`
	`171`	`+ cv::merge(channels, mMat);`
	`172`	`+ }`
	`173`	`+`
	`174`	`+ break;`
`160`	`175`
`161`	`176`	`case ReplaceBackgroundVideoFilter::FilterMethod::NONE:`
`162`	`177`	`break;`