This model is used by duckpuc.com to detect dicks on an edge device and apply amusing filters to them. The trained model can be downloaded in PyTorch or Darknet formats. You can learn more on how to integrate it with a ncnn+wasm app in the dick-detection-app.
Please use Python 3.8 with all requirements.txt dependencies installed, including torch>=1.6
. Do not use python 3.9.
$ pip install -r requirements.txt
The data for this model was created and labeled using roboflow
, a platform used for training of computer vision models. If you want to explore the data used to train this model (beware, adult content), please download it from here.
If you want to create your own data set, you can use roboflow
to create a new Image Dataset
.
To use it in this project, use the Export > YOLO Darknet
format, download the zip
file and extract it into the modeling\data\roboflow
folder.
The file structure needs to be slightly adjusted to be used by this model. Do so by running:
$ cd modeling
$ python prepare_roboflow_data.py
This will populate the modeling/data/images
and modeling/data/labels
folders.
It will also update the content modeling.names
and modeling.names
files.
✔️ Check
By the end of this step you should have images and labels in
modeling/data/images
andmodeling/data/labels
respectivey. You should also have the configurations inmodeling.names
andmodeling.names
updated according to your model.
Run this code to train the model based on the pretrained weights yolo-fastest.weights from COCO.
$ python3 train.py --cfg yolo-fastest.cfg --data data/modeling.data --weights weights/yolo-fastest.weights --epochs 120
The training process would cost several hours. When the training ended, you can use the following script to get the training graphs.
$ python3 -c "from utils import utils; utils.plot_results()"
At this stage, the model weights file best.pt is stored in PyTorch
format.
If you want to use the model with the ncnn+wasm app, you need to first convert them to Darknet
format by running:
$ python3 -c "from models import *; convert('cfg/yolo-fastest.cfg', 'weights/best.pt')"
✔️ Check
By the end of this step you should have a file named
best.weights
in themodeling\weights
directory.
With the model you got, the inference could be performed directly in this format: python3 detect.py --source ...
For instance, if you want to use your webcam, please run python3 detect.py --source 0
.
The deployment part works based on NCNN and WASM.
At first, you need to compile the NCNN library. For more details, use the tutorial for compiling NCNN library by @waittim.
When the compilation process of NCNN has been completed, you can start to use various tools in the ncnn/build/tools folder to help us convert the model.
For example, you can copy the yolo-fastest.cfg and best.weights files of the darknet model to the ncnn/build/tools/darknet, and use this code to convert to the NCNN model.
./darknet2ncnn yolo-fastest.cfg best.weights yolo-fastest.param yolo-fastest.bin 1
For compacting the model size, you can move the yolo-fastest.param and yolo-fastest.bin to ncnn/build/tools, then run the ncnnoptimize program.
./ncnnoptimize yolo-fastest.param yolo-fastest.bin yolo-fastest-opt.param yolo-fastest-opt.bin 65536
✔️ Check
By the end of this step you should have two files named
yolo-fastest-opt.bin
andyolo-fastest-opt.param
in yourncnn\build\tools\
directory.
Now you have the yolo-fastest-opt.param and yolo-fastest-opt.bin as our final model. Please continue the deployment as documented in the dick-detection-app repository.
The modeling part is modified based on the code from waittim which is in turn based on Ultralytics. The model used is modified from the Yolo-Fastest model shared by @dog-qiuqiu.