-
Notifications
You must be signed in to change notification settings - Fork 182
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add object detection training docs #1435
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -7,6 +7,7 @@ PhotonVision supports object detection using neural network accelerator hardware | |
For the 2024 season, PhotonVision ships with a **pre-trained NOTE detector** (shown above), as well as a mechanism for swapping in custom models. Future development will focus on enabling lower friction management of multiple custom models. | ||
|
||
```{image} images/notes-ui.png | ||
|
||
``` | ||
|
||
## Tracking Objects | ||
|
@@ -34,7 +35,51 @@ The same area, aspect ratio, and target orientation/sort parameters from {ref}`r | |
|
||
## Training Custom Models | ||
|
||
Coming soon! | ||
You can build your own YOLOv5 model for use with PhotonVision's Object Detection. However, this process is complex and requires a good understanding of machine learning, but it can be done. Here's an overview of the process. | ||
|
||
### Step 1: Model Training | ||
|
||
You can train the Object Detection model using [Ultralytics' YOLOv5 repository](https://github.com/ultralytics/yolov5), then use [airockchip's YOLOv5 fork](https://github.com/airockchip/yolov5) to export the model currectly before the export to a deployable RKNN file. The specific training setup will vary depending on your dataset and configuration, but here’s an example using `yolov5s`. Other models should work as well. You **will** need [PyTorch](https://pytorch.org/get-started/locally/). This example has been tested on Ubuntu 22.04. | ||
|
||
#### Downloading nessessary files | ||
|
||
```bash | ||
git clone https://github.com/ultralytics/yolov5.git | ||
git clone https://github.com/airockchip/yolov5.git airockchip-yolov5 | ||
wget https://gist.githubusercontent.com/Alex-idk/9a512ca7bd263892ff6991a856f1a458/raw/e8c0c9d8d5a1a60a2bbe72c065e04a261300baac/onnx2rknn.py # This is the onnx to rknn convertion script | ||
``` | ||
|
||
#### Training Command | ||
|
||
Please research what each of these parameters do and adjust them to fit your dataset and training needs. Make sure to change the number of classes in the `models/yolov5s.yaml` file to how many classes are in your dataset, otherwise you will run into problems with class labeling. Currently as of `September 2024` only YOLOv5s models have been tested. | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. "you must know what you are doing" lolz. Is there anything major you could recommend an initial user look at, just to get them started? Things like "Be careful touching these parameters" or "These parameters are inter-linked" or "Here's the values we started with for the note detector" or similar? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The main parameters for training are dependent on your dataset and hardware. I can add a link to a docs page that explains what they do. But I'm also assuming that the people who are making a custom model already have an idea of what they are doing. The special thing about doing it for PV is the conversion to RKNN. |
||
|
||
```bash | ||
python train.py --img 640 --batch 16 --epochs 10 --data path/to/dataset/data.yaml --cfg 'models/yolov5s.yaml' --weights '' --cache | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Might still need a concrete example of what data.yaml is There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It might just be best to point people towards some other resources for preparing yolov5 datasets. |
||
``` | ||
|
||
### Step 2: Exporting the Model to ONNX | ||
|
||
Once your model is trained, the next step is to export it to ONNX format using airockchip's YOLOv5 fork. | ||
|
||
#### Export Command | ||
|
||
```bash | ||
cd /path/to/airockchip-yolov5 && python export.py --weights '/path/to/best.pt' --rknpu --include 'onnx' | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. If the instructions above were used, the path should just be There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Rather than including this as a separate step could it be merged into the current conversion command? My thought is people who want to share models (say the PV devs, or community members) could benefit from a single script that goes from the weights to all of our supported platforms. |
||
``` | ||
|
||
### Step 3: Converting ONNX to RKNN | ||
|
||
Using the `onnx2rknn.py` script, convert the ONNX model to an RKNN file. This script was downloaded in a previous step. | ||
|
||
#### Conversion Command | ||
|
||
Run the script, passing in the ONNX model and a text file containing paths to images from your dataset: | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
So we've got to create this file ourselves? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yeah, I can probably make a revision of the conversion script to just take a path and make the file itself. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. That works too. |
||
|
||
```bash | ||
python onnx2rknn.py /path/to/best.onnx /path/to/export/best.rknn /path/to/imagePaths.txt | ||
``` | ||
|
||
If you have any questions about this process feel free to mention `alex_idk` in the PhotonVision Discord server. | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. A wise man once told me to never put my name on anything :) |
||
|
||
## Uploading Custom Models | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Script looks reasonable, seems like the same thing most online resources have been hinting at.
Do we wanna pull this into the photonvision repo proper?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think this script should be owned by the PV organization. Once we start supporting more platforms it can be updated to go from a single
.onnx
file to all of our supported platform.