Skip to content

乔峰昇 / tamper_det

Go to a project
Toggle navigation
Toggle navigation pinning
94ecd1ae by 乔峰昇
Options

update code

1 parent b3277458
Inline Side-by-side
Showing 37 changed files with 265 additions and 2610 deletions
.idea/.gitignore 0 → 100644
1 # Default ignored files
2 /shelf/
3 /workspace.xml
.idea/inspectionProfiles/Project_Default.xml 0 → 100644
1 <component name="InspectionProjectProfileManager">
2 <profile version="1.0">
3 <option name="myName" value="Project Default" />
4 <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
5 <option name="ignoredPackages">
6 <value>
7 <list size="2">
8 <item index="0" class="java.lang.String" itemvalue="psutil" />
9 <item index="1" class="java.lang.String" itemvalue="thop" />
10 </list>
11 </value>
12 </option>
13 </inspection_tool>
14 </profile>
15 </component>
...\ No newline at end of file ...\ No newline at end of file
.idea/inspectionProfiles/profiles_settings.xml 0 → 100644
1 <component name="InspectionProjectProfileManager">
2 <settings>
3 <option name="USE_PROJECT_PROFILE" value="false" />
4 <version value="1.0" />
5 </settings>
6 </component>
...\ No newline at end of file ...\ No newline at end of file
.idea/misc.xml 0 → 100644
1 <?xml version="1.0" encoding="UTF-8"?>
2 <project version="4">
3 <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.6 (workenv)" project-jdk-type="Python SDK" />
4 </project>
...\ No newline at end of file ...\ No newline at end of file
.idea/modules.xml 0 → 100644
1 <?xml version="1.0" encoding="UTF-8"?>
2 <project version="4">
3 <component name="ProjectModuleManager">
4 <modules>
5 <module fileurl="file://$PROJECT_DIR$/.idea/tamper_det.iml" filepath="$PROJECT_DIR$/.idea/tamper_det.iml" />
6 </modules>
7 </component>
8 </project>
...\ No newline at end of file ...\ No newline at end of file
.idea/tamper_det.iml 0 → 100644
1 <?xml version="1.0" encoding="UTF-8"?>
2 <module type="PYTHON_MODULE" version="4">
3 <component name="NewModuleRootManager">
4 <content url="file://$MODULE_DIR$" />
5 <orderEntry type="jdk" jdkName="Python 3.6 (workenv)" jdkType="Python SDK" />
6 <orderEntry type="sourceFolder" forTests="false" />
7 </component>
8 <component name="PyDocumentationSettings">
9 <option name="format" value="PLAIN" />
10 <option name="myDocStringFormat" value="Plain" />
11 </component>
12 </module>
...\ No newline at end of file ...\ No newline at end of file
.idea/vcs.xml 0 → 100644
1 <?xml version="1.0" encoding="UTF-8"?>
2 <project version="4">
3 <component name="VcsDirectoryMappings">
4 <mapping directory="$PROJECT_DIR$" vcs="Git" />
5 </component>
6 </project>
...\ No newline at end of file ...\ No newline at end of file
README.md deleted 100644 → 0
1 <div align="center">
2
3 Hi, I'm [Glenn Jocher](https://www.linkedin.com/in/glenn-jocher/), author of [YOLOv5](https://github.com/ultralytics/yolov5) 🚀.
4
5 I'd like to invite you to attend the world's first-ever YOLO conference: [#YOLOVISION22](https://ultralytics.com/yolo-vision)!
6
7 This virtual event takes place on **September 27th, 2022** with talks from the world's leading vision AI experts from Google, OpenMMLab's MMDetection, Baidu's PaddlePaddle, Meituan's YOLOv6, Weight & Biases, Roboflow, Neural Magic, OctoML and of course Ultralytics YOLOv5 and many others.
8
9 Save your spot at https://ultralytics.com/yolo-vision!
10
11 <a align="center" href="https://ultralytics.com/yolo-vision" target="_blank">
12 <img width="850" src="https://user-images.githubusercontent.com/26833433/192165285-a13a370c-60dc-4bf2-9d63-0573fb8e4d44.jpg"></a>
13
14 ## <div align="center"></div>
15 <br>
16
17 <p>
18 <a align="center" href="https://ultralytics.com/yolov5" target="_blank">
19 <img width="850" src="https://github.com/ultralytics/assets/raw/master/yolov5/v62/splash_readme.png"></a>
20 <br><br>
21 <a href="https://play.google.com/store/apps/details?id=com.ultralytics.ultralytics_app" style="text-decoration:none;">
22 <img src="https://raw.githubusercontent.com/ultralytics/assets/master/app/google-play.svg" width="15%" alt="" /></a>&nbsp;
23 <a href="https://apps.apple.com/xk/app/ultralytics/id1583935240" style="text-decoration:none;">
24 <img src="https://raw.githubusercontent.com/ultralytics/assets/master/app/app-store.svg" width="15%" alt="" /></a>
25 </p>
26
27 English | [简体中文](.github/README_cn.md)
28 <br>
29 <div>
30 <a href="https://github.com/ultralytics/yolov5/actions/workflows/ci-testing.yml"><img src="https://github.com/ultralytics/yolov5/actions/workflows/ci-testing.yml/badge.svg" alt="YOLOv5 CI"></a>
31 <a href="https://zenodo.org/badge/latestdoi/264818686"><img src="https://zenodo.org/badge/264818686.svg" alt="YOLOv5 Citation"></a>
32 <a href="https://hub.docker.com/r/ultralytics/yolov5"><img src="https://img.shields.io/docker/pulls/ultralytics/yolov5?logo=docker" alt="Docker Pulls"></a>
33 <br>
34 <a href="https://bit.ly/yolov5-paperspace-notebook"><img src="https://assets.paperspace.io/img/gradient-badge.svg" alt="Run on Gradient"></a>
35 <a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
36 <a href="https://www.kaggle.com/ultralytics/yolov5"><img src="https://kaggle.com/static/images/open-in-kaggle.svg" alt="Open In Kaggle"></a>
37 </div>
38
39 <br>
40 <p>
41 YOLOv5 🚀 is a family of object detection architectures and models pretrained on the COCO dataset, and represents <a href="https://ultralytics.com">Ultralytics</a>
42 open-source research into future vision AI methods, incorporating lessons learned and best practices evolved over thousands of hours of research and development.
43 </p>
44
45 <div align="center">
46 <a href="https://github.com/ultralytics" style="text-decoration:none;">
47 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-github.png" width="2%" alt="" /></a>
48 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="2%" alt="" />
49 <a href="https://www.linkedin.com/company/ultralytics" style="text-decoration:none;">
50 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-linkedin.png" width="2%" alt="" /></a>
51 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="2%" alt="" />
52 <a href="https://twitter.com/ultralytics" style="text-decoration:none;">
53 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-twitter.png" width="2%" alt="" /></a>
54 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="2%" alt="" />
55 <a href="https://www.producthunt.com/@glenn_jocher" style="text-decoration:none;">
56 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-producthunt.png" width="2%" alt="" /></a>
57 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="2%" alt="" />
58 <a href="https://youtube.com/ultralytics" style="text-decoration:none;">
59 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-youtube.png" width="2%" alt="" /></a>
60 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="2%" alt="" />
61 <a href="https://www.facebook.com/ultralytics" style="text-decoration:none;">
62 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-facebook.png" width="2%" alt="" /></a>
63 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="2%" alt="" />
64 <a href="https://www.instagram.com/ultralytics/" style="text-decoration:none;">
65 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-instagram.png" width="2%" alt="" /></a>
66 </div>
67 </div>
68
69
70 ## <div align="center">Documentation</div>
71
72 See the [YOLOv5 Docs](https://docs.ultralytics.com) for full documentation on training, testing and deployment.
73
74 ## <div align="center">Quick Start Examples</div>
75
76 <details open>
77 <summary>Install</summary>
78
79 Clone repo and install [requirements.txt](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) in a
80 [**Python>=3.7.0**](https://www.python.org/) environment, including
81 [**PyTorch>=1.7**](https://pytorch.org/get-started/locally/).
82
83 ```bash
84 git clone https://github.com/ultralytics/yolov5 # clone
85 cd yolov5
86 pip install -r requirements.txt # install
87 ```
88
89 </details>
90
91 <details open>
92 <summary>Inference</summary>
93
94 YOLOv5 [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36) inference. [Models](https://github.com/ultralytics/yolov5/tree/master/models) download automatically from the latest
95 YOLOv5 [release](https://github.com/ultralytics/yolov5/releases).
96
97 ```python
98 import torch
99
100 # Model
101 model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # or yolov5n - yolov5x6, custom
102
103 # Images
104 img = 'https://ultralytics.com/images/zidane.jpg' # or file, Path, PIL, OpenCV, numpy, list
105
106 # Inference
107 results = model(img)
108
109 # Results
110 results.print() # or .show(), .save(), .crop(), .pandas(), etc.
111 ```
112
113 </details>
114
115 <details>
116 <summary>Inference with detect.py</summary>
117
118 `detect.py` runs inference on a variety of sources, downloading [models](https://github.com/ultralytics/yolov5/tree/master/models) automatically from
119 the latest YOLOv5 [release](https://github.com/ultralytics/yolov5/releases) and saving results to `runs/detect`.
120
121 ```bash
122 python detect.py --source 0 # webcam
123 img.jpg # image
124 vid.mp4 # video
125 screen # screenshot
126 path/ # directory
127 'path/*.jpg' # glob
128 'https://youtu.be/Zgi9g1ksQHc' # YouTube
129 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
130 ```
131
132 </details>
133
134 <details>
135 <summary>Training</summary>
136
137 The commands below reproduce YOLOv5 [COCO](https://github.com/ultralytics/yolov5/blob/master/data/scripts/get_coco.sh)
138 results. [Models](https://github.com/ultralytics/yolov5/tree/master/models)
139 and [datasets](https://github.com/ultralytics/yolov5/tree/master/data) download automatically from the latest
140 YOLOv5 [release](https://github.com/ultralytics/yolov5/releases). Training times for YOLOv5n/s/m/l/x are
141 1/2/4/6/8 days on a V100 GPU ([Multi-GPU](https://github.com/ultralytics/yolov5/issues/475) times faster). Use the
142 largest `--batch-size` possible, or pass `--batch-size -1` for
143 YOLOv5 [AutoBatch](https://github.com/ultralytics/yolov5/pull/5092). Batch sizes shown for V100-16GB.
144
145 ```bash
146 python train.py --data coco.yaml --cfg yolov5n.yaml --weights '' --batch-size 128
147 yolov5s 64
148 yolov5m 40
149 yolov5l 24
150 yolov5x 16
151 ```
152
153 <img width="800" src="https://user-images.githubusercontent.com/26833433/90222759-949d8800-ddc1-11ea-9fa1-1c97eed2b963.png">
154
155 </details>
156
157 <details open>
158 <summary>Tutorials</summary>
159
160 - [Train Custom Data](https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data)  🚀 RECOMMENDED
161 - [Tips for Best Training Results](https://github.com/ultralytics/yolov5/wiki/Tips-for-Best-Training-Results)  ☘️
162 RECOMMENDED
163 - [Multi-GPU Training](https://github.com/ultralytics/yolov5/issues/475)
164 - [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36) 🌟 NEW
165 - [TFLite, ONNX, CoreML, TensorRT Export](https://github.com/ultralytics/yolov5/issues/251) 🚀
166 - [NVIDIA Jetson Nano Deployment](https://github.com/ultralytics/yolov5/issues/9627) 🌟 NEW
167 - [Test-Time Augmentation (TTA)](https://github.com/ultralytics/yolov5/issues/303)
168 - [Model Ensembling](https://github.com/ultralytics/yolov5/issues/318)
169 - [Model Pruning/Sparsity](https://github.com/ultralytics/yolov5/issues/304)
170 - [Hyperparameter Evolution](https://github.com/ultralytics/yolov5/issues/607)
171 - [Transfer Learning with Frozen Layers](https://github.com/ultralytics/yolov5/issues/1314)
172 - [Architecture Summary](https://github.com/ultralytics/yolov5/issues/6998) 🌟 NEW
173 - [Weights & Biases Logging](https://github.com/ultralytics/yolov5/issues/1289)
174 - [Roboflow for Datasets, Labeling, and Active Learning](https://github.com/ultralytics/yolov5/issues/4975)  🌟 NEW
175 - [ClearML Logging](https://github.com/ultralytics/yolov5/tree/master/utils/loggers/clearml) 🌟 NEW
176 - [Deci Platform](https://github.com/ultralytics/yolov5/wiki/Deci-Platform) 🌟 NEW
177 - [Comet Logging](https://github.com/ultralytics/yolov5/tree/master/utils/loggers/comet) 🌟 NEW
178
179 </details>
180
181
182 ## <div align="center">Integrations</div>
183
184 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/image-integrations-loop.png" width="100%" />
185
186 <div align="center">
187 <a href="https://bit.ly/yolov5-readme-comet">
188 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-comet.png" width="10%" /></a>
189 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="9%" height="0" alt="" />
190 <a href="https://bit.ly/yolov5-deci-platform">
191 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-deci.png" width="10%" /></a>
192 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="9%" height="0" alt="" />
193 <a href="https://cutt.ly/yolov5-readme-clearml">
194 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-clearml.png" width="10%" /></a>
195 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="9%" height="0" alt="" />
196 <a href="https://roboflow.com/?ref=ultralytics">
197 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-roboflow.png" width="10%" /></a>
198 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="9%" height="0" alt="" />
199 <a href="https://wandb.ai/site?utm_campaign=repo_yolo_readme">
200 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-wb.png" width="10%" /></a>
201 </div>
202
203 |Comet ⭐ NEW|Deci ⭐ NEW|ClearML ⭐ NEW|Roboflow|Weights & Biases
204 |:-:|:-:|:-:|:-:|:-:|
205 |Visualize model metrics and predictions and upload models and datasets in realtime with [Comet](https://bit.ly/yolov5-readme-comet)|Automatically compile and quantize YOLOv5 for better inference performance in one click at [Deci](https://bit.ly/yolov5-deci-platform)|Automatically track, visualize and even remotely train YOLOv5 using [ClearML](https://cutt.ly/yolov5-readme-clearml) (open-source!)|Label and export your custom datasets directly to YOLOv5 for training with [Roboflow](https://roboflow.com/?ref=ultralytics) |Automatically track and visualize all your YOLOv5 training runs in the cloud with [Weights & Biases](https://wandb.ai/site?utm_campaign=repo_yolo_readme)
206
207
208 ## <div align="center">Why YOLOv5</div>
209
210 YOLOv5 has been designed to be super easy to get started and simple to learn. We prioritize real-world results.
211
212 <p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/155040763-93c22a27-347c-4e3c-847a-8094621d3f4e.png"></p>
213 <details>
214 <summary>YOLOv5-P5 640 Figure (click to expand)</summary>
215
216 <p align="left"><img width="800" src="https://user-images.githubusercontent.com/26833433/155040757-ce0934a3-06a6-43dc-a979-2edbbd69ea0e.png"></p>
217 </details>
218 <details>
219 <summary>Figure Notes (click to expand)</summary>
220
221 - **COCO AP val** denotes mAP@0.5:0.95 metric measured on the 5000-image [COCO val2017](http://cocodataset.org) dataset over various inference sizes from 256 to 1536.
222 - **GPU Speed** measures average inference time per image on [COCO val2017](http://cocodataset.org) dataset using a [AWS p3.2xlarge](https://aws.amazon.com/ec2/instance-types/p3/) V100 instance at batch-size 32.
223 - **EfficientDet** data from [google/automl](https://github.com/google/automl) at batch size 8.
224 - **Reproduce** by `python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n6.pt yolov5s6.pt yolov5m6.pt yolov5l6.pt yolov5x6.pt`
225
226 </details>
227
228 ### Pretrained Checkpoints
229
230 | Model | size<br><sup>(pixels) | mAP<sup>val<br>0.5:0.95 | mAP<sup>val<br>0.5 | Speed<br><sup>CPU b1<br>(ms) | Speed<br><sup>V100 b1<br>(ms) | Speed<br><sup>V100 b32<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>@640 (B) |
231 |------------------------------------------------------------------------------------------------------|-----------------------|-------------------------|--------------------|------------------------------|-------------------------------|--------------------------------|--------------------|------------------------|
232 | [YOLOv5n](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5n.pt) | 640 | 28.0 | 45.7 | **45** | **6.3** | **0.6** | **1.9** | **4.5** |
233 | [YOLOv5s](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5s.pt) | 640 | 37.4 | 56.8 | 98 | 6.4 | 0.9 | 7.2 | 16.5 |
234 | [YOLOv5m](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5m.pt) | 640 | 45.4 | 64.1 | 224 | 8.2 | 1.7 | 21.2 | 49.0 |
235 | [YOLOv5l](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5l.pt) | 640 | 49.0 | 67.3 | 430 | 10.1 | 2.7 | 46.5 | 109.1 |
236 | [YOLOv5x](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5x.pt) | 640 | 50.7 | 68.9 | 766 | 12.1 | 4.8 | 86.7 | 205.7 |
237 | | | | | | | | | |
238 | [YOLOv5n6](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5n6.pt) | 1280 | 36.0 | 54.4 | 153 | 8.1 | 2.1 | 3.2 | 4.6 |
239 | [YOLOv5s6](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5s6.pt) | 1280 | 44.8 | 63.7 | 385 | 8.2 | 3.6 | 12.6 | 16.8 |
240 | [YOLOv5m6](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5m6.pt) | 1280 | 51.3 | 69.3 | 887 | 11.1 | 6.8 | 35.7 | 50.0 |
241 | [YOLOv5l6](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5l6.pt) | 1280 | 53.7 | 71.3 | 1784 | 15.8 | 10.5 | 76.8 | 111.4 |
242 | [YOLOv5x6](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5x6.pt)<br>+ [TTA][TTA] | 1280<br>1536 | 55.0<br>**55.8** | 72.7<br>**72.7** | 3136<br>- | 26.2<br>- | 19.4<br>- | 140.7<br>- | 209.8<br>- |
243
244 <details>
245 <summary>Table Notes (click to expand)</summary>
246
247 - All checkpoints are trained to 300 epochs with default settings. Nano and Small models use [hyp.scratch-low.yaml](https://github.com/ultralytics/yolov5/blob/master/data/hyps/hyp.scratch-low.yaml) hyps, all others use [hyp.scratch-high.yaml](https://github.com/ultralytics/yolov5/blob/master/data/hyps/hyp.scratch-high.yaml).
248 - **mAP<sup>val</sup>** values are for single-model single-scale on [COCO val2017](http://cocodataset.org) dataset.<br>Reproduce by `python val.py --data coco.yaml --img 640 --conf 0.001 --iou 0.65`
249 - **Speed** averaged over COCO val images using a [AWS p3.2xlarge](https://aws.amazon.com/ec2/instance-types/p3/) instance. NMS times (~1 ms/img) not included.<br>Reproduce by `python val.py --data coco.yaml --img 640 --task speed --batch 1`
250 - **TTA** [Test Time Augmentation](https://github.com/ultralytics/yolov5/issues/303) includes reflection and scale augmentations.<br>Reproduce by `python val.py --data coco.yaml --img 1536 --iou 0.7 --augment`
251
252 </details>
253
254 ## <div align="center">Classification ⭐ NEW</div>
255
256 YOLOv5 [release v6.2](https://github.com/ultralytics/yolov5/releases) brings support for classification model training, validation, prediction and export! We've made training classifier models super simple. Click below to get started.
257
258 <details>
259 <summary>Classification Checkpoints (click to expand)</summary>
260
261 <br>
262
263 We trained YOLOv5-cls classification models on ImageNet for 90 epochs using a 4xA100 instance, and we trained ResNet and EfficientNet models alongside with the same default training settings to compare. We exported all models to ONNX FP32 for CPU speed tests and to TensorRT FP16 for GPU speed tests. We ran all speed tests on Google [Colab Pro](https://colab.research.google.com/signup) for easy reproducibility.
264
265 | Model | size<br><sup>(pixels) | acc<br><sup>top1 | acc<br><sup>top5 | Training<br><sup>90 epochs<br>4xA100 (hours) | Speed<br><sup>ONNX CPU<br>(ms) | Speed<br><sup>TensorRT V100<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>@224 (B) |
266 |----------------------------------------------------------------------------------------------------|-----------------------|------------------|------------------|----------------------------------------------|--------------------------------|-------------------------------------|--------------------|------------------------|
267 | [YOLOv5n-cls](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5n-cls.pt) | 224 | 64.6 | 85.4 | 7:59 | **3.3** | **0.5** | **2.5** | **0.5** |
268 | [YOLOv5s-cls](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5s-cls.pt) | 224 | 71.5 | 90.2 | 8:09 | 6.6 | 0.6 | 5.4 | 1.4 |
269 | [YOLOv5m-cls](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5m-cls.pt) | 224 | 75.9 | 92.9 | 10:06 | 15.5 | 0.9 | 12.9 | 3.9 |
270 | [YOLOv5l-cls](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5l-cls.pt) | 224 | 78.0 | 94.0 | 11:56 | 26.9 | 1.4 | 26.5 | 8.5 |
271 | [YOLOv5x-cls](https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5x-cls.pt) | 224 | **79.0** | **94.4** | 15:04 | 54.3 | 1.8 | 48.1 | 15.9 |
272 | |
273 | [ResNet18](https://github.com/ultralytics/yolov5/releases/download/v6.2/resnet18.pt) | 224 | 70.3 | 89.5 | **6:47** | 11.2 | 0.5 | 11.7 | 3.7 |
274 | [ResNet34](https://github.com/ultralytics/yolov5/releases/download/v6.2/resnet34.pt) | 224 | 73.9 | 91.8 | 8:33 | 20.6 | 0.9 | 21.8 | 7.4 |
275 | [ResNet50](https://github.com/ultralytics/yolov5/releases/download/v6.2/resnet50.pt) | 224 | 76.8 | 93.4 | 11:10 | 23.4 | 1.0 | 25.6 | 8.5 |
276 | [ResNet101](https://github.com/ultralytics/yolov5/releases/download/v6.2/resnet101.pt) | 224 | 78.5 | 94.3 | 17:10 | 42.1 | 1.9 | 44.5 | 15.9 |
277 | |
278 | [EfficientNet_b0](https://github.com/ultralytics/yolov5/releases/download/v6.2/efficientnet_b0.pt) | 224 | 75.1 | 92.4 | 13:03 | 12.5 | 1.3 | 5.3 | 1.0 |
279 | [EfficientNet_b1](https://github.com/ultralytics/yolov5/releases/download/v6.2/efficientnet_b1.pt) | 224 | 76.4 | 93.2 | 17:04 | 14.9 | 1.6 | 7.8 | 1.5 |
280 | [EfficientNet_b2](https://github.com/ultralytics/yolov5/releases/download/v6.2/efficientnet_b2.pt) | 224 | 76.6 | 93.4 | 17:10 | 15.9 | 1.6 | 9.1 | 1.7 |
281 | [EfficientNet_b3](https://github.com/ultralytics/yolov5/releases/download/v6.2/efficientnet_b3.pt) | 224 | 77.7 | 94.0 | 19:19 | 18.9 | 1.9 | 12.2 | 2.4 |
282
283 <details>
284 <summary>Table Notes (click to expand)</summary>
285
286 - All checkpoints are trained to 90 epochs with SGD optimizer with `lr0=0.001` and `weight_decay=5e-5` at image size 224 and all default settings.<br>Runs logged to https://wandb.ai/glenn-jocher/YOLOv5-Classifier-v6-2
287 - **Accuracy** values are for single-model single-scale on [ImageNet-1k](https://www.image-net.org/index.php) dataset.<br>Reproduce by `python classify/val.py --data ../datasets/imagenet --img 224`
288 - **Speed** averaged over 100 inference images using a Google [Colab Pro](https://colab.research.google.com/signup) V100 High-RAM instance.<br>Reproduce by `python classify/val.py --data ../datasets/imagenet --img 224 --batch 1`
289 - **Export** to ONNX at FP32 and TensorRT at FP16 done with `export.py`. <br>Reproduce by `python export.py --weights yolov5s-cls.pt --include engine onnx --imgsz 224`
290 </details>
291 </details>
292
293 <details>
294 <summary>Classification Usage Examples (click to expand)</summary>
295
296 ### Train
297 YOLOv5 classification training supports auto-download of MNIST, Fashion-MNIST, CIFAR10, CIFAR100, Imagenette, Imagewoof, and ImageNet datasets with the `--data` argument. To start training on MNIST for example use `--data mnist`.
298
299 ```bash
300 # Single-GPU
301 python classify/train.py --model yolov5s-cls.pt --data cifar100 --epochs 5 --img 224 --batch 128
302
303 # Multi-GPU DDP
304 python -m torch.distributed.run --nproc_per_node 4 --master_port 1 classify/train.py --model yolov5s-cls.pt --data imagenet --epochs 5 --img 224 --device 0,1,2,3
305 ```
306
307 ### Val
308 Validate YOLOv5m-cls accuracy on ImageNet-1k dataset:
309 ```bash
310 bash data/scripts/get_imagenet.sh --val # download ImageNet val split (6.3G, 50000 images)
311 python classify/val.py --weights yolov5m-cls.pt --data ../datasets/imagenet --img 224 # validate
312 ```
313
314 ### Predict
315 Use pretrained YOLOv5s-cls.pt to predict bus.jpg:
316 ```bash
317 python classify/predict.py --weights yolov5s-cls.pt --data data/images/bus.jpg
318 ```
319 ```python
320 model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s-cls.pt') # load from PyTorch Hub
321 ```
322
323 ### Export
324 Export a group of trained YOLOv5s-cls, ResNet and EfficientNet models to ONNX and TensorRT:
325 ```bash
326 python export.py --weights yolov5s-cls.pt resnet50.pt efficientnet_b0.pt --include onnx engine --img 224
327 ```
328 </details>
329
330
331 ## <div align="center">Environments</div>
332
333 Get started in seconds with our verified environments. Click each icon below for details.
334
335 <div align="center">
336 <a href="https://bit.ly/yolov5-paperspace-notebook">
337 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-gradient.png" width="10%" /></a>
338 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="5%" alt="" />
339 <a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb">
340 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-colab-small.png" width="10%" /></a>
341 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="5%" alt="" />
342 <a href="https://www.kaggle.com/ultralytics/yolov5">
343 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-kaggle-small.png" width="10%" /></a>
344 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="5%" alt="" />
345 <a href="https://hub.docker.com/r/ultralytics/yolov5">
346 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-docker-small.png" width="10%" /></a>
347 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="5%" alt="" />
348 <a href="https://github.com/ultralytics/yolov5/wiki/AWS-Quickstart">
349 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-aws-small.png" width="10%" /></a>
350 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="5%" alt="" />
351 <a href="https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart">
352 <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-gcp-small.png" width="10%" /></a>
353 </div>
354
355
356 ## <div align="center">Contribute</div>
357
358 We love your input! We want to make contributing to YOLOv5 as easy and transparent as possible. Please see our [Contributing Guide](CONTRIBUTING.md) to get started, and fill out the [YOLOv5 Survey](https://ultralytics.com/survey?utm_source=github&utm_medium=social&utm_campaign=Survey) to send us feedback on your experiences. Thank you to all our contributors!
359
360 <!-- SVG image from https://opencollective.com/ultralytics/contributors.svg?width=990 -->
361 <a href="https://github.com/ultralytics/yolov5/graphs/contributors"><img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/image-contributors-1280.png" /></a>
362
363 ## <div align="center">Contact</div>
364
365 For YOLOv5 bugs and feature requests please visit [GitHub Issues](https://github.com/ultralytics/yolov5/issues). For business inquiries or
366 professional support requests please visit [https://ultralytics.com/contact](https://ultralytics.com/contact).
367
368 <br>
369 <div align="center">
370 <a href="https://github.com/ultralytics" style="text-decoration:none;">
371 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-github.png" width="3%" alt="" /></a>
372 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="3%" alt="" />
373 <a href="https://www.linkedin.com/company/ultralytics" style="text-decoration:none;">
374 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-linkedin.png" width="3%" alt="" /></a>
375 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="3%" alt="" />
376 <a href="https://twitter.com/ultralytics" style="text-decoration:none;">
377 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-twitter.png" width="3%" alt="" /></a>
378 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="3%" alt="" />
379 <a href="https://www.producthunt.com/@glenn_jocher" style="text-decoration:none;">
380 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-producthunt.png" width="3%" alt="" /></a>
381 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="3%" alt="" />
382 <a href="https://youtube.com/ultralytics" style="text-decoration:none;">
383 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-youtube.png" width="3%" alt="" /></a>
384 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="3%" alt="" />
385 <a href="https://www.facebook.com/ultralytics" style="text-decoration:none;">
386 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-facebook.png" width="3%" alt="" /></a>
387 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-transparent.png" width="3%" alt="" />
388 <a href="https://www.instagram.com/ultralytics/" style="text-decoration:none;">
389 <img src="https://github.com/ultralytics/assets/raw/master/social/logo-social-instagram.png" width="3%" alt="" /></a>
390 </div>
391
392 [assets]: https://github.com/ultralytics/yolov5/releases
393 [tta]: https://github.com/ultralytics/yolov5/issues/303
classify/predict.py deleted 100644 → 0
1 # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
2 """
3 Run YOLOv5 classification inference on images, videos, directories, globs, YouTube, webcam, streams, etc.
4
5 Usage - sources:
6 $ python classify/predict.py --weights yolov5s-cls.pt --source 0 # webcam
7 img.jpg # image
8 vid.mp4 # video
9 path/ # directory
10 'path/*.jpg' # glob
11 'https://youtu.be/Zgi9g1ksQHc' # YouTube
12 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
13
14 Usage - formats:
15 $ python classify/predict.py --weights yolov5s-cls.pt # PyTorch
16 yolov5s-cls.torchscript # TorchScript
17 yolov5s-cls.onnx # ONNX Runtime or OpenCV DNN with --dnn
18 yolov5s-cls.xml # OpenVINO
19 yolov5s-cls.engine # TensorRT
20 yolov5s-cls.mlmodel # CoreML (macOS-only)
21 yolov5s-cls_saved_model # TensorFlow SavedModel
22 yolov5s-cls.pb # TensorFlow GraphDef
23 yolov5s-cls.tflite # TensorFlow Lite
24 yolov5s-cls_edgetpu.tflite # TensorFlow Edge TPU
25 yolov5s-cls_paddle_model # PaddlePaddle
26 """
27
28 import argparse
29 import os
30 import platform
31 import sys
32 from pathlib import Path
33
34 import torch
35 import torch.nn.functional as F
36
37 FILE = Path(__file__).resolve()
38 ROOT = FILE.parents[1] # YOLOv5 root directory
39 if str(ROOT) not in sys.path:
40 sys.path.append(str(ROOT)) # add ROOT to PATH
41 ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
42
43 from models.common import DetectMultiBackend
44 from utils.augmentations import classify_transforms
45 from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
46 from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
47 increment_path, print_args, strip_optimizer)
48 from utils.plots import Annotator
49 from utils.torch_utils import select_device, smart_inference_mode
50
51
52 @smart_inference_mode()
53 def run(
54 weights=ROOT / 'yolov5s-cls.pt', # model.pt path(s)
55 source=ROOT / 'data/images', # file/dir/URL/glob/screen/0(webcam)
56 data=ROOT / 'data/coco128.yaml', # dataset.yaml path
57 imgsz=(224, 224), # inference size (height, width)
58 device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
59 view_img=False, # show results
60 save_txt=False, # save results to *.txt
61 nosave=False, # do not save images/videos
62 augment=False, # augmented inference
63 visualize=False, # visualize features
64 update=False, # update all models
65 project=ROOT / 'runs/predict-cls', # save results to project/name
66 name='exp', # save results to project/name
67 exist_ok=False, # existing project/name ok, do not increment
68 half=False, # use FP16 half-precision inference
69 dnn=False, # use OpenCV DNN for ONNX inference
70 vid_stride=1, # video frame-rate stride
71 ):
72 source = str(source)
73 save_img = not nosave and not source.endswith('.txt') # save inference images
74 is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
75 is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
76 webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
77 screenshot = source.lower().startswith('screen')
78 if is_url and is_file:
79 source = check_file(source) # download
80
81 # Directories
82 save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
83 (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
84
85 # Load model
86 device = select_device(device)
87 model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
88 stride, names, pt = model.stride, model.names, model.pt
89 imgsz = check_img_size(imgsz, s=stride) # check image size
90
91 # Dataloader
92 bs = 1 # batch_size
93 if webcam:
94 view_img = check_imshow()
95 dataset = LoadStreams(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride)
96 bs = len(dataset)
97 elif screenshot:
98 dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
99 else:
100 dataset = LoadImages(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride)
101 vid_path, vid_writer = [None] * bs, [None] * bs
102
103 # Run inference
104 model.warmup(imgsz=(1 if pt else bs, 3, *imgsz)) # warmup
105 seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
106 for path, im, im0s, vid_cap, s in dataset:
107 with dt[0]:
108 im = torch.Tensor(im).to(model.device)
109 im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
110 if len(im.shape) == 3:
111 im = im[None] # expand for batch dim
112
113 # Inference
114 with dt[1]:
115 results = model(im)
116
117 # Post-process
118 with dt[2]:
119 pred = F.softmax(results, dim=1) # probabilities
120
121 # Process predictions
122 for i, prob in enumerate(pred): # per image
123 seen += 1
124 if webcam: # batch_size >= 1
125 p, im0, frame = path[i], im0s[i].copy(), dataset.count
126 s += f'{i}: '
127 else:
128 p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
129
130 p = Path(p) # to Path
131 save_path = str(save_dir / p.name) # im.jpg
132 txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # im.txt
133
134 s += '%gx%g ' % im.shape[2:] # print string
135 annotator = Annotator(im0, example=str(names), pil=True)
136
137 # Print results
138 top5i = prob.argsort(0, descending=True)[:5].tolist() # top 5 indices
139 s += f"{', '.join(f'{names[j]} {prob[j]:.2f}' for j in top5i)}, "
140
141 # Write results
142 text = '\n'.join(f'{prob[j]:.2f} {names[j]}' for j in top5i)
143 if save_img or view_img: # Add bbox to image
144 annotator.text((32, 32), text, txt_color=(255, 255, 255))
145 if save_txt: # Write to file
146 with open(f'{txt_path}.txt', 'a') as f:
147 f.write(text + '\n')
148
149 # Stream results
150 im0 = annotator.result()
151 if view_img:
152 if platform.system() == 'Linux' and p not in windows:
153 windows.append(p)
154 cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO) # allow window resize (Linux)
155 cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
156 cv2.imshow(str(p), im0)
157 cv2.waitKey(1) # 1 millisecond
158
159 # Save results (image with detections)
160 if save_img:
161 if dataset.mode == 'image':
162 cv2.imwrite(save_path, im0)
163 else: # 'video' or 'stream'
164 if vid_path[i] != save_path: # new video
165 vid_path[i] = save_path
166 if isinstance(vid_writer[i], cv2.VideoWriter):
167 vid_writer[i].release() # release previous video writer
168 if vid_cap: # video
169 fps = vid_cap.get(cv2.CAP_PROP_FPS)
170 w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
171 h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
172 else: # stream
173 fps, w, h = 30, im0.shape[1], im0.shape[0]
174 save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos
175 vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
176 vid_writer[i].write(im0)
177
178 # Print time (inference-only)
179 LOGGER.info(f"{s}{dt[1].dt * 1E3:.1f}ms")
180
181 # Print results
182 t = tuple(x.t / seen * 1E3 for x in dt) # speeds per image
183 LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
184 if save_txt or save_img:
185 s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
186 LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
187 if update:
188 strip_optimizer(weights[0]) # update model (to fix SourceChangeWarning)
189
190
191 def parse_opt():
192 parser = argparse.ArgumentParser()
193 parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-cls.pt', help='model path(s)')
194 parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)')
195 parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
196 parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[224], help='inference size h,w')
197 parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
198 parser.add_argument('--view-img', action='store_true', help='show results')
199 parser.add_argument('--save-txt', action='store_false', help='save results to *.txt')
200 parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
201 parser.add_argument('--augment', action='store_true', help='augmented inference')
202 parser.add_argument('--visualize', action='store_true', help='visualize features')
203 parser.add_argument('--update', action='store_true', help='update all models')
204 parser.add_argument('--project', default=ROOT / 'runs/predict-cls', help='save results to project/name')
205 parser.add_argument('--name', default='exp', help='save results to project/name')
206 parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
207 parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
208 parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
209 parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
210 opt = parser.parse_args()
211 opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
212 print_args(vars(opt))
213 return opt
214
215
216 def main(opt):
217 check_requirements(exclude=('tensorboard', 'thop'))
218 run(**vars(opt))
219
220
221 if __name__ == "__main__":
222 opt = parse_opt()
223 main(opt)
classify/train.py deleted 100644 → 0
1 # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
2 """
3 Train a YOLOv5 classifier model on a classification dataset
4
5 Usage - Single-GPU training:
6 $ python classify/train.py --model yolov5s-cls.pt --data imagenette160 --epochs 5 --img 224
7
8 Usage - Multi-GPU DDP training:
9 $ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 classify/train.py --model yolov5s-cls.pt --data imagenet --epochs 5 --img 224 --device 0,1,2,3
10
11 Datasets: --data mnist, fashion-mnist, cifar10, cifar100, imagenette, imagewoof, imagenet, or 'path/to/data'
12 YOLOv5-cls models: --model yolov5n-cls.pt, yolov5s-cls.pt, yolov5m-cls.pt, yolov5l-cls.pt, yolov5x-cls.pt
13 Torchvision models: --model resnet50, efficientnet_b0, etc. See https://pytorch.org/vision/stable/models.html
14 """
15
16 import argparse
17 import os
18 import subprocess
19 import sys
20 import time
21 from copy import deepcopy
22 from datetime import datetime
23 from pathlib import Path
24
25 import torch
26 import torch.distributed as dist
27 import torch.hub as hub
28 import torch.optim.lr_scheduler as lr_scheduler
29 import torchvision
30 from torch.cuda import amp
31 from tqdm import tqdm
32
33 FILE = Path(__file__).resolve()
34 ROOT = FILE.parents[1] # YOLOv5 root directory
35 if str(ROOT) not in sys.path:
36 sys.path.append(str(ROOT)) # add ROOT to PATH
37 ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
38
39 from classify import val as validate
40 from models.experimental import attempt_load
41 from models.yolo import ClassificationModel, DetectionModel
42 from utils.dataloaders import create_classification_dataloader
43 from utils.general import (DATASETS_DIR, LOGGER, WorkingDirectory, check_git_status, check_requirements, colorstr,
44 download, increment_path, init_seeds, print_args, yaml_save)
45 from utils.loggers import GenericLogger
46 from utils.plots import imshow_cls
47 from utils.torch_utils import (ModelEMA, model_info, reshape_classifier_output, select_device, smart_DDP,
48 smart_optimizer, smartCrossEntropyLoss, torch_distributed_zero_first)
49
50 LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html
51 RANK = int(os.getenv('RANK', -1))
52 WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
53
54
55 def train(opt, device):
56 init_seeds(opt.seed + 1 + RANK, deterministic=True)
57 save_dir, data, bs, epochs, nw, imgsz, pretrained = \
58 opt.save_dir, Path(opt.data), opt.batch_size, opt.epochs, min(os.cpu_count() - 1, opt.workers), \
59 opt.imgsz, str(opt.pretrained).lower() == 'true'
60 cuda = device.type != 'cpu'
61
62 # Directories
63 wdir = save_dir / 'weights'
64 wdir.mkdir(parents=True, exist_ok=True) # make dir
65 last, best = wdir / 'last.pt', wdir / 'best.pt'
66
67 # Save run settings
68 yaml_save(save_dir / 'opt.yaml', vars(opt))
69
70 # Logger
71 logger = GenericLogger(opt=opt, console_logger=LOGGER) if RANK in {-1, 0} else None
72
73 # Download Dataset
74 with torch_distributed_zero_first(LOCAL_RANK), WorkingDirectory(ROOT):
75 data_dir = data if data.is_dir() else (DATASETS_DIR / data)
76 if not data_dir.is_dir():
77 LOGGER.info(f'\nDataset not found ⚠️, missing path {data_dir}, attempting download...')
78 t = time.time()
79 if str(data) == 'imagenet':
80 subprocess.run(f"bash {ROOT / 'data/scripts/get_imagenet.sh'}", shell=True, check=True)
81 else:
82 url = f'https://github.com/ultralytics/yolov5/releases/download/v1.0/{data}.zip'
83 download(url, dir=data_dir.parent)
84 s = f"Dataset download success ✅ ({time.time() - t:.1f}s), saved to {colorstr('bold', data_dir)}\n"
85 LOGGER.info(s)
86
87 # Dataloaders
88 nc = len([x for x in (data_dir / 'train').glob('*') if x.is_dir()]) # number of classes
89 trainloader = create_classification_dataloader(path=data_dir / 'train',
90 imgsz=imgsz,
91 batch_size=bs // WORLD_SIZE,
92 augment=True,
93 cache=opt.cache,
94 rank=LOCAL_RANK,
95 workers=nw)
96
97 test_dir = data_dir / 'test' if (data_dir / 'test').exists() else data_dir / 'val' # data/test or data/val
98 if RANK in {-1, 0}:
99 testloader = create_classification_dataloader(path=test_dir,
100 imgsz=imgsz,
101 batch_size=bs // WORLD_SIZE * 2,
102 augment=False,
103 cache=opt.cache,
104 rank=-1,
105 workers=nw)
106
107 # Model
108 with torch_distributed_zero_first(LOCAL_RANK), WorkingDirectory(ROOT):
109 if Path(opt.model).is_file() or opt.model.endswith('.pt'):
110 model = attempt_load(opt.model, device='cpu', fuse=False)
111 elif opt.model in torchvision.models.__dict__: # TorchVision models i.e. resnet50, efficientnet_b0
112 model = torchvision.models.__dict__[opt.model](weights='IMAGENET1K_V1' if pretrained else None)
113 else:
114 m = hub.list('ultralytics/yolov5') # + hub.list('pytorch/vision') # models
115 raise ModuleNotFoundError(f'--model {opt.model} not found. Available models are: \n' + '\n'.join(m))
116 if isinstance(model, DetectionModel):
117 LOGGER.warning("WARNING ⚠️ pass YOLOv5 classifier model with '-cls' suffix, i.e. '--model yolov5s-cls.pt'")
118 model = ClassificationModel(model=model, nc=nc, cutoff=opt.cutoff or 10) # convert to classification model
119 reshape_classifier_output(model, nc) # update class count
120 for m in model.modules():
121 if not pretrained and hasattr(m, 'reset_parameters'):
122 m.reset_parameters()
123 if isinstance(m, torch.nn.Dropout) and opt.dropout is not None:
124 m.p = opt.dropout # set dropout
125 for p in model.parameters():
126 p.requires_grad = True # for training
127 model = model.to(device)
128
129 # Info
130 if RANK in {-1, 0}:
131 model.names = trainloader.dataset.classes # attach class names
132 model.transforms = testloader.dataset.torch_transforms # attach inference transforms
133 model_info(model)
134 if opt.verbose:
135 LOGGER.info(model)
136 images, labels = next(iter(trainloader))
137 file = imshow_cls(images[:25], labels[:25], names=model.names, f=save_dir / 'train_images.jpg')
138 logger.log_images(file, name='Train Examples')
139 logger.log_graph(model, imgsz) # log model
140
141 # Optimizer
142 optimizer = smart_optimizer(model, opt.optimizer, opt.lr0, momentum=0.9, decay=opt.decay)
143
144 # Scheduler
145 lrf = 0.01 # final lr (fraction of lr0)
146 # lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - lrf) + lrf # cosine
147 lf = lambda x: (1 - x / epochs) * (1 - lrf) + lrf # linear
148 scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
149 # scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=lr0, total_steps=epochs, pct_start=0.1,
150 # final_div_factor=1 / 25 / lrf)
151
152 # EMA
153 ema = ModelEMA(model) if RANK in {-1, 0} else None
154
155 # DDP mode
156 if cuda and RANK != -1:
157 model = smart_DDP(model)
158
159 # Train
160 t0 = time.time()
161 criterion = smartCrossEntropyLoss(label_smoothing=opt.label_smoothing) # loss function
162 best_fitness = 0.0
163 scaler = amp.GradScaler(enabled=cuda)
164 val = test_dir.stem # 'val' or 'test'
165 LOGGER.info(f'Image sizes {imgsz} train, {imgsz} test\n'
166 f'Using {nw * WORLD_SIZE} dataloader workers\n'
167 f"Logging results to {colorstr('bold', save_dir)}\n"
168 f'Starting {opt.model} training on {data} dataset with {nc} classes for {epochs} epochs...\n\n'
169 f"{'Epoch':>10}{'GPU_mem':>10}{'train_loss':>12}{f'{val}_loss':>12}{'top1_acc':>12}{'top5_acc':>12}")
170 for epoch in range(epochs): # loop over the dataset multiple times
171 tloss, vloss, fitness = 0.0, 0.0, 0.0 # train loss, val loss, fitness
172 model.train()
173 if RANK != -1:
174 trainloader.sampler.set_epoch(epoch)
175 pbar = enumerate(trainloader)
176 if RANK in {-1, 0}:
177 pbar = tqdm(enumerate(trainloader), total=len(trainloader), bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}')
178 for i, (images, labels) in pbar: # progress bar
179 images, labels = images.to(device, non_blocking=True), labels.to(device)
180
181 # Forward
182 with amp.autocast(enabled=cuda): # stability issues when enabled
183 loss = criterion(model(images), labels)
184
185 # Backward
186 scaler.scale(loss).backward()
187
188 # Optimize
189 scaler.unscale_(optimizer) # unscale gradients
190 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradients
191 scaler.step(optimizer)
192 scaler.update()
193 optimizer.zero_grad()
194 if ema:
195 ema.update(model)
196
197 if RANK in {-1, 0}:
198 # Print
199 tloss = (tloss * i + loss.item()) / (i + 1) # update mean losses
200 mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) # (GB)
201 pbar.desc = f"{f'{epoch + 1}/{epochs}':>10}{mem:>10}{tloss:>12.3g}" + ' ' * 36
202
203 # Test
204 if i == len(pbar) - 1: # last batch
205 top1, top5, vloss = validate.run(model=ema.ema,
206 dataloader=testloader,
207 criterion=criterion,
208 pbar=pbar) # test accuracy, loss
209 fitness = top1 # define fitness as top1 accuracy
210
211 # Scheduler
212 scheduler.step()
213
214 # Log metrics
215 if RANK in {-1, 0}:
216 # Best fitness
217 if fitness > best_fitness:
218 best_fitness = fitness
219
220 # Log
221 metrics = {
222 "train/loss": tloss,
223 f"{val}/loss": vloss,
224 "metrics/accuracy_top1": top1,
225 "metrics/accuracy_top5": top5,
226 "lr/0": optimizer.param_groups[0]['lr']} # learning rate
227 logger.log_metrics(metrics, epoch)
228
229 # Save model
230 final_epoch = epoch + 1 == epochs
231 if (not opt.nosave) or final_epoch:
232 ckpt = {
233 'epoch': epoch,
234 'best_fitness': best_fitness,
235 'model': deepcopy(ema.ema).half(), # deepcopy(de_parallel(model)).half(),
236 'ema': None, # deepcopy(ema.ema).half(),
237 'updates': ema.updates,
238 'optimizer': None, # optimizer.state_dict(),
239 'opt': vars(opt),
240 'date': datetime.now().isoformat()}
241
242 # Save last, best and delete
243 torch.save(ckpt, last)
244 if best_fitness == fitness:
245 torch.save(ckpt, best)
246 del ckpt
247
248 # Train complete
249 if RANK in {-1, 0} and final_epoch:
250 LOGGER.info(f'\nTraining complete ({(time.time() - t0) / 3600:.3f} hours)'
251 f"\nResults saved to {colorstr('bold', save_dir)}"
252 f"\nPredict: python classify/predict.py --weights {best} --source im.jpg"
253 f"\nValidate: python classify/val.py --weights {best} --data {data_dir}"
254 f"\nExport: python export.py --weights {best} --include onnx"
255 f"\nPyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', '{best}')"
256 f"\nVisualize: https://netron.app\n")
257
258 # Plot examples
259 images, labels = (x[:25] for x in next(iter(testloader))) # first 25 images and labels
260 pred = torch.max(ema.ema(images.to(device)), 1)[1]
261 file = imshow_cls(images, labels, pred, model.names, verbose=False, f=save_dir / 'test_images.jpg')
262
263 # Log results
264 meta = {"epochs": epochs, "top1_acc": best_fitness, "date": datetime.now().isoformat()}
265 logger.log_images(file, name='Test Examples (true-predicted)', epoch=epoch)
266 logger.log_model(best, epochs, metadata=meta)
267
268
269 def parse_opt(known=False):
270 parser = argparse.ArgumentParser()
271 parser.add_argument('--model', type=str, default='yolov5s-cls.pt', help='initial weights path')
272 parser.add_argument('--data', type=str, default='imagenette160', help='cifar10, cifar100, mnist, imagenet, ...')
273 parser.add_argument('--epochs', type=int, default=10, help='total training epochs')
274 parser.add_argument('--batch-size', type=int, default=64, help='total batch size for all GPUs')
275 parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=224, help='train, val image size (pixels)')
276 parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
277 parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
278 parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
279 parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
280 parser.add_argument('--project', default=ROOT / 'runs/train-cls', help='save to project/name')
281 parser.add_argument('--name', default='exp', help='save to project/name')
282 parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
283 parser.add_argument('--pretrained', nargs='?', const=True, default=True, help='start from i.e. --pretrained False')
284 parser.add_argument('--optimizer', choices=['SGD', 'Adam', 'AdamW', 'RMSProp'], default='Adam', help='optimizer')
285 parser.add_argument('--lr0', type=float, default=0.001, help='initial learning rate')
286 parser.add_argument('--decay', type=float, default=5e-5, help='weight decay')
287 parser.add_argument('--label-smoothing', type=float, default=0.1, help='Label smoothing epsilon')
288 parser.add_argument('--cutoff', type=int, default=None, help='Model layer cutoff index for Classify() head')
289 parser.add_argument('--dropout', type=float, default=None, help='Dropout (fraction)')
290 parser.add_argument('--verbose', action='store_true', help='Verbose mode')
291 parser.add_argument('--seed', type=int, default=0, help='Global training seed')
292 parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
293 return parser.parse_known_args()[0] if known else parser.parse_args()
294
295
296 def main(opt):
297 # Checks
298 if RANK in {-1, 0}:
299 print_args(vars(opt))
300 check_git_status()
301 check_requirements()
302
303 # DDP mode
304 device = select_device(opt.device, batch_size=opt.batch_size)
305 if LOCAL_RANK != -1:
306 assert opt.batch_size != -1, 'AutoBatch is coming soon for classification, please pass a valid --batch-size'
307 assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
308 assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
309 torch.cuda.set_device(LOCAL_RANK)
310 device = torch.device('cuda', LOCAL_RANK)
311 dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
312
313 # Parameters
314 opt.save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok) # increment run
315
316 # Train
317 train(opt, device)
318
319
320 def run(**kwargs):
321 # Usage: from yolov5 import classify; classify.train.run(data=mnist, imgsz=320, model='yolov5m')
322 opt = parse_opt(True)
323 for k, v in kwargs.items():
324 setattr(opt, k, v)
325 main(opt)
326 return opt
327
328
329 if __name__ == "__main__":
330 opt = parse_opt()
331 main(opt)
classify/val.py deleted 100644 → 0
1 # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
2 """
3 Validate a trained YOLOv5 classification model on a classification dataset
4
5 Usage:
6 $ bash data/scripts/get_imagenet.sh --val # download ImageNet val split (6.3G, 50000 images)
7 $ python classify/val.py --weights yolov5m-cls.pt --data ../datasets/imagenet --img 224 # validate ImageNet
8
9 Usage - formats:
10 $ python classify/val.py --weights yolov5s-cls.pt # PyTorch
11 yolov5s-cls.torchscript # TorchScript
12 yolov5s-cls.onnx # ONNX Runtime or OpenCV DNN with --dnn
13 yolov5s-cls.xml # OpenVINO
14 yolov5s-cls.engine # TensorRT
15 yolov5s-cls.mlmodel # CoreML (macOS-only)
16 yolov5s-cls_saved_model # TensorFlow SavedModel
17 yolov5s-cls.pb # TensorFlow GraphDef
18 yolov5s-cls.tflite # TensorFlow Lite
19 yolov5s-cls_edgetpu.tflite # TensorFlow Edge TPU
20 yolov5s-cls_paddle_model # PaddlePaddle
21 """
22
23 import argparse
24 import os
25 import sys
26 from pathlib import Path
27
28 import torch
29 from tqdm import tqdm
30
31 FILE = Path(__file__).resolve()
32 ROOT = FILE.parents[1] # YOLOv5 root directory
33 if str(ROOT) not in sys.path:
34 sys.path.append(str(ROOT)) # add ROOT to PATH
35 ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
36
37 from models.common import DetectMultiBackend
38 from utils.dataloaders import create_classification_dataloader
39 from utils.general import LOGGER, Profile, check_img_size, check_requirements, colorstr, increment_path, print_args
40 from utils.torch_utils import select_device, smart_inference_mode
41
42
43 @smart_inference_mode()
44 def run(
45 data=ROOT / '../datasets/mnist', # dataset dir
46 weights=ROOT / 'yolov5s-cls.pt', # model.pt path(s)
47 batch_size=128, # batch size
48 imgsz=224, # inference size (pixels)
49 device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
50 workers=8, # max dataloader workers (per RANK in DDP mode)
51 verbose=False, # verbose output
52 project=ROOT / 'runs/val-cls', # save to project/name
53 name='exp', # save to project/name
54 exist_ok=False, # existing project/name ok, do not increment
55 half=False, # use FP16 half-precision inference
56 dnn=False, # use OpenCV DNN for ONNX inference
57 model=None,
58 dataloader=None,
59 criterion=None,
60 pbar=None,
61 ):
62 # Initialize/load model and set device
63 training = model is not None
64 if training: # called by train.py
65 device, pt, jit, engine = next(model.parameters()).device, True, False, False # get model device, PyTorch model
66 half &= device.type != 'cpu' # half precision only supported on CUDA
67 model.half() if half else model.float()
68 else: # called directly
69 device = select_device(device, batch_size=batch_size)
70
71 # Directories
72 save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
73 save_dir.mkdir(parents=True, exist_ok=True) # make dir
74
75 # Load model
76 model = DetectMultiBackend(weights, device=device, dnn=dnn, fp16=half)
77 stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
78 imgsz = check_img_size(imgsz, s=stride) # check image size
79 half = model.fp16 # FP16 supported on limited backends with CUDA
80 if engine:
81 batch_size = model.batch_size
82 else:
83 device = model.device
84 if not (pt or jit):
85 batch_size = 1 # export.py models default to batch-size 1
86 LOGGER.info(f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')
87
88 # Dataloader
89 data = Path(data)
90 test_dir = data / 'test' if (data / 'test').exists() else data / 'val' # data/test or data/val
91 dataloader = create_classification_dataloader(path=test_dir,
92 imgsz=imgsz,
93 batch_size=batch_size,
94 augment=False,
95 rank=-1,
96 workers=workers)
97
98 model.eval()
99 pred, targets, loss, dt = [], [], 0, (Profile(), Profile(), Profile())
100 n = len(dataloader) # number of batches
101 action = 'validating' if dataloader.dataset.root.stem == 'val' else 'testing'
102 desc = f"{pbar.desc[:-36]}{action:>36}" if pbar else f"{action}"
103 bar = tqdm(dataloader, desc, n, not training, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}', position=0)
104 with torch.cuda.amp.autocast(enabled=device.type != 'cpu'):
105 for images, labels in bar:
106 with dt[0]:
107 images, labels = images.to(device, non_blocking=True), labels.to(device)
108
109 with dt[1]:
110 y = model(images)
111
112 with dt[2]:
113 pred.append(y.argsort(1, descending=True)[:, :5])
114 targets.append(labels)
115 if criterion:
116 loss += criterion(y, labels)
117
118 loss /= n
119 pred, targets = torch.cat(pred), torch.cat(targets)
120 correct = (targets[:, None] == pred).float()
121 acc = torch.stack((correct[:, 0], correct.max(1).values), dim=1) # (top1, top5) accuracy
122 top1, top5 = acc.mean(0).tolist()
123
124 if pbar:
125 pbar.desc = f"{pbar.desc[:-36]}{loss:>12.3g}{top1:>12.3g}{top5:>12.3g}"
126 if verbose: # all classes
127 LOGGER.info(f"{'Class':>24}{'Images':>12}{'top1_acc':>12}{'top5_acc':>12}")
128 LOGGER.info(f"{'all':>24}{targets.shape[0]:>12}{top1:>12.3g}{top5:>12.3g}")
129 for i, c in model.names.items():
130 aci = acc[targets == i]
131 top1i, top5i = aci.mean(0).tolist()
132 LOGGER.info(f"{c:>24}{aci.shape[0]:>12}{top1i:>12.3g}{top5i:>12.3g}")
133
134 # Print results
135 t = tuple(x.t / len(dataloader.dataset.samples) * 1E3 for x in dt) # speeds per image
136 shape = (1, 3, imgsz, imgsz)
137 LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms post-process per image at shape {shape}' % t)
138 LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
139
140 return top1, top5, loss
141
142
143 def parse_opt():
144 parser = argparse.ArgumentParser()
145 parser.add_argument('--data', type=str, default=ROOT / '../datasets/mnist', help='dataset path')
146 parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-cls.pt', help='model.pt path(s)')
147 parser.add_argument('--batch-size', type=int, default=128, help='batch size')
148 parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=224, help='inference size (pixels)')
149 parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
150 parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
151 parser.add_argument('--verbose', nargs='?', const=True, default=True, help='verbose output')
152 parser.add_argument('--project', default=ROOT / 'runs/val-cls', help='save to project/name')
153 parser.add_argument('--name', default='exp', help='save to project/name')
154 parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
155 parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
156 parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
157 opt = parser.parse_args()
158 print_args(vars(opt))
159 return opt
160
161
162 def main(opt):
163 check_requirements(exclude=('tensorboard', 'thop'))
164 run(**vars(opt))
165
166
167 if __name__ == "__main__":
168 opt = parse_opt()
169 main(opt)
detect.py
...@@ -215,11 +215,11 @@ def run( ...@@ -215,11 +215,11 @@ def run(
215 215
216 def parse_opt(): 216 def parse_opt():
217 parser = argparse.ArgumentParser() 217 parser = argparse.ArgumentParser()
218 parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path or triton URL') 218 parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'runs/train/exp/weights/best.pt', help='model path or triton URL')
219 parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)') 219 parser.add_argument('--source', type=str, default=ROOT / 'data/images/crop_img', help='file/dir/URL/glob/screen/0(webcam)')
220 parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path') 220 parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
221 parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w') 221 parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
222 parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold') 222 parser.add_argument('--conf-thres', type=float, default=0.3, help='confidence threshold')
223 parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold') 223 parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
224 parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image') 224 parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
225 parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 225 parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
......
get_pr.py
...@@ -17,7 +17,8 @@ def iou(box, boxes): ...@@ -17,7 +17,8 @@ def iou(box, boxes):
17 inner = np.maximum(0, (xx2 - xx1) * (yy2 - yy1)) 17 inner = np.maximum(0, (xx2 - xx1) * (yy2 - yy1))
18 return inner / (area1 + areas - inner) 18 return inner / (area1 + areas - inner)
19 19
20 def get_evaluate_score(true_image_path,true_label_path,predict_label_path,threshold): 20
21 def get_evaluate_score(true_image_path, true_label_path, predict_label_path, threshold):
21 true_labels = os.listdir(true_label_path) 22 true_labels = os.listdir(true_label_path)
22 predict_labels = os.listdir(predict_label_path) 23 predict_labels = os.listdir(predict_label_path)
23 targets, predicts = [], [] 24 targets, predicts = [], []
...@@ -38,7 +39,7 @@ def get_evaluate_score(true_image_path,true_label_path,predict_label_path,thresh ...@@ -38,7 +39,7 @@ def get_evaluate_score(true_image_path,true_label_path,predict_label_path,thresh
38 predicts.append(0) 39 predicts.append(0)
39 else: 40 else:
40 tmp = 0 41 tmp = 0
41 predict_label = open(os.path.join(predict_label_path,label)).readlines() 42 predict_label = open(os.path.join(predict_label_path, label)).readlines()
42 boxes = [] 43 boxes = []
43 for pl in predict_label: 44 for pl in predict_label:
44 cls, x1, y1, w1, h1 = [float(i) for i in pl.strip().split(' ')] 45 cls, x1, y1, w1, h1 = [float(i) for i in pl.strip().split(' ')]
...@@ -50,23 +51,28 @@ def get_evaluate_score(true_image_path,true_label_path,predict_label_path,thresh ...@@ -50,23 +51,28 @@ def get_evaluate_score(true_image_path,true_label_path,predict_label_path,thresh
50 x1, y1, w1, h1 = int(x1 * w), int(y1 * h), int(w1 * w), int(h1 * h) 51 x1, y1, w1, h1 = int(x1 * w), int(y1 * h), int(w1 * w), int(h1 * h)
51 xx1, yy1, xx2, yy2 = x1 - w1 // 2, y1 - h1 // 2, x1 + w1 // 2, y1 + h1 // 2 52 xx1, yy1, xx2, yy2 = x1 - w1 // 2, y1 - h1 // 2, x1 + w1 // 2, y1 + h1 // 2
52 box1 = [xx1, yy1, xx2, yy2] 53 box1 = [xx1, yy1, xx2, yy2]
53 inner_score = iou(np.array(box1),np.array(boxes)) 54 inner_score = iou(np.array(box1), np.array(boxes))
54 if max(inner_score)>threshold: 55 if max(inner_score) > threshold:
55 tmp=1 56 tmp = 1
56 predicts.append(1) 57 predicts.append(1)
57 break 58 break
58 if tmp==0: 59 if tmp == 0:
59 predicts.append(0) 60 predicts.append(0)
60 p = precision_score(targets,predicts) 61 p = precision_score(targets, predicts)
61 r = recall_score(targets,predicts) 62 r = recall_score(targets, predicts)
62 conf = confusion_matrix(targets,predicts) 63 conf = confusion_matrix(targets, predicts)
63 print('precison:',p) 64 print('precison:', p)
64 print('recall:',r) 65 print('recall:', r)
65 print(conf) 66 print(conf)
66 67 print(f' 预 测 ')
68 print(f' authentic tampered ')
69 print(f'真 authentic \t\t{conf[0, 0]} \t\t{conf[0,1]}')
70 print(f'实 tempered \t\t{conf[1, 0]} \t\t\t{conf[1,1]}')
71 print(f'authentic precision:{conf[0,0]/(conf[0,0]+conf[1,0])}\trecall:{conf[0, 0]/(conf[0, 0]+conf[0, 1])}')
72 print(f'tampered precision:{conf[1, 1]/(conf[0, 1]+conf[1, 1])}\trecall:{conf[1, 1]/(conf[1, 0]+conf[1, 1])}')
67 if __name__ == '__main__': 73 if __name__ == '__main__':
68 true_image_path = '/home/qfs/WorkSpace/ps_tamper/yolov5_ps/val_data/minsheng/images' 74 true_image_path = '/data/situ_invoice_bill_data/qfs_train_val_data/gongshang/images/val'
69 true_label_path = '/home/qfs/WorkSpace/ps_tamper/yolov5_ps/val_data/minsheng/labels' 75 true_label_path = '/data/situ_invoice_bill_data/qfs_train_val_data/gongshang/labels/val'
70 predict_label_path = '/home/qfs/WorkSpace/ps_tamper/yolov5/runs/detect/exp/labels' 76 predict_label_path = '/home/situ/qfs/invoice_tamper/09_project/project/tamper_det/runs/detect/exp4/labels'
71 threshold = 0.1 77 threshold = 0.1
72 get_evaluate_score(true_image_path,true_label_path,predict_label_path,threshold)
...\ No newline at end of file ...\ No newline at end of file
78 get_evaluate_score(true_image_path, true_label_path, predict_label_path, threshold)
......
inference.py 0 → 100644
1 import copy
2 import os
3 import sys
4 from pathlib import Path
5 import numpy as np
6 import torch
7
8 from utils.augmentations import letterbox
9
10 FILE = Path(__file__).resolve()
11 ROOT = FILE.parents[0] # YOLOv5 root directory
12 if str(ROOT) not in sys.path:
13 sys.path.append(str(ROOT)) # add ROOT to PATH
14 ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
15 from models.common import DetectMultiBackend
16 from utils.general import (check_img_size, cv2, non_max_suppression, scale_boxes)
17 from utils.torch_utils import select_device, smart_inference_mode
18 from models.yolov5_config import config
19
20 classes = ['tampered']
21
22
23 def gen_result_dict(boxes, label_list=[], std=False):
24 result = {
25 "error_code": 1,
26 "result": []
27 }
28 rs_box = {
29 "class": '',
30 "score": 0,
31 "left": 0,
32 "top": 0,
33 "width": 0,
34 "height": 0
35 }
36
37 if not label_list:
38 label_list = classes
39
40 for box in boxes:
41 result['error_code'] = 0
42 box_dict = copy.deepcopy(rs_box)
43 if std:
44 box_dict['class'] = str(int(box[-1]))
45 else:
46 box_dict['class'] = label_list[int(box[-1])]
47
48 box_dict['left'] = int(round(box[0], 0))
49 box_dict['top'] = int(round(box[1], 0))
50 box_dict['width'] = int(round(box[2], 0) - round(box[0], 0))
51 box_dict['height'] = int(round(box[3], 0) - (round(box[1], 0)))
52 box_dict['score'] = box[-2]
53 result['result'].append(box_dict)
54 return result
55
56
57 class Yolov5:
58 def __init__(self, cfg=None):
59 self.cfg = cfg
60 self.device = select_device(self.cfg.device)
61 self.model = DetectMultiBackend(self.cfg.weights, device=self.device, dnn=False, data=self.cfg.data, fp16=False)
62
63 def detect(self, image):
64 image0 = image.copy()
65 stride, names, pt = self.model.stride, self.model.names, self.model.pt
66 imgsz = check_img_size(self.cfg.imgsz, s=stride) # check image size
67 # Dataloader
68 bs = 1 # batch_size
69 im = letterbox(image, imgsz, stride=stride, auto=True)[0] # padded resize
70 im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
71 im = np.ascontiguousarray(im) # contiguous
72 # Run inference
73 self.model.warmup(imgsz=(1 if pt or self.model.triton else bs, 3, *imgsz)) # warmup
74 im = torch.from_numpy(im).to(self.model.device)
75 im = im.half() if self.model.fp16 else im.float() # uint8 to fp16/32
76 im /= 255 # 0 - 255 to 0.0 - 1.0
77 if len(im.shape) == 3:
78 im = im[None] # expand for batch dim
79 # Inference
80 pred = self.model(im, augment=False, visualize=False)
81 # NMS
82 pred = non_max_suppression(pred, self.cfg.conf_thres, self.cfg.iou_thres, None, False, max_det=self.cfg.max_det)
83
84 det = pred[0]
85 # if len(det):
86 det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], image0.shape).round()
87 result = gen_result_dict(det.cpu().numpy().tolist())
88 return result
89
90 def plot(self, image, boxes):
91 for box in boxes:
92 cv2.rectangle(image, (box[0], box[1], box[2], box[3]), (0, 0, 255), 2)
93 return image
94
95
96 if __name__ == "__main__":
97 img = cv2.imread(
98 '/home/situ/qfs/invoice_tamper/09_project/project/yolov5_inference/data/images/crop_img/_1594890230.8032346page_10_img_0_hname.jpg')
99 detector = Yolov5(config)
100 result = detector.detect(img)
101 print(result)
models/__pycache__/__init__.cpython-36.pyc 0 → 100644
No preview for this file type
models/__pycache__/common.cpython-36.pyc 0 → 100644
No preview for this file type
models/__pycache__/experimental.cpython-36.pyc 0 → 100644
No preview for this file type
models/__pycache__/yolo.cpython-36.pyc 0 → 100644
No preview for this file type
models/yolov5_config.py 0 → 100644
1 from easydict import EasyDict as edict
2
3 config = edict(
4 weights='runs/train/exp/weights/best.pt', # model path or triton URL
5 data='data/VOC.yaml', # dataset.yaml path
6 imgsz=(640, 640), # inference size (height, width)
7 conf_thres=0.5, # confidence threshold
8 iou_thres=0.45, # NMS IOU threshold
9 max_det=1000, # maximum detections per image
10 device='' # cuda device, i.e. 0 or 0,1,2,3 or cpu
11 )
pipeline.py 0 → 100644
File mode changed
plot_sourece_labels.py 0 → 100644
1 import os
2
3 import cv2
4 import numpy as np
5 import pandas as pd
6 import tqdm
7
8
9 def get_source_image_det(crop_position, predict_positions):
10 result = []
11 x1, y1, x2, y2 = crop_position
12 for p in predict_positions:
13 px1, py1, px2, py2,score = p
14 w, h = px2 - px1, py2 - py1
15 result.append([x1 + px1, y1 + py1, x1 + px1 + w, y1 + py1 + h,score])
16 return result
17
18
19 def decode_label(image, label_path):
20 data = open(label_path).readlines()
21 h, w, c = image.shape
22 result = []
23 for d in data:
24 d = [float(i) for i in d.strip().split(' ')]
25 cls, cx, cy, cw, ch,score = d
26 cx, cy, cw, ch = cx * w, cy * h, cw * w, ch * h
27 result.append([int(cx - cw // 2), int(cy - ch // 2), int(cx + cw // 2), int(cy + ch // 2),score])
28 return result
29
30
31 if __name__ == '__main__':
32 source_image_path = '/data/situ_invoice_bill_data/new_data/qfs_bank_bill_data/gongshang/authentic/images/val'
33 val_image_path = '/home/situ/qfs/invoice_tamper/09_project/project/tamper_det/data/images/crop_img'
34 predict_label_path = '/home/situ/qfs/invoice_tamper/09_project/project/tamper_det/runs/detect/exp/labels'
35 crop_csv_path = '/data/situ_invoice_bill_data/new_data/qfs_bank_bill_data/gongshang/croped_merge.csv'
36 predict_labels = os.listdir(predict_label_path)
37 source_images = os.listdir(source_image_path)
38 data = pd.read_csv(crop_csv_path)
39 img_name = data.loc[:, 'img_name'].tolist()
40 crop_position1 = data.loc[:, 'name_crop_coord'].tolist()
41 crop_position2 = data.loc[:,'number_crop_coord'].tolist()
42 cc='/data/situ_invoice_bill_data/new_data/qfs_bank_bill_data/gongshang/tampered/images/val/ps3'
43 for im in os.listdir(cc):
44 print(im)
45 img = cv2.imread(os.path.join(cc,im))
46 img_=img.copy()
47 id = img_name.index(im)
48 name_crop_position=[int(i) for i in crop_position1[id].split(',')]
49 number_crop_position=[int(i) for i in crop_position2[id].split(',')]
50 nx1,ny1,nx2,ny2=name_crop_position
51 nux1,nuy1,nux2,nuy2=number_crop_position
52 if im[:-4]+'_hname.txt' in predict_labels:
53
54 h, w, c = img[ny1:ny2, nx1:nx2, :].shape
55 data = open(os.path.join(predict_label_path,im[:-4]+'_hname.txt')).readlines()
56 for d in data:
57 cls,cx,cy,cw,ch,score = [float(i) for i in d.strip().split(' ')]
58 cx,cy,cw,ch=int(cx*w),int(cy*h),int(cw*w),int(ch*h)
59 cx1,cy1=cx-cw//2,cy-ch//2
60 x1,y1,x2,y2=nx1+cx1,ny1+cy1,nx1+cx1+cw,ny1+cy1+ch
61 cv2.rectangle(img,(x1,y1),(x2,y2),(0,0,255),2)
62 cv2.putText(img,f'tampered:{score}',(x1,y1-5),cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,0,255),1)
63 if im[:-4] + '_hnumber.txt' in predict_labels:
64 h, w, c = img[nuy1:nuy2, nux1:nux2, :].shape
65 data = open(os.path.join(predict_label_path, im[:-4] + '_hname.txt')).readlines()
66 for d in data:
67 cls, cx, cy, cw, ch, score = [float(i) for i in d.strip().split(' ')]
68 cx, cy, cw, ch = int(cx * w), int(cy * h), int(cw * w), int(ch * h)
69 cx1, cy1 = cx - cw // 2, cy - ch // 2
70 x1, y1, x2, y2 = nux1 + cx1, nuy1 + cy1, nux1 + cx1 + cw, nuy1 + cy1 + ch
71 cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
72 cv2.putText(img, f'tampered:{score}', (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
73 result = np.vstack((img_,img))
74 cv2.imwrite(f'z/{im}',result)
segment/predict.py deleted 100644 → 0
1 # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
2 """
3 Run YOLOv5 segmentation inference on images, videos, directories, streams, etc.
4
5 Usage - sources:
6 $ python segment/predict.py --weights yolov5s-seg.pt --source 0 # webcam
7 img.jpg # image
8 vid.mp4 # video
9 path/ # directory
10 'path/*.jpg' # glob
11 'https://youtu.be/Zgi9g1ksQHc' # YouTube
12 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
13
14 Usage - formats:
15 $ python segment/predict.py --weights yolov5s-seg.pt # PyTorch
16 yolov5s-seg.torchscript # TorchScript
17 yolov5s-seg.onnx # ONNX Runtime or OpenCV DNN with --dnn
18 yolov5s-seg.xml # OpenVINO
19 yolov5s-seg.engine # TensorRT
20 yolov5s-seg.mlmodel # CoreML (macOS-only)
21 yolov5s-seg_saved_model # TensorFlow SavedModel
22 yolov5s-seg.pb # TensorFlow GraphDef
23 yolov5s-seg.tflite # TensorFlow Lite
24 yolov5s-seg_edgetpu.tflite # TensorFlow Edge TPU
25 yolov5s-seg_paddle_model # PaddlePaddle
26 """
27
28 import argparse
29 import os
30 import platform
31 import sys
32 from pathlib import Path
33
34 import torch
35
36 FILE = Path(__file__).resolve()
37 ROOT = FILE.parents[1] # YOLOv5 root directory
38 if str(ROOT) not in sys.path:
39 sys.path.append(str(ROOT)) # add ROOT to PATH
40 ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
41
42 from models.common import DetectMultiBackend
43 from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
44 from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
45 increment_path, non_max_suppression, print_args, scale_boxes, scale_segments,
46 strip_optimizer, xyxy2xywh)
47 from utils.plots import Annotator, colors, save_one_box
48 from utils.segment.general import masks2segments, process_mask
49 from utils.torch_utils import select_device, smart_inference_mode
50
51
52 @smart_inference_mode()
53 def run(
54 weights=ROOT / 'yolov5s-seg.pt', # model.pt path(s)
55 source=ROOT / 'data/images', # file/dir/URL/glob/screen/0(webcam)
56 data=ROOT / 'data/coco128.yaml', # dataset.yaml path
57 imgsz=(640, 640), # inference size (height, width)
58 conf_thres=0.25, # confidence threshold
59 iou_thres=0.45, # NMS IOU threshold
60 max_det=1000, # maximum detections per image
61 device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
62 view_img=False, # show results
63 save_txt=False, # save results to *.txt
64 save_conf=False, # save confidences in --save-txt labels
65 save_crop=False, # save cropped prediction boxes
66 nosave=False, # do not save images/videos
67 classes=None, # filter by class: --class 0, or --class 0 2 3
68 agnostic_nms=False, # class-agnostic NMS
69 augment=False, # augmented inference
70 visualize=False, # visualize features
71 update=False, # update all models
72 project=ROOT / 'runs/predict-seg', # save results to project/name
73 name='exp', # save results to project/name
74 exist_ok=False, # existing project/name ok, do not increment
75 line_thickness=3, # bounding box thickness (pixels)
76 hide_labels=False, # hide labels
77 hide_conf=False, # hide confidences
78 half=False, # use FP16 half-precision inference
79 dnn=False, # use OpenCV DNN for ONNX inference
80 vid_stride=1, # video frame-rate stride
81 retina_masks=False,
82 ):
83 source = str(source)
84 save_img = not nosave and not source.endswith('.txt') # save inference images
85 is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
86 is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
87 webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
88 screenshot = source.lower().startswith('screen')
89 if is_url and is_file:
90 source = check_file(source) # download
91
92 # Directories
93 save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
94 (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
95
96 # Load model
97 device = select_device(device)
98 model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
99 stride, names, pt = model.stride, model.names, model.pt
100 imgsz = check_img_size(imgsz, s=stride) # check image size
101
102 # Dataloader
103 bs = 1 # batch_size
104 if webcam:
105 view_img = check_imshow()
106 dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
107 bs = len(dataset)
108 elif screenshot:
109 dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
110 else:
111 dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
112 vid_path, vid_writer = [None] * bs, [None] * bs
113
114 # Run inference
115 model.warmup(imgsz=(1 if pt else bs, 3, *imgsz)) # warmup
116 seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
117 for path, im, im0s, vid_cap, s in dataset:
118 with dt[0]:
119 im = torch.from_numpy(im).to(model.device)
120 im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
121 im /= 255 # 0 - 255 to 0.0 - 1.0
122 if len(im.shape) == 3:
123 im = im[None] # expand for batch dim
124
125 # Inference
126 with dt[1]:
127 visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
128 pred, proto = model(im, augment=augment, visualize=visualize)[:2]
129
130 # NMS
131 with dt[2]:
132 pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det, nm=32)
133
134 # Second-stage classifier (optional)
135 # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
136
137 # Process predictions
138 for i, det in enumerate(pred): # per image
139 seen += 1
140 if webcam: # batch_size >= 1
141 p, im0, frame = path[i], im0s[i].copy(), dataset.count
142 s += f'{i}: '
143 else:
144 p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
145
146 p = Path(p) # to Path
147 save_path = str(save_dir / p.name) # im.jpg
148 txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # im.txt
149 s += '%gx%g ' % im.shape[2:] # print string
150 imc = im0.copy() if save_crop else im0 # for save_crop
151 annotator = Annotator(im0, line_width=line_thickness, example=str(names))
152 if len(det):
153 masks = process_mask(proto[i], det[:, 6:], det[:, :4], im.shape[2:], upsample=True) # HWC
154 det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round() # rescale boxes to im0 size
155
156 # Segments
157 if save_txt:
158 segments = reversed(masks2segments(masks))
159 segments = [scale_segments(im.shape[2:], x, im0.shape).round() for x in segments]
160
161 # Print results
162 for c in det[:, 5].unique():
163 n = (det[:, 5] == c).sum() # detections per class
164 s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
165
166 # Mask plotting
167 annotator.masks(masks,
168 colors=[colors(x, True) for x in det[:, 5]],
169 im_gpu=None if retina_masks else im[i])
170
171 # Write results
172 for j, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])):
173 if save_txt: # Write to file
174 segj = segments[j].reshape(-1) # (n,2) to (n*2)
175 line = (cls, *segj, conf) if save_conf else (cls, *segj) # label format
176 with open(f'{txt_path}.txt', 'a') as f:
177 f.write(('%g ' * len(line)).rstrip() % line + '\n')
178
179 if save_img or save_crop or view_img: # Add bbox to image
180 c = int(cls) # integer class
181 label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
182 annotator.box_label(xyxy, label, color=colors(c, True))
183 # annotator.draw.polygon(segments[j], outline=colors(c, True), width=3)
184 if save_crop:
185 save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
186
187 # Stream results
188 im0 = annotator.result()
189 if view_img:
190 if platform.system() == 'Linux' and p not in windows:
191 windows.append(p)
192 cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO) # allow window resize (Linux)
193 cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
194 cv2.imshow(str(p), im0)
195 if cv2.waitKey(1) == ord('q'): # 1 millisecond
196 exit()
197
198 # Save results (image with detections)
199 if save_img:
200 if dataset.mode == 'image':
201 cv2.imwrite(save_path, im0)
202 else: # 'video' or 'stream'
203 if vid_path[i] != save_path: # new video
204 vid_path[i] = save_path
205 if isinstance(vid_writer[i], cv2.VideoWriter):
206 vid_writer[i].release() # release previous video writer
207 if vid_cap: # video
208 fps = vid_cap.get(cv2.CAP_PROP_FPS)
209 w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
210 h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
211 else: # stream
212 fps, w, h = 30, im0.shape[1], im0.shape[0]
213 save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos
214 vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
215 vid_writer[i].write(im0)
216
217 # Print time (inference-only)
218 LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
219
220 # Print results
221 t = tuple(x.t / seen * 1E3 for x in dt) # speeds per image
222 LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
223 if save_txt or save_img:
224 s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
225 LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
226 if update:
227 strip_optimizer(weights[0]) # update model (to fix SourceChangeWarning)
228
229
230 def parse_opt():
231 parser = argparse.ArgumentParser()
232 parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-seg.pt', help='model path(s)')
233 parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)')
234 parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
235 parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
236 parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')
237 parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
238 parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
239 parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
240 parser.add_argument('--view-img', action='store_true', help='show results')
241 parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
242 parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
243 parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
244 parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
245 parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
246 parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
247 parser.add_argument('--augment', action='store_true', help='augmented inference')
248 parser.add_argument('--visualize', action='store_true', help='visualize features')
249 parser.add_argument('--update', action='store_true', help='update all models')
250 parser.add_argument('--project', default=ROOT / 'runs/predict-seg', help='save results to project/name')
251 parser.add_argument('--name', default='exp', help='save results to project/name')
252 parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
253 parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
254 parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
255 parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
256 parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
257 parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
258 parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
259 parser.add_argument('--retina-masks', action='store_true', help='whether to plot masks in native resolution')
260 opt = parser.parse_args()
261 opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
262 print_args(vars(opt))
263 return opt
264
265
266 def main(opt):
267 check_requirements(exclude=('tensorboard', 'thop'))
268 run(**vars(opt))
269
270
271 if __name__ == "__main__":
272 opt = parse_opt()
273 main(opt)
segment/train.py deleted 100644 → 0
1 # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
2 """
3 Train a YOLOv5 segment model on a segment dataset
4 Models and datasets download automatically from the latest YOLOv5 release.
5
6 Usage - Single-GPU training:
7 $ python segment/train.py --data coco128-seg.yaml --weights yolov5s-seg.pt --img 640 # from pretrained (recommended)
8 $ python segment/train.py --data coco128-seg.yaml --weights '' --cfg yolov5s-seg.yaml --img 640 # from scratch
9
10 Usage - Multi-GPU DDP training:
11 $ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 segment/train.py --data coco128-seg.yaml --weights yolov5s-seg.pt --img 640 --device 0,1,2,3
12
13 Models: https://github.com/ultralytics/yolov5/tree/master/models
14 Datasets: https://github.com/ultralytics/yolov5/tree/master/data
15 Tutorial: https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data
16 """
17
18 import argparse
19 import math
20 import os
21 import random
22 import sys
23 import time
24 from copy import deepcopy
25 from datetime import datetime
26 from pathlib import Path
27
28 import numpy as np
29 import torch
30 import torch.distributed as dist
31 import torch.nn as nn
32 import yaml
33 from torch.optim import lr_scheduler
34 from tqdm import tqdm
35
36 FILE = Path(__file__).resolve()
37 ROOT = FILE.parents[1] # YOLOv5 root directory
38 if str(ROOT) not in sys.path:
39 sys.path.append(str(ROOT)) # add ROOT to PATH
40 ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
41
42 import segment.val as validate # for end-of-epoch mAP
43 from models.experimental import attempt_load
44 from models.yolo import SegmentationModel
45 from utils.autoanchor import check_anchors
46 from utils.autobatch import check_train_batch_size
47 from utils.callbacks import Callbacks
48 from utils.downloads import attempt_download, is_url
49 from utils.general import (LOGGER, check_amp, check_dataset, check_file, check_git_status, check_img_size,
50 check_requirements, check_suffix, check_yaml, colorstr, get_latest_run, increment_path,
51 init_seeds, intersect_dicts, labels_to_class_weights, labels_to_image_weights, one_cycle,
52 print_args, print_mutation, strip_optimizer, yaml_save)
53 from utils.loggers import GenericLogger
54 from utils.plots import plot_evolve, plot_labels
55 from utils.segment.dataloaders import create_dataloader
56 from utils.segment.loss import ComputeLoss
57 from utils.segment.metrics import KEYS, fitness
58 from utils.segment.plots import plot_images_and_masks, plot_results_with_masks
59 from utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, select_device, smart_DDP, smart_optimizer,
60 smart_resume, torch_distributed_zero_first)
61
62 LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html
63 RANK = int(os.getenv('RANK', -1))
64 WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
65
66
67 def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary
68 save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze, mask_ratio = \
69 Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
70 opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze, opt.mask_ratio
71 # callbacks.run('on_pretrain_routine_start')
72
73 # Directories
74 w = save_dir / 'weights' # weights dir
75 (w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir
76 last, best = w / 'last.pt', w / 'best.pt'
77
78 # Hyperparameters
79 if isinstance(hyp, str):
80 with open(hyp, errors='ignore') as f:
81 hyp = yaml.safe_load(f) # load hyps dict
82 LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
83 opt.hyp = hyp.copy() # for saving hyps to checkpoints
84
85 # Save run settings
86 if not evolve:
87 yaml_save(save_dir / 'hyp.yaml', hyp)
88 yaml_save(save_dir / 'opt.yaml', vars(opt))
89
90 # Loggers
91 data_dict = None
92 if RANK in {-1, 0}:
93 logger = GenericLogger(opt=opt, console_logger=LOGGER)
94 # loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance
95 # if loggers.clearml:
96 # data_dict = loggers.clearml.data_dict # None if no ClearML dataset or filled in by ClearML
97 # if loggers.wandb:
98 # data_dict = loggers.wandb.data_dict
99 # if resume:
100 # weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
101 #
102 # # Register actions
103 # for k in methods(loggers):
104 # callbacks.register_action(k, callback=getattr(loggers, k))
105
106 # Config
107 plots = not evolve and not opt.noplots # create plots
108 overlap = not opt.no_overlap
109 cuda = device.type != 'cpu'
110 init_seeds(opt.seed + 1 + RANK, deterministic=True)
111 with torch_distributed_zero_first(LOCAL_RANK):
112 data_dict = data_dict or check_dataset(data) # check if None
113 train_path, val_path = data_dict['train'], data_dict['val']
114 nc = 1 if single_cls else int(data_dict['nc']) # number of classes
115 names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
116 is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset
117
118 # Model
119 check_suffix(weights, '.pt') # check weights
120 pretrained = weights.endswith('.pt')
121 if pretrained:
122 with torch_distributed_zero_first(LOCAL_RANK):
123 weights = attempt_download(weights) # download if not found locally
124 ckpt = torch.load(weights, map_location='cpu') # load checkpoint to CPU to avoid CUDA memory leak
125 model = SegmentationModel(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)
126 exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys
127 csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
128 csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect
129 model.load_state_dict(csd, strict=False) # load
130 LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # report
131 else:
132 model = SegmentationModel(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
133 amp = check_amp(model) # check AMP
134
135 # Freeze
136 freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freeze
137 for k, v in model.named_parameters():
138 v.requires_grad = True # train all layers
139 # v.register_hook(lambda x: torch.nan_to_num(x)) # NaN to 0 (commented for erratic training results)
140 if any(x in k for x in freeze):
141 LOGGER.info(f'freezing {k}')
142 v.requires_grad = False
143
144 # Image size
145 gs = max(int(model.stride.max()), 32) # grid size (max stride)
146 imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple
147
148 # Batch size
149 if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch size
150 batch_size = check_train_batch_size(model, imgsz, amp)
151 logger.update_params({"batch_size": batch_size})
152 # loggers.on_params_update({"batch_size": batch_size})
153
154 # Optimizer
155 nbs = 64 # nominal batch size
156 accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing
157 hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay
158 optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
159
160 # Scheduler
161 if opt.cos_lr:
162 lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
163 else:
164 lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
165 scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
166
167 # EMA
168 ema = ModelEMA(model) if RANK in {-1, 0} else None
169
170 # Resume
171 best_fitness, start_epoch = 0.0, 0
172 if pretrained:
173 if resume:
174 best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
175 del ckpt, csd
176
177 # DP mode
178 if cuda and RANK == -1 and torch.cuda.device_count() > 1:
179 LOGGER.warning('WARNING ⚠️ DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n'
180 'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.')
181 model = torch.nn.DataParallel(model)
182
183 # SyncBatchNorm
184 if opt.sync_bn and cuda and RANK != -1:
185 model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
186 LOGGER.info('Using SyncBatchNorm()')
187
188 # Trainloader
189 train_loader, dataset = create_dataloader(
190 train_path,
191 imgsz,
192 batch_size // WORLD_SIZE,
193 gs,
194 single_cls,
195 hyp=hyp,
196 augment=True,
197 cache=None if opt.cache == 'val' else opt.cache,
198 rect=opt.rect,
199 rank=LOCAL_RANK,
200 workers=workers,
201 image_weights=opt.image_weights,
202 quad=opt.quad,
203 prefix=colorstr('train: '),
204 shuffle=True,
205 mask_downsample_ratio=mask_ratio,
206 overlap_mask=overlap,
207 )
208 labels = np.concatenate(dataset.labels, 0)
209 mlc = int(labels[:, 0].max()) # max label class
210 assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
211
212 # Process 0
213 if RANK in {-1, 0}:
214 val_loader = create_dataloader(val_path,
215 imgsz,
216 batch_size // WORLD_SIZE * 2,
217 gs,
218 single_cls,
219 hyp=hyp,
220 cache=None if noval else opt.cache,
221 rect=True,
222 rank=-1,
223 workers=workers * 2,
224 pad=0.5,
225 mask_downsample_ratio=mask_ratio,
226 overlap_mask=overlap,
227 prefix=colorstr('val: '))[0]
228
229 if not resume:
230 if not opt.noautoanchor:
231 check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) # run AutoAnchor
232 model.half().float() # pre-reduce anchor precision
233
234 if plots:
235 plot_labels(labels, names, save_dir)
236 # callbacks.run('on_pretrain_routine_end', labels, names)
237
238 # DDP mode
239 if cuda and RANK != -1:
240 model = smart_DDP(model)
241
242 # Model attributes
243 nl = de_parallel(model).model[-1].nl # number of detection layers (to scale hyps)
244 hyp['box'] *= 3 / nl # scale to layers
245 hyp['cls'] *= nc / 80 * 3 / nl # scale to classes and layers
246 hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layers
247 hyp['label_smoothing'] = opt.label_smoothing
248 model.nc = nc # attach number of classes to model
249 model.hyp = hyp # attach hyperparameters to model
250 model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights
251 model.names = names
252
253 # Start training
254 t0 = time.time()
255 nb = len(train_loader) # number of batches
256 nw = max(round(hyp['warmup_epochs'] * nb), 100) # number of warmup iterations, max(3 epochs, 100 iterations)
257 # nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
258 last_opt_step = -1
259 maps = np.zeros(nc) # mAP per class
260 results = (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
261 scheduler.last_epoch = start_epoch - 1 # do not move
262 scaler = torch.cuda.amp.GradScaler(enabled=amp)
263 stopper, stop = EarlyStopping(patience=opt.patience), False
264 compute_loss = ComputeLoss(model, overlap=overlap) # init loss class
265 # callbacks.run('on_train_start')
266 LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
267 f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'
268 f"Logging results to {colorstr('bold', save_dir)}\n"
269 f'Starting training for {epochs} epochs...')
270 for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
271 # callbacks.run('on_train_epoch_start')
272 model.train()
273
274 # Update image weights (optional, single-GPU only)
275 if opt.image_weights:
276 cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights
277 iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights
278 dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx
279
280 # Update mosaic border (optional)
281 # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
282 # dataset.mosaic_border = [b - imgsz, -b] # height, width borders
283
284 mloss = torch.zeros(4, device=device) # mean losses
285 if RANK != -1:
286 train_loader.sampler.set_epoch(epoch)
287 pbar = enumerate(train_loader)
288 LOGGER.info(('\n' + '%11s' * 8) %
289 ('Epoch', 'GPU_mem', 'box_loss', 'seg_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size'))
290 if RANK in {-1, 0}:
291 pbar = tqdm(pbar, total=nb, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
292 optimizer.zero_grad()
293 for i, (imgs, targets, paths, _, masks) in pbar: # batch ------------------------------------------------------
294 # callbacks.run('on_train_batch_start')
295 ni = i + nb * epoch # number integrated batches (since train start)
296 imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0
297
298 # Warmup
299 if ni <= nw:
300 xi = [0, nw] # x interp
301 # compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
302 accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
303 for j, x in enumerate(optimizer.param_groups):
304 # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
305 x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
306 if 'momentum' in x:
307 x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
308
309 # Multi-scale
310 if opt.multi_scale:
311 sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
312 sf = sz / max(imgs.shape[2:]) # scale factor
313 if sf != 1:
314 ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
315 imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
316
317 # Forward
318 with torch.cuda.amp.autocast(amp):
319 pred = model(imgs) # forward
320 loss, loss_items = compute_loss(pred, targets.to(device), masks=masks.to(device).float())
321 if RANK != -1:
322 loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
323 if opt.quad:
324 loss *= 4.
325
326 # Backward
327 scaler.scale(loss).backward()
328
329 # Optimize - https://pytorch.org/docs/master/notes/amp_examples.html
330 if ni - last_opt_step >= accumulate:
331 scaler.unscale_(optimizer) # unscale gradients
332 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradients
333 scaler.step(optimizer) # optimizer.step
334 scaler.update()
335 optimizer.zero_grad()
336 if ema:
337 ema.update(model)
338 last_opt_step = ni
339
340 # Log
341 if RANK in {-1, 0}:
342 mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
343 mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)
344 pbar.set_description(('%11s' * 2 + '%11.4g' * 6) %
345 (f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
346 # callbacks.run('on_train_batch_end', model, ni, imgs, targets, paths)
347 # if callbacks.stop_training:
348 # return
349
350 # Mosaic plots
351 if plots:
352 if ni < 3:
353 plot_images_and_masks(imgs, targets, masks, paths, save_dir / f"train_batch{ni}.jpg")
354 if ni == 10:
355 files = sorted(save_dir.glob('train*.jpg'))
356 logger.log_images(files, "Mosaics", epoch)
357 # end batch ------------------------------------------------------------------------------------------------
358
359 # Scheduler
360 lr = [x['lr'] for x in optimizer.param_groups] # for loggers
361 scheduler.step()
362
363 if RANK in {-1, 0}:
364 # mAP
365 # callbacks.run('on_train_epoch_end', epoch=epoch)
366 ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])
367 final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
368 if not noval or final_epoch: # Calculate mAP
369 results, maps, _ = validate.run(data_dict,
370 batch_size=batch_size // WORLD_SIZE * 2,
371 imgsz=imgsz,
372 half=amp,
373 model=ema.ema,
374 single_cls=single_cls,
375 dataloader=val_loader,
376 save_dir=save_dir,
377 plots=False,
378 callbacks=callbacks,
379 compute_loss=compute_loss,
380 mask_downsample_ratio=mask_ratio,
381 overlap=overlap)
382
383 # Update best mAP
384 fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
385 stop = stopper(epoch=epoch, fitness=fi) # early stop check
386 if fi > best_fitness:
387 best_fitness = fi
388 log_vals = list(mloss) + list(results) + lr
389 # callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
390 # Log val metrics and media
391 metrics_dict = dict(zip(KEYS, log_vals))
392 logger.log_metrics(metrics_dict, epoch)
393
394 # Save model
395 if (not nosave) or (final_epoch and not evolve): # if save
396 ckpt = {
397 'epoch': epoch,
398 'best_fitness': best_fitness,
399 'model': deepcopy(de_parallel(model)).half(),
400 'ema': deepcopy(ema.ema).half(),
401 'updates': ema.updates,
402 'optimizer': optimizer.state_dict(),
403 # 'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None,
404 'opt': vars(opt),
405 'date': datetime.now().isoformat()}
406
407 # Save last, best and delete
408 torch.save(ckpt, last)
409 if best_fitness == fi:
410 torch.save(ckpt, best)
411 if opt.save_period > 0 and epoch % opt.save_period == 0:
412 torch.save(ckpt, w / f'epoch{epoch}.pt')
413 logger.log_model(w / f'epoch{epoch}.pt')
414 del ckpt
415 # callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
416
417 # EarlyStopping
418 if RANK != -1: # if DDP training
419 broadcast_list = [stop if RANK == 0 else None]
420 dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks
421 if RANK != 0:
422 stop = broadcast_list[0]
423 if stop:
424 break # must break all DDP ranks
425
426 # end epoch ----------------------------------------------------------------------------------------------------
427 # end training -----------------------------------------------------------------------------------------------------
428 if RANK in {-1, 0}:
429 LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
430 for f in last, best:
431 if f.exists():
432 strip_optimizer(f) # strip optimizers
433 if f is best:
434 LOGGER.info(f'\nValidating {f}...')
435 results, _, _ = validate.run(
436 data_dict,
437 batch_size=batch_size // WORLD_SIZE * 2,
438 imgsz=imgsz,
439 model=attempt_load(f, device).half(),
440 iou_thres=0.65 if is_coco else 0.60, # best pycocotools at iou 0.65
441 single_cls=single_cls,
442 dataloader=val_loader,
443 save_dir=save_dir,
444 save_json=is_coco,
445 verbose=True,
446 plots=plots,
447 callbacks=callbacks,
448 compute_loss=compute_loss,
449 mask_downsample_ratio=mask_ratio,
450 overlap=overlap) # val best model with plots
451 if is_coco:
452 # callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
453 metrics_dict = dict(zip(KEYS, list(mloss) + list(results) + lr))
454 logger.log_metrics(metrics_dict, epoch)
455
456 # callbacks.run('on_train_end', last, best, epoch, results)
457 # on train end callback using genericLogger
458 logger.log_metrics(dict(zip(KEYS[4:16], results)), epochs)
459 if not opt.evolve:
460 logger.log_model(best, epoch)
461 if plots:
462 plot_results_with_masks(file=save_dir / 'results.csv') # save results.png
463 files = ['results.png', 'confusion_matrix.png', *(f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R'))]
464 files = [(save_dir / f) for f in files if (save_dir / f).exists()] # filter
465 LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
466 logger.log_images(files, "Results", epoch + 1)
467 logger.log_images(sorted(save_dir.glob('val*.jpg')), "Validation", epoch + 1)
468 torch.cuda.empty_cache()
469 return results
470
471
472 def parse_opt(known=False):
473 parser = argparse.ArgumentParser()
474 parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s-seg.pt', help='initial weights path')
475 parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
476 parser.add_argument('--data', type=str, default=ROOT / 'data/coco128-seg.yaml', help='dataset.yaml path')
477 parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
478 parser.add_argument('--epochs', type=int, default=300, help='total training epochs')
479 parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')
480 parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
481 parser.add_argument('--rect', action='store_true', help='rectangular training')
482 parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
483 parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
484 parser.add_argument('--noval', action='store_true', help='only validate final epoch')
485 parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
486 parser.add_argument('--noplots', action='store_true', help='save no plot files')
487 parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
488 parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
489 parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
490 parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
491 parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
492 parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
493 parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
494 parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
495 parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
496 parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
497 parser.add_argument('--project', default=ROOT / 'runs/train-seg', help='save to project/name')
498 parser.add_argument('--name', default='exp', help='save to project/name')
499 parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
500 parser.add_argument('--quad', action='store_true', help='quad dataloader')
501 parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
502 parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
503 parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
504 parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
505 parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
506 parser.add_argument('--seed', type=int, default=0, help='Global training seed')
507 parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
508
509 # Instance Segmentation Args
510 parser.add_argument('--mask-ratio', type=int, default=4, help='Downsample the truth masks to saving memory')
511 parser.add_argument('--no-overlap', action='store_true', help='Overlap masks train faster at slightly less mAP')
512
513 # Weights & Biases arguments
514 # parser.add_argument('--entity', default=None, help='W&B: Entity')
515 # parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='W&B: Upload data, "val" option')
516 # parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval')
517 # parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use')
518
519 return parser.parse_known_args()[0] if known else parser.parse_args()
520
521
522 def main(opt, callbacks=Callbacks()):
523 # Checks
524 if RANK in {-1, 0}:
525 print_args(vars(opt))
526 check_git_status()
527 check_requirements()
528
529 # Resume
530 if opt.resume and not opt.evolve: # resume from specified or most recent last.pt
531 last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
532 opt_yaml = last.parent.parent / 'opt.yaml' # train options yaml
533 opt_data = opt.data # original dataset
534 if opt_yaml.is_file():
535 with open(opt_yaml, errors='ignore') as f:
536 d = yaml.safe_load(f)
537 else:
538 d = torch.load(last, map_location='cpu')['opt']
539 opt = argparse.Namespace(**d) # replace
540 opt.cfg, opt.weights, opt.resume = '', str(last), True # reinstate
541 if is_url(opt_data):
542 opt.data = check_file(opt_data) # avoid HUB resume auth timeout
543 else:
544 opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \
545 check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project) # checks
546 assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
547 if opt.evolve:
548 if opt.project == str(ROOT / 'runs/train'): # if default project name, rename to runs/evolve
549 opt.project = str(ROOT / 'runs/evolve')
550 opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume
551 if opt.name == 'cfg':
552 opt.name = Path(opt.cfg).stem # use model.yaml as name
553 opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
554
555 # DDP mode
556 device = select_device(opt.device, batch_size=opt.batch_size)
557 if LOCAL_RANK != -1:
558 msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'
559 assert not opt.image_weights, f'--image-weights {msg}'
560 assert not opt.evolve, f'--evolve {msg}'
561 assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'
562 assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
563 assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
564 torch.cuda.set_device(LOCAL_RANK)
565 device = torch.device('cuda', LOCAL_RANK)
566 dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
567
568 # Train
569 if not opt.evolve:
570 train(opt.hyp, opt, device, callbacks)
571
572 # Evolve hyperparameters (optional)
573 else:
574 # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
575 meta = {
576 'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)
577 'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)
578 'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1
579 'weight_decay': (1, 0.0, 0.001), # optimizer weight decay
580 'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok)
581 'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum
582 'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr
583 'box': (1, 0.02, 0.2), # box loss gain
584 'cls': (1, 0.2, 4.0), # cls loss gain
585 'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight
586 'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels)
587 'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight
588 'iou_t': (0, 0.1, 0.7), # IoU training threshold
589 'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold
590 'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore)
591 'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)
592 'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction)
593 'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction)
594 'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction)
595 'degrees': (1, 0.0, 45.0), # image rotation (+/- deg)
596 'translate': (1, 0.0, 0.9), # image translation (+/- fraction)
597 'scale': (1, 0.0, 0.9), # image scale (+/- gain)
598 'shear': (1, 0.0, 10.0), # image shear (+/- deg)
599 'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001
600 'flipud': (1, 0.0, 1.0), # image flip up-down (probability)
601 'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)
602 'mosaic': (1, 0.0, 1.0), # image mixup (probability)
603 'mixup': (1, 0.0, 1.0), # image mixup (probability)
604 'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability)
605
606 with open(opt.hyp, errors='ignore') as f:
607 hyp = yaml.safe_load(f) # load hyps dict
608 if 'anchors' not in hyp: # anchors commented in hyp.yaml
609 hyp['anchors'] = 3
610 if opt.noautoanchor:
611 del hyp['anchors'], meta['anchors']
612 opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # only val/save final epoch
613 # ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices
614 evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
615 if opt.bucket:
616 os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {evolve_csv}') # download evolve.csv if exists
617
618 for _ in range(opt.evolve): # generations to evolve
619 if evolve_csv.exists(): # if evolve.csv exists: select best hyps and mutate
620 # Select parent(s)
621 parent = 'single' # parent selection method: 'single' or 'weighted'
622 x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
623 n = min(5, len(x)) # number of previous results to consider
624 x = x[np.argsort(-fitness(x))][:n] # top n mutations
625 w = fitness(x) - fitness(x).min() + 1E-6 # weights (sum > 0)
626 if parent == 'single' or len(x) == 1:
627 # x = x[random.randint(0, n - 1)] # random selection
628 x = x[random.choices(range(n), weights=w)[0]] # weighted selection
629 elif parent == 'weighted':
630 x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination
631
632 # Mutate
633 mp, s = 0.8, 0.2 # mutation probability, sigma
634 npr = np.random
635 npr.seed(int(time.time()))
636 g = np.array([meta[k][0] for k in hyp.keys()]) # gains 0-1
637 ng = len(meta)
638 v = np.ones(ng)
639 while all(v == 1): # mutate until a change occurs (prevent duplicates)
640 v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)
641 for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)
642 hyp[k] = float(x[i + 7] * v[i]) # mutate
643
644 # Constrain to limits
645 for k, v in meta.items():
646 hyp[k] = max(hyp[k], v[1]) # lower limit
647 hyp[k] = min(hyp[k], v[2]) # upper limit
648 hyp[k] = round(hyp[k], 5) # significant digits
649
650 # Train mutation
651 results = train(hyp.copy(), opt, device, callbacks)
652 callbacks = Callbacks()
653 # Write mutation results
654 print_mutation(KEYS, results, hyp.copy(), save_dir, opt.bucket)
655
656 # Plot results
657 plot_evolve(evolve_csv)
658 LOGGER.info(f'Hyperparameter evolution finished {opt.evolve} generations\n'
659 f"Results saved to {colorstr('bold', save_dir)}\n"
660 f'Usage example: $ python train.py --hyp {evolve_yaml}')
661
662
663 def run(**kwargs):
664 # Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
665 opt = parse_opt(True)
666 for k, v in kwargs.items():
667 setattr(opt, k, v)
668 main(opt)
669 return opt
670
671
672 if __name__ == "__main__":
673 opt = parse_opt()
674 main(opt)
segment/val.py deleted 100644 → 0
1 # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
2 """
3 Validate a trained YOLOv5 segment model on a segment dataset
4
5 Usage:
6 $ bash data/scripts/get_coco.sh --val --segments # download COCO-segments val split (1G, 5000 images)
7 $ python segment/val.py --weights yolov5s-seg.pt --data coco.yaml --img 640- # validate COCO-segments
8
9 Usage - formats:
10 $ python segment/val.py --weights yolov5s-seg.pt # PyTorch
11 yolov5s-seg.torchscript # TorchScript
12 yolov5s-seg.onnx # ONNX Runtime or OpenCV DNN with --dnn
13 yolov5s-seg.xml # OpenVINO
14 yolov5s-seg.engine # TensorRT
15 yolov5s-seg.mlmodel # CoreML (macOS-only)
16 yolov5s-seg_saved_model # TensorFlow SavedModel
17 yolov5s-seg.pb # TensorFlow GraphDef
18 yolov5s-seg.tflite # TensorFlow Lite
19 yolov5s-seg_edgetpu.tflite # TensorFlow Edge TPU
20 yolov5s-seg_paddle_model # PaddlePaddle
21 """
22
23 import argparse
24 import json
25 import os
26 import sys
27 from multiprocessing.pool import ThreadPool
28 from pathlib import Path
29
30 import numpy as np
31 import torch
32 from tqdm import tqdm
33
34 FILE = Path(__file__).resolve()
35 ROOT = FILE.parents[1] # YOLOv5 root directory
36 if str(ROOT) not in sys.path:
37 sys.path.append(str(ROOT)) # add ROOT to PATH
38 ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
39
40 import torch.nn.functional as F
41
42 from models.common import DetectMultiBackend
43 from models.yolo import SegmentationModel
44 from utils.callbacks import Callbacks
45 from utils.general import (LOGGER, NUM_THREADS, Profile, check_dataset, check_img_size, check_requirements, check_yaml,
46 coco80_to_coco91_class, colorstr, increment_path, non_max_suppression, print_args,
47 scale_boxes, xywh2xyxy, xyxy2xywh)
48 from utils.metrics import ConfusionMatrix, box_iou
49 from utils.plots import output_to_target, plot_val_study
50 from utils.segment.dataloaders import create_dataloader
51 from utils.segment.general import mask_iou, process_mask, process_mask_upsample, scale_image
52 from utils.segment.metrics import Metrics, ap_per_class_box_and_mask
53 from utils.segment.plots import plot_images_and_masks
54 from utils.torch_utils import de_parallel, select_device, smart_inference_mode
55
56
57 def save_one_txt(predn, save_conf, shape, file):
58 # Save one txt result
59 gn = torch.tensor(shape)[[1, 0, 1, 0]] # normalization gain whwh
60 for *xyxy, conf, cls in predn.tolist():
61 xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
62 line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
63 with open(file, 'a') as f:
64 f.write(('%g ' * len(line)).rstrip() % line + '\n')
65
66
67 def save_one_json(predn, jdict, path, class_map, pred_masks):
68 # Save one JSON result {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}
69 from pycocotools.mask import encode
70
71 def single_encode(x):
72 rle = encode(np.asarray(x[:, :, None], order="F", dtype="uint8"))[0]
73 rle["counts"] = rle["counts"].decode("utf-8")
74 return rle
75
76 image_id = int(path.stem) if path.stem.isnumeric() else path.stem
77 box = xyxy2xywh(predn[:, :4]) # xywh
78 box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
79 pred_masks = np.transpose(pred_masks, (2, 0, 1))
80 with ThreadPool(NUM_THREADS) as pool:
81 rles = pool.map(single_encode, pred_masks)
82 for i, (p, b) in enumerate(zip(predn.tolist(), box.tolist())):
83 jdict.append({
84 'image_id': image_id,
85 'category_id': class_map[int(p[5])],
86 'bbox': [round(x, 3) for x in b],
87 'score': round(p[4], 5),
88 'segmentation': rles[i]})
89
90
91 def process_batch(detections, labels, iouv, pred_masks=None, gt_masks=None, overlap=False, masks=False):
92 """
93 Return correct prediction matrix
94 Arguments:
95 detections (array[N, 6]), x1, y1, x2, y2, conf, class
96 labels (array[M, 5]), class, x1, y1, x2, y2
97 Returns:
98 correct (array[N, 10]), for 10 IoU levels
99 """
100 if masks:
101 if overlap:
102 nl = len(labels)
103 index = torch.arange(nl, device=gt_masks.device).view(nl, 1, 1) + 1
104 gt_masks = gt_masks.repeat(nl, 1, 1) # shape(1,640,640) -> (n,640,640)
105 gt_masks = torch.where(gt_masks == index, 1.0, 0.0)
106 if gt_masks.shape[1:] != pred_masks.shape[1:]:
107 gt_masks = F.interpolate(gt_masks[None], pred_masks.shape[1:], mode="bilinear", align_corners=False)[0]
108 gt_masks = gt_masks.gt_(0.5)
109 iou = mask_iou(gt_masks.view(gt_masks.shape[0], -1), pred_masks.view(pred_masks.shape[0], -1))
110 else: # boxes
111 iou = box_iou(labels[:, 1:], detections[:, :4])
112
113 correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
114 correct_class = labels[:, 0:1] == detections[:, 5]
115 for i in range(len(iouv)):
116 x = torch.where((iou >= iouv[i]) & correct_class) # IoU > threshold and classes match
117 if x[0].shape[0]:
118 matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() # [label, detect, iou]
119 if x[0].shape[0] > 1:
120 matches = matches[matches[:, 2].argsort()[::-1]]
121 matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
122 # matches = matches[matches[:, 2].argsort()[::-1]]
123 matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
124 correct[matches[:, 1].astype(int), i] = True
125 return torch.tensor(correct, dtype=torch.bool, device=iouv.device)
126
127
128 @smart_inference_mode()
129 def run(
130 data,
131 weights=None, # model.pt path(s)
132 batch_size=32, # batch size
133 imgsz=640, # inference size (pixels)
134 conf_thres=0.001, # confidence threshold
135 iou_thres=0.6, # NMS IoU threshold
136 max_det=300, # maximum detections per image
137 task='val', # train, val, test, speed or study
138 device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
139 workers=8, # max dataloader workers (per RANK in DDP mode)
140 single_cls=False, # treat as single-class dataset
141 augment=False, # augmented inference
142 verbose=False, # verbose output
143 save_txt=False, # save results to *.txt
144 save_hybrid=False, # save label+prediction hybrid results to *.txt
145 save_conf=False, # save confidences in --save-txt labels
146 save_json=False, # save a COCO-JSON results file
147 project=ROOT / 'runs/val-seg', # save to project/name
148 name='exp', # save to project/name
149 exist_ok=False, # existing project/name ok, do not increment
150 half=True, # use FP16 half-precision inference
151 dnn=False, # use OpenCV DNN for ONNX inference
152 model=None,
153 dataloader=None,
154 save_dir=Path(''),
155 plots=True,
156 overlap=False,
157 mask_downsample_ratio=1,
158 compute_loss=None,
159 callbacks=Callbacks(),
160 ):
161 if save_json:
162 check_requirements(['pycocotools'])
163 process = process_mask_upsample # more accurate
164 else:
165 process = process_mask # faster
166
167 # Initialize/load model and set device
168 training = model is not None
169 if training: # called by train.py
170 device, pt, jit, engine = next(model.parameters()).device, True, False, False # get model device, PyTorch model
171 half &= device.type != 'cpu' # half precision only supported on CUDA
172 model.half() if half else model.float()
173 nm = de_parallel(model).model[-1].nm # number of masks
174 else: # called directly
175 device = select_device(device, batch_size=batch_size)
176
177 # Directories
178 save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
179 (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
180
181 # Load model
182 model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
183 stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
184 imgsz = check_img_size(imgsz, s=stride) # check image size
185 half = model.fp16 # FP16 supported on limited backends with CUDA
186 nm = de_parallel(model).model.model[-1].nm if isinstance(model, SegmentationModel) else 32 # number of masks
187 if engine:
188 batch_size = model.batch_size
189 else:
190 device = model.device
191 if not (pt or jit):
192 batch_size = 1 # export.py models default to batch-size 1
193 LOGGER.info(f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')
194
195 # Data
196 data = check_dataset(data) # check
197
198 # Configure
199 model.eval()
200 cuda = device.type != 'cpu'
201 is_coco = isinstance(data.get('val'), str) and data['val'].endswith(f'coco{os.sep}val2017.txt') # COCO dataset
202 nc = 1 if single_cls else int(data['nc']) # number of classes
203 iouv = torch.linspace(0.5, 0.95, 10, device=device) # iou vector for mAP@0.5:0.95
204 niou = iouv.numel()
205
206 # Dataloader
207 if not training:
208 if pt and not single_cls: # check --weights are trained on --data
209 ncm = model.model.nc
210 assert ncm == nc, f'{weights} ({ncm} classes) trained on different --data than what you passed ({nc} ' \
211 f'classes). Pass correct combination of --weights and --data that are trained together.'
212 model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz)) # warmup
213 pad, rect = (0.0, False) if task == 'speed' else (0.5, pt) # square inference for benchmarks
214 task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images
215 dataloader = create_dataloader(data[task],
216 imgsz,
217 batch_size,
218 stride,
219 single_cls,
220 pad=pad,
221 rect=rect,
222 workers=workers,
223 prefix=colorstr(f'{task}: '),
224 overlap_mask=overlap,
225 mask_downsample_ratio=mask_downsample_ratio)[0]
226
227 seen = 0
228 confusion_matrix = ConfusionMatrix(nc=nc)
229 names = model.names if hasattr(model, 'names') else model.module.names # get class names
230 if isinstance(names, (list, tuple)): # old format
231 names = dict(enumerate(names))
232 class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
233 s = ('%22s' + '%11s' * 10) % ('Class', 'Images', 'Instances', 'Box(P', "R", "mAP50", "mAP50-95)", "Mask(P", "R",
234 "mAP50", "mAP50-95)")
235 dt = Profile(), Profile(), Profile()
236 metrics = Metrics()
237 loss = torch.zeros(4, device=device)
238 jdict, stats = [], []
239 # callbacks.run('on_val_start')
240 pbar = tqdm(dataloader, desc=s, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
241 for batch_i, (im, targets, paths, shapes, masks) in enumerate(pbar):
242 # callbacks.run('on_val_batch_start')
243 with dt[0]:
244 if cuda:
245 im = im.to(device, non_blocking=True)
246 targets = targets.to(device)
247 masks = masks.to(device)
248 masks = masks.float()
249 im = im.half() if half else im.float() # uint8 to fp16/32
250 im /= 255 # 0 - 255 to 0.0 - 1.0
251 nb, _, height, width = im.shape # batch size, channels, height, width
252
253 # Inference
254 with dt[1]:
255 preds, protos, train_out = model(im) if compute_loss else (*model(im, augment=augment)[:2], None)
256
257 # Loss
258 if compute_loss:
259 loss += compute_loss((train_out, protos), targets, masks)[1] # box, obj, cls
260
261 # NMS
262 targets[:, 2:] *= torch.tensor((width, height, width, height), device=device) # to pixels
263 lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
264 with dt[2]:
265 preds = non_max_suppression(preds,
266 conf_thres,
267 iou_thres,
268 labels=lb,
269 multi_label=True,
270 agnostic=single_cls,
271 max_det=max_det,
272 nm=nm)
273
274 # Metrics
275 plot_masks = [] # masks for plotting
276 for si, (pred, proto) in enumerate(zip(preds, protos)):
277 labels = targets[targets[:, 0] == si, 1:]
278 nl, npr = labels.shape[0], pred.shape[0] # number of labels, predictions
279 path, shape = Path(paths[si]), shapes[si][0]
280 correct_masks = torch.zeros(npr, niou, dtype=torch.bool, device=device) # init
281 correct_bboxes = torch.zeros(npr, niou, dtype=torch.bool, device=device) # init
282 seen += 1
283
284 if npr == 0:
285 if nl:
286 stats.append((correct_masks, correct_bboxes, *torch.zeros((2, 0), device=device), labels[:, 0]))
287 if plots:
288 confusion_matrix.process_batch(detections=None, labels=labels[:, 0])
289 continue
290
291 # Masks
292 midx = [si] if overlap else targets[:, 0] == si
293 gt_masks = masks[midx]
294 pred_masks = process(proto, pred[:, 6:], pred[:, :4], shape=im[si].shape[1:])
295
296 # Predictions
297 if single_cls:
298 pred[:, 5] = 0
299 predn = pred.clone()
300 scale_boxes(im[si].shape[1:], predn[:, :4], shape, shapes[si][1]) # native-space pred
301
302 # Evaluate
303 if nl:
304 tbox = xywh2xyxy(labels[:, 1:5]) # target boxes
305 scale_boxes(im[si].shape[1:], tbox, shape, shapes[si][1]) # native-space labels
306 labelsn = torch.cat((labels[:, 0:1], tbox), 1) # native-space labels
307 correct_bboxes = process_batch(predn, labelsn, iouv)
308 correct_masks = process_batch(predn, labelsn, iouv, pred_masks, gt_masks, overlap=overlap, masks=True)
309 if plots:
310 confusion_matrix.process_batch(predn, labelsn)
311 stats.append((correct_masks, correct_bboxes, pred[:, 4], pred[:, 5], labels[:, 0])) # (conf, pcls, tcls)
312
313 pred_masks = torch.as_tensor(pred_masks, dtype=torch.uint8)
314 if plots and batch_i < 3:
315 plot_masks.append(pred_masks[:15].cpu()) # filter top 15 to plot
316
317 # Save/log
318 if save_txt:
319 save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt')
320 if save_json:
321 pred_masks = scale_image(im[si].shape[1:],
322 pred_masks.permute(1, 2, 0).contiguous().cpu().numpy(), shape, shapes[si][1])
323 save_one_json(predn, jdict, path, class_map, pred_masks) # append to COCO-JSON dictionary
324 # callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
325
326 # Plot images
327 if plots and batch_i < 3:
328 if len(plot_masks):
329 plot_masks = torch.cat(plot_masks, dim=0)
330 plot_images_and_masks(im, targets, masks, paths, save_dir / f'val_batch{batch_i}_labels.jpg', names)
331 plot_images_and_masks(im, output_to_target(preds, max_det=15), plot_masks, paths,
332 save_dir / f'val_batch{batch_i}_pred.jpg', names) # pred
333
334 # callbacks.run('on_val_batch_end')
335
336 # Compute metrics
337 stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)] # to numpy
338 if len(stats) and stats[0].any():
339 results = ap_per_class_box_and_mask(*stats, plot=plots, save_dir=save_dir, names=names)
340 metrics.update(results)
341 nt = np.bincount(stats[4].astype(int), minlength=nc) # number of targets per class
342
343 # Print results
344 pf = '%22s' + '%11i' * 2 + '%11.3g' * 8 # print format
345 LOGGER.info(pf % ("all", seen, nt.sum(), *metrics.mean_results()))
346 if nt.sum() == 0:
347 LOGGER.warning(f'WARNING ⚠️ no labels found in {task} set, can not compute metrics without labels')
348
349 # Print results per class
350 if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
351 for i, c in enumerate(metrics.ap_class_index):
352 LOGGER.info(pf % (names[c], seen, nt[c], *metrics.class_result(i)))
353
354 # Print speeds
355 t = tuple(x.t / seen * 1E3 for x in dt) # speeds per image
356 if not training:
357 shape = (batch_size, 3, imgsz, imgsz)
358 LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}' % t)
359
360 # Plots
361 if plots:
362 confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
363 # callbacks.run('on_val_end')
364
365 mp_bbox, mr_bbox, map50_bbox, map_bbox, mp_mask, mr_mask, map50_mask, map_mask = metrics.mean_results()
366
367 # Save JSON
368 if save_json and len(jdict):
369 w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights
370 anno_json = str(Path(data.get('path', '../coco')) / 'annotations/instances_val2017.json') # annotations json
371 pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
372 LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
373 with open(pred_json, 'w') as f:
374 json.dump(jdict, f)
375
376 try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
377 from pycocotools.coco import COCO
378 from pycocotools.cocoeval import COCOeval
379
380 anno = COCO(anno_json) # init annotations api
381 pred = anno.loadRes(pred_json) # init predictions api
382 results = []
383 for eval in COCOeval(anno, pred, 'bbox'), COCOeval(anno, pred, 'segm'):
384 if is_coco:
385 eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.im_files] # img ID to evaluate
386 eval.evaluate()
387 eval.accumulate()
388 eval.summarize()
389 results.extend(eval.stats[:2]) # update results (mAP@0.5:0.95, mAP@0.5)
390 map_bbox, map50_bbox, map_mask, map50_mask = results
391 except Exception as e:
392 LOGGER.info(f'pycocotools unable to run: {e}')
393
394 # Return results
395 model.float() # for training
396 if not training:
397 s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
398 LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
399 final_metric = mp_bbox, mr_bbox, map50_bbox, map_bbox, mp_mask, mr_mask, map50_mask, map_mask
400 return (*final_metric, *(loss.cpu() / len(dataloader)).tolist()), metrics.get_maps(nc), t
401
402
403 def parse_opt():
404 parser = argparse.ArgumentParser()
405 parser.add_argument('--data', type=str, default=ROOT / 'data/coco128-seg.yaml', help='dataset.yaml path')
406 parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-seg.pt', help='model path(s)')
407 parser.add_argument('--batch-size', type=int, default=32, help='batch size')
408 parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
409 parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')
410 parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')
411 parser.add_argument('--max-det', type=int, default=300, help='maximum detections per image')
412 parser.add_argument('--task', default='val', help='train, val, test, speed or study')
413 parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
414 parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
415 parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
416 parser.add_argument('--augment', action='store_true', help='augmented inference')
417 parser.add_argument('--verbose', action='store_true', help='report mAP by class')
418 parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
419 parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
420 parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
421 parser.add_argument('--save-json', action='store_true', help='save a COCO-JSON results file')
422 parser.add_argument('--project', default=ROOT / 'runs/val-seg', help='save results to project/name')
423 parser.add_argument('--name', default='exp', help='save to project/name')
424 parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
425 parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
426 parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
427 opt = parser.parse_args()
428 opt.data = check_yaml(opt.data) # check YAML
429 # opt.save_json |= opt.data.endswith('coco.yaml')
430 opt.save_txt |= opt.save_hybrid
431 print_args(vars(opt))
432 return opt
433
434
435 def main(opt):
436 check_requirements(requirements=ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
437
438 if opt.task in ('train', 'val', 'test'): # run normally
439 if opt.conf_thres > 0.001: # https://github.com/ultralytics/yolov5/issues/1466
440 LOGGER.warning(f'WARNING ⚠️ confidence threshold {opt.conf_thres} > 0.001 produces invalid results')
441 if opt.save_hybrid:
442 LOGGER.warning('WARNING ⚠️ --save-hybrid returns high mAP from hybrid labels, not from predictions alone')
443 run(**vars(opt))
444
445 else:
446 weights = opt.weights if isinstance(opt.weights, list) else [opt.weights]
447 opt.half = True # FP16 for fastest results
448 if opt.task == 'speed': # speed benchmarks
449 # python val.py --task speed --data coco.yaml --batch 1 --weights yolov5n.pt yolov5s.pt...
450 opt.conf_thres, opt.iou_thres, opt.save_json = 0.25, 0.45, False
451 for opt.weights in weights:
452 run(**vars(opt), plots=False)
453
454 elif opt.task == 'study': # speed vs mAP benchmarks
455 # python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n.pt yolov5s.pt...
456 for opt.weights in weights:
457 f = f'study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt' # filename to save to
458 x, y = list(range(256, 1536 + 128, 128)), [] # x axis (image sizes), y axis
459 for opt.imgsz in x: # img-size
460 LOGGER.info(f'\nRunning {f} --imgsz {opt.imgsz}...')
461 r, _, t = run(**vars(opt), plots=False)
462 y.append(r + t) # results and times
463 np.savetxt(f, y, fmt='%10.4g') # save
464 os.system('zip -r study.zip study_*.txt')
465 plot_val_study(x=x) # plot
466
467
468 if __name__ == "__main__":
469 opt = parse_opt()
470 main(opt)
setup.cfg deleted 100644 → 0
1 # Project-wide configuration file, can be used for package metadata and other toll configurations
2 # Example usage: global configuration for PEP8 (via flake8) setting or default pytest arguments
3 # Local usage: pip install pre-commit, pre-commit run --all-files
4
5 [metadata]
6 license_file = LICENSE
7 description_file = README.md
8
9
10 [tool:pytest]
11 norecursedirs =
12 .git
13 dist
14 build
15 addopts =
16 --doctest-modules
17 --durations=25
18 --color=yes
19
20
21 [flake8]
22 max-line-length = 120
23 exclude = .tox,*.egg,build,temp
24 select = E,W,F
25 doctests = True
26 verbose = 2
27 # https://pep8.readthedocs.io/en/latest/intro.html#error-codes
28 format = pylint
29 # see: https://www.flake8rules.com/
30 ignore =
31 E731 # Do not assign a lambda expression, use a def
32 F405 # name may be undefined, or defined from star imports: module
33 E402 # module level import not at top of file
34 F401 # module imported but unused
35 W504 # line break after binary operator
36 E127 # continuation line over-indented for visual indent
37 E231 # missing whitespace after ‘,’, ‘;’, or ‘:’
38 E501 # line too long
39 F403 # ‘from module import *’ used; unable to detect undefined names
40
41
42 [isort]
43 # https://pycqa.github.io/isort/docs/configuration/options.html
44 line_length = 120
45 # see: https://pycqa.github.io/isort/docs/configuration/multi_line_output_modes.html
46 multi_line_output = 0
47
48
49 [yapf]
50 based_on_style = pep8
51 spaces_before_comment = 2
52 COLUMN_LIMIT = 120
53 COALESCE_BRACKETS = True
54 SPACES_AROUND_POWER_OPERATOR = True
55 SPACE_BETWEEN_ENDING_COMMA_AND_CLOSING_BRACKET = False
56 SPLIT_BEFORE_CLOSING_BRACKET = False
57 SPLIT_BEFORE_FIRST_ARGUMENT = False
58 # EACH_DICT_ENTRY_ON_SEPARATE_LINE = False
utils/__pycache__/__init__.cpython-36.pyc 0 → 100644
No preview for this file type
utils/__pycache__/augmentations.cpython-36.pyc 0 → 100644
No preview for this file type
utils/__pycache__/autoanchor.cpython-36.pyc 0 → 100644
No preview for this file type
utils/__pycache__/dataloaders.cpython-36.pyc 0 → 100644
No preview for this file type
utils/__pycache__/downloads.cpython-36.pyc 0 → 100644
No preview for this file type
utils/__pycache__/general.cpython-36.pyc 0 → 100644
No preview for this file type
utils/__pycache__/metrics.cpython-36.pyc 0 → 100644
No preview for this file type
utils/__pycache__/plots.cpython-36.pyc 0 → 100644
No preview for this file type
utils/__pycache__/torch_utils.cpython-36.pyc 0 → 100644
No preview for this file type
Styling with Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!