You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
This repository was archived by the owner on Aug 28, 2024. It is now read-only.
Copy file name to clipboardExpand all lines: ObjectDetection/README.md
+79-6Lines changed: 79 additions & 6 deletions
Display the source diff
Display the rich diff
Original file line number
Diff line number
Diff line change
@@ -4,6 +4,8 @@
4
4
5
5
[YOLO](https://pjreddie.com/darknet/yolo/) (You Only Look Once) is one of the fastest and most popular object detection models. [YOLOv5](https://github.com/ultralytics/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](https://pytorch.org/hub/ultralytics_yolov5/#load-from-pytorch-hub)). This Object Detection with YOLOv5 iOS sample app uses the PyTorch scripted YOLOv5 model to detect objects of the [80 classes](https://github.com/ultralytics/yolov5/blob/master/data/coco.yaml) trained with the model.
6
6
7
+
**Update 10-07-2021**: A new section of using a custom dataset to fine-tune the YOLOv5 model (aka transfer learning) with steps to change the iOS demo app to use the custom model was added.
8
+
7
9
## Prerequisites
8
10
9
11
* PyTorch 1.9 and torchvision 0.10 (Optional)
@@ -15,24 +17,32 @@
15
17
16
18
To Test Run the Object Detection iOS App, follow the steps below:
17
19
18
-
### 1. Prepare the Model
20
+
### 1. Prepare the model
19
21
20
22
If you don't have the PyTorch environment set up to run the script, you can download the model file [here](https://pytorch-mobile-demo-apps.s3.us-east-2.amazonaws.com/yolov5s.torchscript.ptl) to the `ios-demo-app/ObjectDetection/ObjectDetection` folder, then skip the rest of this step and go to step 2 directly.
21
23
22
24
The Python script `export.py` in the `models` folder of the [YOLOv5 repo](https://github.com/ultralytics/yolov5) is used to generate a TorchScript-formatted YOLOv5 model named `yolov5s.torchscript.ptl` for mobile apps.
23
25
24
-
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):
26
+
Open a Mac/Linux/Windows Terminal, run the following commands:
25
27
26
28
```
27
-
git clone https://github.com/jeffxtang/yolov5
29
+
git clone https://github.com/ultralytics/yolov5
28
30
cd yolov5
29
-
pip install -r requirements.txt
31
+
pip install -r requirements.txt wanb
30
32
```
31
33
32
-
Finally, run the script below to generate the optimized TorchScript Lite Interpreter model and copy the generated model file `yolov5s.torchscript.ptl` to the `ios-demo-app/ObjectDetection/ObjectDetection` folder:
34
+
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`.
35
+
36
+
Edit `export.py` to make the following two changes:
37
+
38
+
* After `f = file.with_suffix('.torchscript.pt')`, add a line `fl = file.with_suffix('.torchscript.ptl')`
39
+
40
+
* 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))`
41
+
42
+
Finally, run the script below to generate the optimized TorchScript Lite Interpreter model and copy the generated model file `yolov5s.torchscript.ptl` to the `ios-demo-app/ObjectDetection/ObjectDetection` folder (the original full JIT model `yolov5s.torchscript.pt` was also generated for comparison):
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.
@@ -56,3 +66,66 @@ Some example images and the detection results are as follows:
56
66
57
67
58
68
69
+
70
+
## Transfer Learning
71
+
72
+
In this section, you'll see how to use an example dataset called [aicook](https://universe.roboflow.com/karel-cornelis-q2qqg/aicook-lcv4d/4), used to detect ingredients in your fridge, to fine-tune the YOLOv5 model. For more info on the YOLOv5 transfer learning, see [here](https://github.com/ultralytics/yolov5/issues/1314). If you use the default YOLOv5 model to do object detection on what's inside your fridge, you'll likely not get good results. That's why you need to have a custom model trained with a dataset like aicook.
73
+
74
+
### 1. Download the custom dataset
75
+
76
+
Simply go to [here](https://universe.roboflow.com/karel-cornelis-q2qqg/aicook-lcv4d/4) 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.
77
+
78
+
### 2. Retrain the YOLOv5 with the custom dataset
79
+
80
+
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](https://wandb.ai), 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 iOS demo app. Below is a sample wandb metrics from 3, 30, and 100 epochs of training:
87
+
88
+
89
+
90
+
### 3. Convert the custom model to lite version
91
+
92
+
With the `export.py` modified as in step 1 `Prepare the model` of the section `Quick Start`, you can convert the new custom model to its TorchScript lite version:
The resulting `best.torchscript.ptl` is located in `runs/train/exp/weights`, which needs to be added to the iOS ObjectDetection demo app project.
99
+
100
+
### 4. Update the demo app
101
+
102
+
In Xcode, first in `ViewController.swift`, change line `private let testImages = ["test1.png", "test2.jpg", "test3.png"]` to `private let testImages = ["aicook1.jpg", "aicook2.jpg", "aicook3.jpg", "test1.png", "test2.jpg", "test3.png"]`
103
+
(The three aicook test images have been added to the repo.)
104
+
105
+
Then change lines in `ObjectDetector.swift`:
106
+
```
107
+
if let filePath = Bundle.main.path(forResource: "yolov5s.torchscript", ofType: "ptl"),
108
+
```
109
+
to:
110
+
```
111
+
if let filePath = Bundle.main.path(forResource: "best.torchscript", ofType: "ptl"),
112
+
```
113
+
and
114
+
```
115
+
if let filePath = Bundle.main.path(forResource: "classes", ofType: "txt"),
116
+
```
117
+
to:
118
+
```
119
+
if let filePath = Bundle.main.path(forResource: "aicook", ofType: "txt"),
120
+
```
121
+
(aicook.txt defines the 30 custom class names, copied from `data.yaml` in the custom dataset downloaded in step 1 of this section.)
122
+
123
+
Finally in `PrePostProcessor.swift`, change line `static let outputColumn = 85` to `static let outputColumn = 35`, which is 5 (left, top, right, bottom, score) + 30 (number of custom classes).
124
+
125
+
Run the app in Xcode and you should see the custom model working on the first three aicook test images:
0 commit comments