YOLO (You Only Look Once) is one of the fastest and most popular object detection models. YOLOv5 is an open-source implementation of the latest version of YOLO (for a quick test of loading YOLOv5 from PyTorch hub for inference, see here). This Object Detection with YOLOv5 Android sample app uses the PyTorch scripted YOLOv5 model to detect objects of the 80 classes trained with the model.
Update 09-30-2021: A new section of using a custom dataset to fine-tune the YOLOv5 model (aka transfer learning) and steps to change the Android project to use the custom model was added.
Prerequisites
PyTorch 1.9.0 or later (Optional)
Python 3.8 (Optional)
Android Pytorch library pytorch_android_lite:1.9.0 and pytorch_android_torchvision:1.9.0
Android Studio 4.0.1 or later
Quick Start
To Test Run the Object Detection Android App, follow the steps below:
1. Prepare the model
If you don't have the PyTorch environment set up to run the script, you can download the model file yolov5s.torchscript.ptlhere to the android-demo-app/ObjectDetection/app/src/main/assets folder, then skip the rest of this step and go to step 2 directly.
Be aware that the downloadable model file was created with PyTorch 1.9.0, matching the PyTorch Android library 1.9.0 specified in the project's build.gradle file as implementation 'org.pytorch:pytorch_android_lite:1.9.0'. If you use a different version of PyTorch to create your model by following the instructions below, make sure you specify the same PyTorch Android library version in the build.gradle file to avoid possible errors caused by the version mismatch. Furthermore, if you want to use the latest PyTorch master code to create the model, follow the steps at Building PyTorch Android from Source and Using the PyTorch Android Libraries Built on how to use the model in Android.
The Python script export.py in the models folder of the YOLOv5 repo is used to generate a TorchScript-formatted YOLOv5 model named yolov5s.torchscript.pt for mobile apps.
Open a Mac/Linux/Windows Terminal, run the following commands (note that we use the fork of the original YOLOv5 repo to make sure the code changes work, but feel free to use the original repo):
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt wanb
Note the steps below have been tested with the commit cd35a009ba964331abccd30f6fa0614224105d39 and if there's any issue with running the script or using the model, try git reset --hard cd35a009ba964331abccd30f6fa0614224105d39.
Edit export.py to make the following two changes:
After f = file.with_suffix('.torchscript.pt'), add a line fl = file.with_suffix('.torchscript.ptl')
After (optimize_for_mobile(ts) if optimize else ts).save(f), add (optimize_for_mobile(ts) if optimize else ts)._save_for_lite_interpreter(str(fl))
Now run the script below to generate the optimized TorchScript lite model and copy the generated model file yolov5s.torchscript.ptl to the android-demo-app/ObjectDetection/app/src/main/assets folder:
Note that small sized version of the YOLOv5 model, which runs faster but with less accuracy, is generated by default when running the export.py. You can also change the value of the weights parameter in the export.py to generate the medium, large, and extra large version of the model.
2. Build with Android Studio
Start Android Studio, then open the project located in android-demo-app/ObjectDetection
3. Run the app
Select an Android emulator or device to run the app. You can go through the included example test images to see the detection results. You can also select a picture from your Android device's Photos library, take a picture with the device camera, or even use live camera to do object detection - see this video for a screencast of the app running.
Some example images and the detection results are as follows:
Transfer Learning
In this section, you'll see how to use an example dataset called aicook, used to detect ingredients in your fridge, to fine-tune the YOLOv5 model. For more info on the YOLOv5 transfer learning, see here.
1. Download the custom dataset
Simply go to here to download the aicook dataset in a zip file. Unzip the file to your yolov5 repo directory, then run cd yolov5; mv train ..; mv valid ..; as the aicook data.yaml specifies the train and val folders to be up one level.
2. Retrain the YOLOv5 with the custom dataset
Run the script below to generate a custom model best.torchscript.pt located in runs/train/exp/weights:
The precision of the model with the epochs set as 3 is very low - less than 0.01 actually; with a tool such as Weights and Biases, which can be set up in a few minutes and has been integrated with YOLOv5, you can find that with --epochs set as 80, the precision gets to be 0.95. But on a CPU machine, you can quickly train a custom model using the command above, then test it in the Android demo app.
3. Convert the custom model to lite version
With the export.py modified in step 1 Prepare the model of the section Quick Start, you can convert the new custom model to its TorchScript lite version:
Object Detection with YOLOv5 on Android
Introduction
YOLO (You Only Look Once) is one of the fastest and most popular object detection models. YOLOv5 is an open-source implementation of the latest version of YOLO (for a quick test of loading YOLOv5 from PyTorch hub for inference, see here). This Object Detection with YOLOv5 Android sample app uses the PyTorch scripted YOLOv5 model to detect objects of the 80 classes trained with the model.
Update 09-30-2021: A new section of using a custom dataset to fine-tune the YOLOv5 model (aka transfer learning) and steps to change the Android project to use the custom model was added.
Prerequisites
Quick Start
To Test Run the Object Detection Android App, follow the steps below:
1. Prepare the model
If you don't have the PyTorch environment set up to run the script, you can download the model file
yolov5s.torchscript.ptlhere to theandroid-demo-app/ObjectDetection/app/src/main/assetsfolder, then skip the rest of this step and go to step 2 directly.Be aware that the downloadable model file was created with PyTorch 1.9.0, matching the PyTorch Android library 1.9.0 specified in the project's
build.gradlefile asimplementation 'org.pytorch:pytorch_android_lite:1.9.0'. If you use a different version of PyTorch to create your model by following the instructions below, make sure you specify the same PyTorch Android library version in thebuild.gradlefile to avoid possible errors caused by the version mismatch. Furthermore, if you want to use the latest PyTorch master code to create the model, follow the steps at Building PyTorch Android from Source and Using the PyTorch Android Libraries Built on how to use the model in Android.The Python script
export.pyin themodelsfolder of the YOLOv5 repo is used to generate a TorchScript-formatted YOLOv5 model namedyolov5s.torchscript.ptfor mobile apps.Open a Mac/Linux/Windows Terminal, run the following commands (note that we use the fork of the original YOLOv5 repo to make sure the code changes work, but feel free to use the original repo):
Note the steps below have been tested with the commit
cd35a009ba964331abccd30f6fa0614224105d39and if there's any issue with running the script or using the model, trygit reset --hard cd35a009ba964331abccd30f6fa0614224105d39.Edit
export.pyto make the following two changes:After
f = file.with_suffix('.torchscript.pt'), add a linefl = file.with_suffix('.torchscript.ptl')After
(optimize_for_mobile(ts) if optimize else ts).save(f), add(optimize_for_mobile(ts) if optimize else ts)._save_for_lite_interpreter(str(fl))Now run the script below to generate the optimized TorchScript lite model and copy the generated model file
yolov5s.torchscript.ptlto theandroid-demo-app/ObjectDetection/app/src/main/assetsfolder:Note that small sized version of the YOLOv5 model, which runs faster but with less accuracy, is generated by default when running the
export.py. You can also change the value of theweightsparameter in theexport.pyto generate the medium, large, and extra large version of the model.2. Build with Android Studio
Start Android Studio, then open the project located in
android-demo-app/ObjectDetection3. Run the app
Select an Android emulator or device to run the app. You can go through the included example test images to see the detection results. You can also select a picture from your Android device's Photos library, take a picture with the device camera, or even use live camera to do object detection - see this video for a screencast of the app running.
Some example images and the detection results are as follows:
Transfer Learning
In this section, you'll see how to use an example dataset called aicook, used to detect ingredients in your fridge, to fine-tune the YOLOv5 model. For more info on the YOLOv5 transfer learning, see here.
1. Download the custom dataset
Simply go to here to download the aicook dataset in a zip file. Unzip the file to your
yolov5repo directory, then runcd yolov5; mv train ..; mv valid ..;as the aicookdata.yamlspecifies thetrainandvalfolders to be up one level.2. Retrain the YOLOv5 with the custom dataset
Run the script below to generate a custom model
best.torchscript.ptlocated inruns/train/exp/weights:The precision of the model with the epochs set as 3 is very low - less than 0.01 actually; with a tool such as Weights and Biases, which can be set up in a few minutes and has been integrated with YOLOv5, you can find that with
--epochsset as 80, the precision gets to be 0.95. But on a CPU machine, you can quickly train a custom model using the command above, then test it in the Android demo app.3. Convert the custom model to lite version
With the
export.pymodified in step 1Prepare the modelof the sectionQuick Start, you can convert the new custom model to its TorchScript lite version:The resulting
best.torchscript.ptlis located inruns/train/exp/weights, which needs to be copied to the Android demo app's assets folder.4. Update the demo app
In Android Studio, first in
MainActivity.java, change line:to:
(The three aicook test images have been added to the repo.)
and change lines:
to:
(aicook.txt defines the 30 custom class names, copied from
data.yamlin the custom dataset downloaded in step 1.)Then in
PrePostProcessor.java, change lineprivate static int mOutputColumn = 85;toprivate static int mOutputColumn = 35;.Run the app in Android Studio and you should see the custom model working on the first three aicook test images.