Skip to content
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

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 46 additions & 1 deletion docs/source/docs/objectDetection/about-object-detection.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Copy link
Contributor

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?

Copy link
Contributor

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.

```

#### 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.
Copy link
Contributor

Choose a reason for hiding this comment

The 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?

Copy link
Author

@Alex-idk Alex-idk Sep 27, 2024

Choose a reason for hiding this comment

The 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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Might still need a concrete example of what data.yaml is

Copy link
Contributor

Choose a reason for hiding this comment

The 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'
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the instructions above were used, the path should just be airockchip-yolov5 , right?

Copy link
Contributor

Choose a reason for hiding this comment

The 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:
Copy link
Contributor

@Alextopher Alextopher Sep 27, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

and a text file containing paths to images from your dataset:

So we've got to create this file ourselves?

Copy link
Author

Choose a reason for hiding this comment

The 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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

find /path -type f -name "*.jpg" > images.txt

Copy link
Author

Choose a reason for hiding this comment

The 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.
Copy link
Contributor

Choose a reason for hiding this comment

The 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

Expand Down
Loading