Merge pull request #9 from occ-ai/roy.external_model_load

External model loading

CMakeLists.txt

@@ -77,6 +77,9 @@ else()
   target_link_libraries(${CMAKE_PROJECT_NAME} PRIVATE OpenCV)
 endif()
 
+# add the vendor folder to the include path
+target_include_directories(${CMAKE_PROJECT_NAME} PRIVATE vendor)
+
 target_sources(${CMAKE_PROJECT_NAME} PRIVATE src/plugin-main.c src/detect-filter.cpp src/detect-filter-info.c
                                              src/obs-utils/obs-utils.cpp src/edgeyolo/edgeyolo_onnxruntime.cpp)
 

README.md

@@ -39,6 +39,7 @@ Current features:
 
 - Detect over 80 categories of objects, using an efficient model ([EdgeYOLO](https://github.com/LSH9832/edgeyolo))
 - 3 Model sizes: Small, Medium and Large
+- Load custom ONNX detection models from disk
 - Control detection threshold
 - Select object category filter (e.g. find only "Person")
 - Masking: Blur, Solid color, Transparent, output binary mask (combine with other plugins!)

data/locale/en-US.ini

@@ -29,3 +29,5 @@ ZoomObject="Zoom Object"
 SingleFirst="Single (First)"
 ZoomSpeed="Zoom Speed"
 DetectedObject="Detected Object"
+ExternalModel="External Model"
+ModelPath="Model Path"

docs/train_model.md

@@ -9,6 +9,12 @@ In this example we will use a COCO dataset from Roboflow. You can get the datase
 
 The dataset is in the COCO format, so we can use it directly with the EdgeYOLO training script.
 
+### Enviroment requirements
+
+You will need a GPU to train the model. The training process is quite slow on a CPU.
+
+On Windows you can should the Windows Subsystem for Linux (WSL) to run the training script.
+
 ## Step 1: Unpack the dataset
 
 Unzip the dataset to a folder. The dataset should have the following structure:
@@ -32,24 +38,34 @@ dataset_folder/
         ...
 ```
 
-## Step 2: Install the required packages
+## Step 2: Install the required files
+
+Go to the EdgeYOLO repository and clone it to your machine:
+
+```bash
+git clone git@github.com:LSH9832/edgeyolo.git
+```
+
+The rest of this guide assumes you have the EdgeYOLO repository cloned to your machine and you are in the root of that repository.
 
 You need to install the required packages to train the model. You can install them using the following command:
 
 ```bash
 pip install -r requirements.txt
 ```
 
+**Make sure to download a pretrained model** from the EdgeYOLO repository. You can download e.g. the EdgeYOLO Tiny LRELU model from [here](https://github.com/LSH9832/edgeyolo/releases/download/v0.0.0/edgeyolo_tiny_lrelu_coco.pth). This will speed up your process tremendously.
+
 ## Step 3: Setup the parameters of the training script
 
-You need to set the parameters of the training script.
+You need to set the parameters of the training script to train on your data.
 Make a copy of the configuration file and set the parameters according to your needs.
 For example for my case, I will set the following parameters in the `params/train/train_coco_aquarium.yaml` file:
 
 ```yaml
 # models & weights------------------------------------------------------------------------------------------------------
-model_cfg: "params/model/edgeyolo_tiny.yaml"         # model structure config file
-weights: "output/train/edgeyolo_tiny_coco_aquarium/last.pth"  # contains model_cfg, set null or a no-exist filename if not use it
+model_cfg: "params/model/edgeyolo_tiny_lrelu_aquarium.yaml"         # model structure config file
+weights: "**!! Set this to the path of your pretrained .pth model !!**"  # contains model_cfg, set null or a no-exist filename if not use it
 use_cfg: false                                       # force using model_cfg instead of cfg in weights to build model
 
 # output----------------------------------------------------------------------------------------------------------------
@@ -63,7 +79,26 @@ batch_size_per_gpu: 8                                # batch size for each GPU
 loader_num_workers: 4                                # number data loader workers for each GPU
 num_threads: 1                                       # pytorch threads number for each GPU
 
-# the rest--------------------------------------------------------------------------------------------------------------
+# device & data type----------------------------------------------------------------------------------------------------
+device: [0]                                 # training device list
+fp16: false                                          # train with fp16 precision
+cudnn_benchmark: false                               # it's useful when multiscale_range is set zero
+
+# the rest of the file ...
+```
+
+Note the gpu device number in the `device` field. You can set it to `[0]` if you have only one GPU.
+
+Note that the `model_cfg` field points to the model configuration file. You can find the model configuration files in the `params/model/` folder. This is required to set the model architecture, but mostly the number of classes. Make a copy of one of the architechtures with a new filename. This is an example of the top of my new `edgeyolo_tiny_lrelu_aquarium.yaml` file:
+
+```yaml
+# parameters
+nc: 6  # number of classes - match the number of classes in the dataset
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+# ...
 ```
 
 You will also need to set up the dataset configuration file `params/dataset/coco_aquarium.yaml`:
@@ -86,13 +121,41 @@ val:
   label: "<...>/Downloads/edgeyolo/Aquarium Combined.v2-raw-1024.coco/valid/_annotations.coco.json"
 
 test:
-  test_dir: "test2017"
+  image_dir: "<...>/Downloads/edgeyolo/Aquarium Combined.v2-raw-1024.coco/test"
+  label: "<...>/Downloads/edgeyolo/Aquarium Combined.v2-raw-1024.coco/test/_annotations.coco.json"
 
 segmentaion_enabled: false
 
 names: ["creatures", "fish", "jellyfish", "penguin", "puffin", "shark", "starfish", "stingray"]
 ```
 
+Notice that you need to provide the list of classes in the `names` field. This should match the classes in the dataset.
+
+For example in the Aquarim dataset we will see in the `train/annotations.json` a field like so:
+
+```json
+    "categories": [
+        {
+            "id": 0,
+            "name": "creatures",
+            "supercategory": "none"
+        },
+        {
+            "id": 1,
+            "name": "fish",
+            "supercategory": "creatures"
+        },
+        {
+            "id": 2,
+            "name": "jellyfish",
+            "supercategory": "creatures"
+        },
+        ...
+    ]
+```
+
+Make sure the order from `categories` (and the `id` field) is maintained in the `names` field.
+
 ## Step 4: Train the model
 
 You can train the model using the following command:
@@ -109,11 +172,15 @@ Best to have a GPU for training.
 After training the model, you can convert it to ONNX format using the `export.py` script from EdgeYOLO.
 
 ```bash
-python export.py --weights output/train/edgeyolo_tiny_coco_aquarium/last.pth --onnx-only 
+python export.py --weights output/train/edgeyolo_tiny_coco_aquarium/best.pth --onnx-only --batch 1
 ```
 
-You will find the ONNX model in the `output/export/` folder.
+You will find the ONNX model in the `output/export/` folder, e.g. `output/export/best/640x640_batch1.onnx`. Rename the file to something more descriptive.
 
 ## Step 6: Use the model with OBS Detect
 
-TBD
+You can now use the ONNX model with the OBS Detect plugin. Just load the model from the plugin settings.
+
+![alt text](image.png)
+
+You will also need a configuration file for the model but that is created automatically by the export script above. It will have the same name as the ONNX model but with a `.json` extension.

src/FilterData.h

@@ -47,6 +47,7 @@ struct filter_data {
 	std::mutex modelMutex;
 
 	std::unique_ptr<edgeyolo_cpp::EdgeYOLOONNXRuntime> edgeyolo;
+	std::vector<std::string> classNames;
 
 #if _WIN32
 	std::wstring modelFilepath;