Commit
•
d857ef1
1
Parent(s):
d249dd9
init model
Browse files- README.md +143 -0
- anchor_grid.npy +3 -0
- coco.yaml +28 -0
- demo.jpg +0 -0
- general_json2yolo.py +141 -0
- grid.npy +3 -0
- onnx_eval.py +270 -0
- onnx_inference.py +137 -0
- requirements.txt +37 -0
- utils.py +990 -0
- yolov5s_qat.onnx +3 -0
README.md
ADDED
@@ -0,0 +1,143 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
license: apache-2.0
|
3 |
+
tags:
|
4 |
+
- RyzenAI
|
5 |
+
- object-detection
|
6 |
+
- vision
|
7 |
+
- YOLO
|
8 |
+
- Pytorch
|
9 |
+
datasets:
|
10 |
+
- COCO
|
11 |
+
metrics:
|
12 |
+
- mAP
|
13 |
+
---
|
14 |
+
# YOLOv5s model trained on COCO
|
15 |
+
|
16 |
+
YOLOv5s is the small version of YOLOv5 model trained on COCO object detection (118k annotated images) at resolution 640x640. It was released in [https://github.com/ultralytics/yolov5](https://github.com/ultralytics/yolov5).
|
17 |
+
|
18 |
+
We develop a modified version that could be supported by [AMD Ryzen AI](https://onnxruntime.ai/docs/execution-providers/Vitis-AI-ExecutionProvider.html).
|
19 |
+
|
20 |
+
|
21 |
+
## Model description
|
22 |
+
|
23 |
+
YOLOv5 🚀 is the world's most loved vision AI, representing Ultralytics open-source research into future vision AI methods, incorporating lessons learned and best practices evolved over thousands of hours of research and development.
|
24 |
+
|
25 |
+
|
26 |
+
## Intended uses & limitations
|
27 |
+
|
28 |
+
You can use the raw model for object detection. See the [model hub](https://huggingface.co/models?search=amd/yolov5) to look for all available YOLOv5 models.
|
29 |
+
|
30 |
+
|
31 |
+
## How to use
|
32 |
+
|
33 |
+
### Installation
|
34 |
+
|
35 |
+
Follow [Ryzen AI Installation](https://ryzenai.docs.amd.com/en/latest/inst.html) to prepare the environment for Ryzen AI.
|
36 |
+
Run the following script to install pre-requisites for this model.
|
37 |
+
```bash
|
38 |
+
pip install -r requirements.txt
|
39 |
+
```
|
40 |
+
|
41 |
+
|
42 |
+
### Data Preparation (optional: for accuracy evaluation)
|
43 |
+
|
44 |
+
The dataset MSCOCO2017 contains 118287 images for training and 5000 images for validation.
|
45 |
+
|
46 |
+
Download COCO dataset and create directories in your code like this:
|
47 |
+
```plain
|
48 |
+
└── datasets
|
49 |
+
└── coco
|
50 |
+
├── annotations
|
51 |
+
| ├── instances_val2017.json
|
52 |
+
| └── ...
|
53 |
+
├── labels
|
54 |
+
| ├── val2017
|
55 |
+
| | ├── 000000000139.txt
|
56 |
+
| ├── 000000000285.txt
|
57 |
+
| └── ...
|
58 |
+
├── images
|
59 |
+
| ├── val2017
|
60 |
+
| | ├── 000000000139.jpg
|
61 |
+
| ├── 000000000285.jpg
|
62 |
+
└── val2017.txt
|
63 |
+
```
|
64 |
+
1. put the val2017 image folder under images directory or use a softlink
|
65 |
+
2. the labels folder and val2017.txt above are generate by **general_json2yolo.py**
|
66 |
+
3. modify the coco.yaml like this:
|
67 |
+
```markdown
|
68 |
+
path: /path/to/your/datasets/coco # dataset root dir
|
69 |
+
train: train2017.txt # train images (relative to 'path') 118287 images
|
70 |
+
val: val2017.txt # val images (relative to 'path') 5000 images
|
71 |
+
```
|
72 |
+
|
73 |
+
|
74 |
+
### Test & Evaluation
|
75 |
+
|
76 |
+
- Code snippet from [`onnx_inference.py`](onnx_inference.py) on how to use
|
77 |
+
```python
|
78 |
+
args = make_parser().parse_args()
|
79 |
+
onnx_path = args.model
|
80 |
+
onnx_model = onnxruntime.InferenceSession(onnx_path)
|
81 |
+
grid = np.load("./grid.npy", allow_pickle=True)
|
82 |
+
anchor_grid = np.load("./anchor_grid.npy", allow_pickle=True)
|
83 |
+
path = args.image_path
|
84 |
+
new_path = args.output_path
|
85 |
+
conf_thres, iou_thres, classes, agnostic_nms, max_det = 0.25, 0.45, None, False, 1000
|
86 |
+
|
87 |
+
img0 = cv2.imread(path)
|
88 |
+
img = pre_process(img0)
|
89 |
+
onnx_input = {onnx_model.get_inputs()[0].name: img}
|
90 |
+
onnx_output = onnx_model.run(None, onnx_input)
|
91 |
+
onnx_output = post_process(onnx_output)
|
92 |
+
pred = non_max_suppression(
|
93 |
+
onnx_output[0], conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det
|
94 |
+
)
|
95 |
+
colors = Colors()
|
96 |
+
det = pred[0]
|
97 |
+
im0 = img0.copy()
|
98 |
+
annotator = Annotator(im0, line_width=2, example=str(names))
|
99 |
+
if len(det):
|
100 |
+
# Rescale boxes from img_size to im0 size
|
101 |
+
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
|
102 |
+
|
103 |
+
# Write results
|
104 |
+
for *xyxy, conf, cls in reversed(det):
|
105 |
+
c = int(cls) # integer class
|
106 |
+
label = f"{names[c]} {conf:.2f}"
|
107 |
+
annotator.box_label(xyxy, label, color=colors(c, True))
|
108 |
+
# Stream results
|
109 |
+
im0 = annotator.result()
|
110 |
+
cv2.imwrite(new_path, im0)
|
111 |
+
```
|
112 |
+
|
113 |
+
- Run inference for a single image
|
114 |
+
```python
|
115 |
+
python onnx_inference.py -m ./yolov5s_qat.onnx -i /Path/To/Your/Image --ipu --provider_config /Path/To/Your/Provider_config
|
116 |
+
```
|
117 |
+
*Note: __vaip_config.json__ is located at the setup package of Ryzen AI (refer to [Installation](#installation))*
|
118 |
+
- Test accuracy of the quantized model
|
119 |
+
```python
|
120 |
+
python onnx_eval.py -m ./yolov5s_qat.onnx --ipu --provider_config /Path/To/Your/Provider_config
|
121 |
+
```
|
122 |
+
|
123 |
+
### Performance
|
124 |
+
|
125 |
+
|Metric |Accuracy on IPU|
|
126 |
+
| :----: | :----: |
|
127 |
+
|AP\@0.50:0.95|0.356|
|
128 |
+
|
129 |
+
|
130 |
+
```bibtex
|
131 |
+
@software{glenn_jocher_2021_5563715,
|
132 |
+
author = {Glenn Jocher et. al.},
|
133 |
+
title = {{ultralytics/yolov5: v6.0 - YOLOv5n 'Nano' models,
|
134 |
+
Roboflow integration, TensorFlow export, OpenCV
|
135 |
+
DNN support}},
|
136 |
+
month = oct,
|
137 |
+
year = 2021,
|
138 |
+
publisher = {Zenodo},
|
139 |
+
version = {v6.0},
|
140 |
+
doi = {10.5281/zenodo.5563715},
|
141 |
+
url = {https://doi.org/10.5281/zenodo.5563715}
|
142 |
+
}
|
143 |
+
```
|
anchor_grid.npy
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:ccef53e0f9fb34a86d85a3a93d79250598b85025e156294afc34d673bf3965ad
|
3 |
+
size 202917
|
coco.yaml
ADDED
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
2 |
+
# COCO 2017 dataset http://cocodataset.org
|
3 |
+
# Example usage: python train.py --data coco.yaml
|
4 |
+
# parent
|
5 |
+
# ├── yolov5
|
6 |
+
# └── datasets
|
7 |
+
# └── coco ← downloads here
|
8 |
+
|
9 |
+
|
10 |
+
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
|
11 |
+
path: ./datasets/coco # dataset root dir
|
12 |
+
train: train2017.txt
|
13 |
+
val: val2017.txt
|
14 |
+
#train: train2017.txt # train images (relative to 'path') 118287 images
|
15 |
+
#val: val2017.txt # train images (relative to 'path') 5000 images
|
16 |
+
#test: test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
|
17 |
+
|
18 |
+
# Classes
|
19 |
+
nc: 80 # number of classes
|
20 |
+
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
21 |
+
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
22 |
+
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
23 |
+
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
24 |
+
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
25 |
+
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
26 |
+
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
27 |
+
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
28 |
+
'hair drier', 'toothbrush'] # class names
|
demo.jpg
ADDED
![]() |
general_json2yolo.py
ADDED
@@ -0,0 +1,141 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import json
|
2 |
+
from collections import defaultdict
|
3 |
+
import sys
|
4 |
+
import pathlib
|
5 |
+
import numpy as np
|
6 |
+
from tqdm import tqdm
|
7 |
+
CURRENT_DIR = pathlib.Path(__file__).parent
|
8 |
+
sys.path.append(str(CURRENT_DIR))
|
9 |
+
from utils import *
|
10 |
+
|
11 |
+
|
12 |
+
def convert_coco_json(json_dir='../coco/annotations/', use_segments=False, cls91to80=False):
|
13 |
+
save_dir = make_dirs() # output directory
|
14 |
+
coco80 = coco91_to_coco80_class()
|
15 |
+
|
16 |
+
# Import json
|
17 |
+
for json_file in sorted(Path(json_dir).resolve().glob('*.json')):
|
18 |
+
if not str(json_file).endswith("instances_val2017.json"):
|
19 |
+
continue
|
20 |
+
fn = Path(save_dir) / 'labels' / json_file.stem.replace('instances_', '') # folder name
|
21 |
+
fn.mkdir()
|
22 |
+
with open(json_file) as f:
|
23 |
+
data = json.load(f)
|
24 |
+
|
25 |
+
# Create image dict
|
26 |
+
images = {'%g' % x['id']: x for x in data['images']}
|
27 |
+
# Create image-annotations dict
|
28 |
+
imgToAnns = defaultdict(list)
|
29 |
+
for ann in data['annotations']:
|
30 |
+
imgToAnns[ann['image_id']].append(ann)
|
31 |
+
|
32 |
+
txt_file = open(Path(save_dir / 'val2017').with_suffix('.txt'), 'a')
|
33 |
+
# Write labels file
|
34 |
+
for img_id, anns in tqdm(imgToAnns.items(), desc=f'Annotations {json_file}'):
|
35 |
+
img = images['%g' % img_id]
|
36 |
+
h, w, f = img['height'], img['width'], img['file_name']
|
37 |
+
bboxes = []
|
38 |
+
segments = []
|
39 |
+
|
40 |
+
txt_file.write('./images/' + '/'.join(img['coco_url'].split('/')[-2:]) + '\n')
|
41 |
+
for ann in anns:
|
42 |
+
if ann['iscrowd']:
|
43 |
+
continue
|
44 |
+
# The COCO box format is [top left x, top left y, width, height]
|
45 |
+
box = np.array(ann['bbox'], dtype=np.float64)
|
46 |
+
box[:2] += box[2:] / 2 # xy top-left corner to center
|
47 |
+
box[[0, 2]] /= w # normalize x
|
48 |
+
box[[1, 3]] /= h # normalize y
|
49 |
+
if box[2] <= 0 or box[3] <= 0: # if w <= 0 and h <= 0
|
50 |
+
continue
|
51 |
+
|
52 |
+
cls = coco80[ann['category_id'] - 1] if cls91to80 else ann['category_id'] - 1 # class
|
53 |
+
box = [cls] + box.tolist()
|
54 |
+
if box not in bboxes:
|
55 |
+
bboxes.append(box)
|
56 |
+
# Segments
|
57 |
+
if use_segments:
|
58 |
+
if len(ann['segmentation']) > 1:
|
59 |
+
s = merge_multi_segment(ann['segmentation'])
|
60 |
+
s = (np.concatenate(s, axis=0) / np.array([w, h])).reshape(-1).tolist()
|
61 |
+
else:
|
62 |
+
s = [j for i in ann['segmentation'] for j in i] # all segments concatenated
|
63 |
+
s = (np.array(s).reshape(-1, 2) / np.array([w, h])).reshape(-1).tolist()
|
64 |
+
s = [cls] + s
|
65 |
+
if s not in segments:
|
66 |
+
segments.append(s)
|
67 |
+
|
68 |
+
# Write
|
69 |
+
with open((fn / f).with_suffix('.txt'), 'a') as file:
|
70 |
+
for i in range(len(bboxes)):
|
71 |
+
line = *(segments[i] if use_segments else bboxes[i]), # cls, box or segments
|
72 |
+
file.write(('%g ' * len(line)).rstrip() % line + '\n')
|
73 |
+
txt_file.close()
|
74 |
+
|
75 |
+
def min_index(arr1, arr2):
|
76 |
+
"""Find a pair of indexes with the shortest distance.
|
77 |
+
Args:
|
78 |
+
arr1: (N, 2).
|
79 |
+
arr2: (M, 2).
|
80 |
+
Return:
|
81 |
+
a pair of indexes(tuple).
|
82 |
+
"""
|
83 |
+
dis = ((arr1[:, None, :] - arr2[None, :, :]) ** 2).sum(-1)
|
84 |
+
return np.unravel_index(np.argmin(dis, axis=None), dis.shape)
|
85 |
+
|
86 |
+
|
87 |
+
def merge_multi_segment(segments):
|
88 |
+
"""Merge multi segments to one list.
|
89 |
+
Find the coordinates with min distance between each segment,
|
90 |
+
then connect these coordinates with one thin line to merge all
|
91 |
+
segments into one.
|
92 |
+
|
93 |
+
Args:
|
94 |
+
segments(List(List)): original segmentations in coco's json file.
|
95 |
+
like [segmentation1, segmentation2,...],
|
96 |
+
each segmentation is a list of coordinates.
|
97 |
+
"""
|
98 |
+
s = []
|
99 |
+
segments = [np.array(i).reshape(-1, 2) for i in segments]
|
100 |
+
idx_list = [[] for _ in range(len(segments))]
|
101 |
+
|
102 |
+
# record the indexes with min distance between each segment
|
103 |
+
for i in range(1, len(segments)):
|
104 |
+
idx1, idx2 = min_index(segments[i - 1], segments[i])
|
105 |
+
idx_list[i - 1].append(idx1)
|
106 |
+
idx_list[i].append(idx2)
|
107 |
+
|
108 |
+
# use two round to connect all the segments
|
109 |
+
for k in range(2):
|
110 |
+
# forward connection
|
111 |
+
if k == 0:
|
112 |
+
for i, idx in enumerate(idx_list):
|
113 |
+
# middle segments have two indexes
|
114 |
+
# reverse the index of middle segments
|
115 |
+
if len(idx) == 2 and idx[0] > idx[1]:
|
116 |
+
idx = idx[::-1]
|
117 |
+
segments[i] = segments[i][::-1, :]
|
118 |
+
|
119 |
+
segments[i] = np.roll(segments[i], -idx[0], axis=0)
|
120 |
+
segments[i] = np.concatenate([segments[i], segments[i][:1]])
|
121 |
+
# deal with the first segment and the last one
|
122 |
+
if i in [0, len(idx_list) - 1]:
|
123 |
+
s.append(segments[i])
|
124 |
+
else:
|
125 |
+
idx = [0, idx[1] - idx[0]]
|
126 |
+
s.append(segments[i][idx[0]:idx[1] + 1])
|
127 |
+
|
128 |
+
else:
|
129 |
+
for i in range(len(idx_list) - 1, -1, -1):
|
130 |
+
if i not in [0, len(idx_list) - 1]:
|
131 |
+
idx = idx_list[i]
|
132 |
+
nidx = abs(idx[1] - idx[0])
|
133 |
+
s.append(segments[i][nidx:])
|
134 |
+
return s
|
135 |
+
|
136 |
+
|
137 |
+
if __name__ == '__main__':
|
138 |
+
|
139 |
+
convert_coco_json('./datasets/coco/annotations', # directory with *.json
|
140 |
+
use_segments=True,
|
141 |
+
cls91to80=True)
|
grid.npy
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:a1eb0131b1a854173fa17c0ddb74004257d389c29051e40e67c67dab230b6dba
|
3 |
+
size 202917
|
onnx_eval.py
ADDED
@@ -0,0 +1,270 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import argparse
|
2 |
+
import json
|
3 |
+
import os
|
4 |
+
import sys
|
5 |
+
from pathlib import Path
|
6 |
+
import onnxruntime
|
7 |
+
import numpy as np
|
8 |
+
import torch
|
9 |
+
from tqdm import tqdm
|
10 |
+
from pycocotools.coco import COCO
|
11 |
+
from pycocotools.cocoeval import COCOeval
|
12 |
+
|
13 |
+
FILE = Path(__file__).resolve()
|
14 |
+
ROOT = FILE.parents[0] # YOLOv5 root directory
|
15 |
+
if str(ROOT) not in sys.path:
|
16 |
+
sys.path.append(str(ROOT)) # add ROOT to PATH
|
17 |
+
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
18 |
+
import sys
|
19 |
+
import pathlib
|
20 |
+
CURRENT_DIR = pathlib.Path(__file__).parent
|
21 |
+
sys.path.append(str(CURRENT_DIR))
|
22 |
+
from utils import create_dataloader, coco80_to_coco91_class, check_dataset, box_iou, non_max_suppression, post_process, scale_coords, xyxy2xywh, xywh2xyxy, \
|
23 |
+
increment_path, colorstr, ap_per_class
|
24 |
+
|
25 |
+
|
26 |
+
def save_one_txt(predn, save_conf, shape, file):
|
27 |
+
# Save one txt result
|
28 |
+
gn = torch.tensor(shape)[[1, 0, 1, 0]] # normalization gain whwh
|
29 |
+
for *xyxy, conf, cls in predn.tolist():
|
30 |
+
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
|
31 |
+
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
|
32 |
+
with open(file, 'a') as f:
|
33 |
+
f.write(('%g ' * len(line)).rstrip() % line + '\n')
|
34 |
+
|
35 |
+
|
36 |
+
def save_one_json(predn, jdict, path, class_map):
|
37 |
+
# Save one JSON result {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}
|
38 |
+
image_id = int(path.stem) if path.stem.isnumeric() else path.stem
|
39 |
+
box = xyxy2xywh(predn[:, :4]) # xywh
|
40 |
+
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
|
41 |
+
for p, b in zip(predn.tolist(), box.tolist()):
|
42 |
+
jdict.append({'image_id': image_id,
|
43 |
+
'category_id': class_map[int(p[5])],
|
44 |
+
'bbox': [round(x, 3) for x in b],
|
45 |
+
'score': round(p[4], 5)})
|
46 |
+
|
47 |
+
|
48 |
+
def process_batch(detections, labels, iouv):
|
49 |
+
"""
|
50 |
+
Return correct predictions matrix. Both sets of boxes are in (x1, y1, x2, y2) format.
|
51 |
+
Arguments:
|
52 |
+
detections (Array[N, 6]), x1, y1, x2, y2, conf, class
|
53 |
+
labels (Array[M, 5]), class, x1, y1, x2, y2
|
54 |
+
Returns:
|
55 |
+
correct (Array[N, 10]), for 10 IoU levels
|
56 |
+
"""
|
57 |
+
correct = torch.zeros(detections.shape[0], iouv.shape[0], dtype=torch.bool, device=iouv.device)
|
58 |
+
iou = box_iou(labels[:, 1:], detections[:, :4])
|
59 |
+
x = torch.where((iou >= iouv[0]) & (labels[:, 0:1] == detections[:, 5])) # IoU above threshold and classes match
|
60 |
+
if x[0].shape[0]:
|
61 |
+
matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() # [label, detection, iou]
|
62 |
+
if x[0].shape[0] > 1:
|
63 |
+
matches = matches[matches[:, 2].argsort()[::-1]]
|
64 |
+
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
|
65 |
+
# matches = matches[matches[:, 2].argsort()[::-1]]
|
66 |
+
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
|
67 |
+
matches = torch.Tensor(matches).to(iouv.device)
|
68 |
+
correct[matches[:, 1].long()] = matches[:, 2:3] >= iouv
|
69 |
+
return correct
|
70 |
+
|
71 |
+
|
72 |
+
def run(data,
|
73 |
+
weights=None, # model.pt path(s)
|
74 |
+
batch_size=32, # batch size
|
75 |
+
imgsz=640, # inference size (pixels)
|
76 |
+
conf_thres=0.001, # confidence threshold
|
77 |
+
iou_thres=0.6, # NMS IoU threshold
|
78 |
+
task='val', # val, test
|
79 |
+
single_cls=False, # treat as single-class dataset
|
80 |
+
save_txt=False, # save results to *.txt
|
81 |
+
save_hybrid=False, # save label+prediction hybrid results to *.txt
|
82 |
+
save_conf=False, # save confidences in --save-txt labels
|
83 |
+
save_json=False, # save a COCO-JSON results file
|
84 |
+
project=ROOT / 'runs/val', # save to project/name
|
85 |
+
name='exp', # save to project/name
|
86 |
+
exist_ok=False, # existing project/name ok, do not increment
|
87 |
+
half=True, # use FP16 half-precision inference
|
88 |
+
plots=False,
|
89 |
+
onnx_weights="./yolov5s_qat.onnx",
|
90 |
+
ipu=False,
|
91 |
+
provider_config='',
|
92 |
+
):
|
93 |
+
# Initialize/load model and set device
|
94 |
+
device = torch.device('cpu')
|
95 |
+
|
96 |
+
# Directories
|
97 |
+
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
|
98 |
+
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
|
99 |
+
|
100 |
+
# Load model
|
101 |
+
if isinstance(onnx_weights, list):
|
102 |
+
onnx_weights = onnx_weights[0]
|
103 |
+
if ipu:
|
104 |
+
providers = ["VitisAIExecutionProvider"]
|
105 |
+
provider_options = [{"config_file": provider_config}]
|
106 |
+
onnx_model = onnxruntime.InferenceSession(onnx_weights, providers=providers, provider_options=provider_options)
|
107 |
+
else:
|
108 |
+
onnx_model = onnxruntime.InferenceSession(onnx_weights)
|
109 |
+
|
110 |
+
# Data
|
111 |
+
data = check_dataset(data) # check
|
112 |
+
gs = 32 # grid size (max stride)
|
113 |
+
|
114 |
+
is_coco = isinstance(data.get('val'), str) and data['val'].endswith('val2017.txt') # COCO dataset
|
115 |
+
nc = 1 if single_cls else int(data['nc']) # number of classes
|
116 |
+
iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95
|
117 |
+
niou = iouv.numel()
|
118 |
+
|
119 |
+
# Dataloader
|
120 |
+
pad = 0.0 if task == 'speed' else 0.5
|
121 |
+
task = 'val' # path to val/test images
|
122 |
+
dataloader = create_dataloader(data[task], imgsz, batch_size, gs, single_cls, pad=pad, rect=False,
|
123 |
+
prefix=colorstr(f'{task}: '), workers=8)[0]
|
124 |
+
|
125 |
+
seen = 0
|
126 |
+
names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
127 |
+
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
128 |
+
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
129 |
+
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
130 |
+
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
131 |
+
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
132 |
+
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
133 |
+
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
134 |
+
'hair drier', 'toothbrush']
|
135 |
+
class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
|
136 |
+
s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
|
137 |
+
dt, p, r, f1, mp, mr, map50, map = [0.0, 0.0, 0.0], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
|
138 |
+
loss = torch.zeros(3, device=device)
|
139 |
+
jdict, stats, ap, ap_class = [], [], [], []
|
140 |
+
|
141 |
+
for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s, total=len(dataloader))):
|
142 |
+
img = img.to(device, non_blocking=True)
|
143 |
+
img = img.half() if half else img.float() # uint8 to fp16/32
|
144 |
+
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
145 |
+
targets = targets.to(device)
|
146 |
+
nb, _, height, width = img.shape # batch size, channels, height, width
|
147 |
+
|
148 |
+
outputs = onnx_model.run(None, {onnx_model.get_inputs()[0].name: img.cpu().numpy()})
|
149 |
+
outputs = [torch.tensor(item).to(device) for item in outputs]
|
150 |
+
outputs = post_process(outputs)
|
151 |
+
out, train_out = outputs[0], outputs[1]
|
152 |
+
|
153 |
+
# Run NMS
|
154 |
+
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
|
155 |
+
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
|
156 |
+
out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
|
157 |
+
|
158 |
+
# Statistics per image
|
159 |
+
for si, pred in enumerate(out):
|
160 |
+
labels = targets[targets[:, 0] == si, 1:]
|
161 |
+
nl = len(labels)
|
162 |
+
tcls = labels[:, 0].tolist() if nl else [] # target class
|
163 |
+
path, shape = Path(paths[si]), shapes[si][0]
|
164 |
+
seen += 1
|
165 |
+
|
166 |
+
if len(pred) == 0:
|
167 |
+
if nl:
|
168 |
+
stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
|
169 |
+
continue
|
170 |
+
|
171 |
+
# Predictions
|
172 |
+
if single_cls:
|
173 |
+
pred[:, 5] = 0
|
174 |
+
predn = pred.clone()
|
175 |
+
scale_coords(img[si].shape[1:], predn[:, :4], shape, shapes[si][1]) # native-space pred
|
176 |
+
|
177 |
+
# Evaluate
|
178 |
+
if nl:
|
179 |
+
tbox = xywh2xyxy(labels[:, 1:5]) # target boxes
|
180 |
+
scale_coords(img[si].shape[1:], tbox, shape, shapes[si][1]) # native-space labels
|
181 |
+
labelsn = torch.cat((labels[:, 0:1], tbox), 1) # native-space labels
|
182 |
+
correct = process_batch(predn, labelsn, iouv)
|
183 |
+
else:
|
184 |
+
correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool)
|
185 |
+
stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls)) # (correct, conf, pcls, tcls)
|
186 |
+
|
187 |
+
# Save/log
|
188 |
+
if save_txt:
|
189 |
+
save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / (path.stem + '.txt'))
|
190 |
+
if save_json:
|
191 |
+
save_one_json(predn, jdict, path, class_map) # append to COCO-JSON dictionary
|
192 |
+
|
193 |
+
# Compute statistics
|
194 |
+
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
|
195 |
+
if len(stats) and stats[0].any():
|
196 |
+
p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
|
197 |
+
ap50, ap = ap[:, 0], ap.mean(1) # AP@0.5, AP@0.5:0.95
|
198 |
+
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
|
199 |
+
nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
|
200 |
+
else:
|
201 |
+
nt = torch.zeros(1)
|
202 |
+
|
203 |
+
# Print results
|
204 |
+
pf = '%20s' + '%11i' * 2 + '%11.3g' * 4 # print format
|
205 |
+
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
|
206 |
+
|
207 |
+
# Save JSON
|
208 |
+
if save_json and len(jdict):
|
209 |
+
w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights
|
210 |
+
anno_json = str(Path(data.get('path', '../coco')) / 'annotations/instances_val2017.json') # annotations json
|
211 |
+
pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
|
212 |
+
print(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
|
213 |
+
with open(pred_json, 'w') as f:
|
214 |
+
json.dump(jdict, f)
|
215 |
+
|
216 |
+
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
|
217 |
+
anno = COCO(anno_json) # init annotations api
|
218 |
+
pred = anno.loadRes(pred_json) # init predictions api
|
219 |
+
eval = COCOeval(anno, pred, 'bbox')
|
220 |
+
if is_coco:
|
221 |
+
eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] # image IDs to evaluate
|
222 |
+
eval.evaluate()
|
223 |
+
eval.accumulate()
|
224 |
+
eval.summarize()
|
225 |
+
map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
|
226 |
+
except Exception as e:
|
227 |
+
print(f'pycocotools unable to run: {e}')
|
228 |
+
|
229 |
+
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
|
230 |
+
print(f"Results saved to {colorstr('bold', save_dir)}{s}")
|
231 |
+
maps = np.zeros(nc) + map
|
232 |
+
for i, c in enumerate(ap_class):
|
233 |
+
maps[c] = ap[i]
|
234 |
+
return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, 0
|
235 |
+
|
236 |
+
|
237 |
+
def parse_opt():
|
238 |
+
parser = argparse.ArgumentParser()
|
239 |
+
parser.add_argument('--data', type=str, default='./coco.yaml', help='path to your dataset.yaml')
|
240 |
+
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
|
241 |
+
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
|
242 |
+
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
|
243 |
+
parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')
|
244 |
+
parser.add_argument('--iou-thres', type=float, default=0.65, help='NMS IoU threshold')
|
245 |
+
parser.add_argument('--task', default='val', help='val, test')
|
246 |
+
parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
|
247 |
+
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
|
248 |
+
parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
|
249 |
+
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
|
250 |
+
parser.add_argument('--save-json', action='store_true', help='save a COCO-JSON results file')
|
251 |
+
parser.add_argument('--project', default=ROOT / 'runs/val', help='save to project/name')
|
252 |
+
parser.add_argument('--name', default='exp', help='save to project/name')
|
253 |
+
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
|
254 |
+
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
|
255 |
+
parser.add_argument('-m', '--onnx_weights', default='./yolov5s_qat.onnx', nargs='+', type=str, help='path to your onnx_weights')
|
256 |
+
parser.add_argument('--ipu', action='store_true', help='flag for ryzen ai')
|
257 |
+
parser.add_argument('--provider_config', default='', type=str, help='provider config for ryzen ai')
|
258 |
+
opt = parser.parse_args()
|
259 |
+
opt.save_json |= opt.data.endswith('coco.yaml')
|
260 |
+
opt.save_txt |= opt.save_hybrid
|
261 |
+
return opt
|
262 |
+
|
263 |
+
|
264 |
+
def main(opt):
|
265 |
+
run(**vars(opt))
|
266 |
+
|
267 |
+
|
268 |
+
if __name__ == "__main__":
|
269 |
+
opt = parse_opt()
|
270 |
+
main(opt)
|
onnx_inference.py
ADDED
@@ -0,0 +1,137 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import onnxruntime
|
2 |
+
import numpy as np
|
3 |
+
import cv2
|
4 |
+
import torch
|
5 |
+
import sys
|
6 |
+
import pathlib
|
7 |
+
CURRENT_DIR = pathlib.Path(__file__).parent
|
8 |
+
sys.path.append(str(CURRENT_DIR))
|
9 |
+
import argparse
|
10 |
+
from utils import (
|
11 |
+
letterbox,
|
12 |
+
non_max_suppression,
|
13 |
+
scale_coords,
|
14 |
+
Annotator,
|
15 |
+
Colors,
|
16 |
+
)
|
17 |
+
|
18 |
+
|
19 |
+
def pre_process(img):
|
20 |
+
img = letterbox(img, [640, 640], stride=32, auto=False)[0]
|
21 |
+
# Convert
|
22 |
+
img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
|
23 |
+
img = np.ascontiguousarray(img)
|
24 |
+
img = img.astype("float32")
|
25 |
+
img = img / 255.0
|
26 |
+
img = img[np.newaxis, :]
|
27 |
+
return img
|
28 |
+
|
29 |
+
|
30 |
+
def post_process(x):
|
31 |
+
x = list(x)
|
32 |
+
z = [] # inference output
|
33 |
+
stride = [8, 16, 32]
|
34 |
+
for i in range(3):
|
35 |
+
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
|
36 |
+
x[i] = (
|
37 |
+
torch.tensor(x[i])
|
38 |
+
.view(bs, 3, 85, ny, nx)
|
39 |
+
.permute(0, 1, 3, 4, 2)
|
40 |
+
.contiguous()
|
41 |
+
)
|
42 |
+
y = x[i].sigmoid()
|
43 |
+
xy = (y[..., 0:2] * 2.0 - 0.5 + grid[i]) * stride[i]
|
44 |
+
wh = (y[..., 2:4] * 2) ** 2 * anchor_grid[i]
|
45 |
+
y = torch.cat((xy, wh, y[..., 4:]), -1)
|
46 |
+
z.append(y.view(bs, -1, 85))
|
47 |
+
|
48 |
+
return (torch.cat(z, 1), x)
|
49 |
+
|
50 |
+
|
51 |
+
def make_parser():
|
52 |
+
parser = argparse.ArgumentParser("onnxruntime inference sample")
|
53 |
+
parser.add_argument(
|
54 |
+
"-m",
|
55 |
+
"--model",
|
56 |
+
type=str,
|
57 |
+
default="./yolov5s_qat.onnx",
|
58 |
+
help="input your onnx model.",
|
59 |
+
)
|
60 |
+
parser.add_argument(
|
61 |
+
"-i",
|
62 |
+
"--image_path",
|
63 |
+
type=str,
|
64 |
+
default='./demo.jpg',
|
65 |
+
help="path to your input image.",
|
66 |
+
)
|
67 |
+
parser.add_argument(
|
68 |
+
"-o",
|
69 |
+
"--output_path",
|
70 |
+
type=str,
|
71 |
+
default='./demo_infer.jpg',
|
72 |
+
help="path to your output directory.",
|
73 |
+
)
|
74 |
+
parser.add_argument(
|
75 |
+
'--ipu',
|
76 |
+
action='store_true',
|
77 |
+
help='flag for ryzen ai'
|
78 |
+
)
|
79 |
+
parser.add_argument(
|
80 |
+
'--provider_config',
|
81 |
+
default='',
|
82 |
+
type=str,
|
83 |
+
help='provider config for ryzen ai'
|
84 |
+
)
|
85 |
+
return parser
|
86 |
+
|
87 |
+
|
88 |
+
names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
89 |
+
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
90 |
+
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
91 |
+
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
92 |
+
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
93 |
+
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
94 |
+
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
95 |
+
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
96 |
+
'hair drier', 'toothbrush']
|
97 |
+
|
98 |
+
|
99 |
+
if __name__ == '__main__':
|
100 |
+
args = make_parser().parse_args()
|
101 |
+
onnx_path = args.model
|
102 |
+
if args.ipu:
|
103 |
+
providers = ["VitisAIExecutionProvider"]
|
104 |
+
provider_options = [{"config_file": args.provider_config}]
|
105 |
+
onnx_model = onnxruntime.InferenceSession(onnx_path, providers=providers, provider_options=provider_options)
|
106 |
+
else:
|
107 |
+
onnx_model = onnxruntime.InferenceSession(onnx_path)
|
108 |
+
grid = np.load("./grid.npy", allow_pickle=True)
|
109 |
+
anchor_grid = np.load("./anchor_grid.npy", allow_pickle=True)
|
110 |
+
path = args.image_path
|
111 |
+
new_path = args.output_path
|
112 |
+
conf_thres, iou_thres, classes, agnostic_nms, max_det = 0.25, 0.45, None, False, 1000
|
113 |
+
|
114 |
+
img0 = cv2.imread(path)
|
115 |
+
img = pre_process(img0)
|
116 |
+
onnx_input = {onnx_model.get_inputs()[0].name: img}
|
117 |
+
onnx_output = onnx_model.run(None, onnx_input)
|
118 |
+
onnx_output = post_process(onnx_output)
|
119 |
+
pred = non_max_suppression(
|
120 |
+
onnx_output[0], conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det
|
121 |
+
)
|
122 |
+
colors = Colors()
|
123 |
+
det = pred[0]
|
124 |
+
im0 = img0.copy()
|
125 |
+
annotator = Annotator(im0, line_width=2, example=str(names))
|
126 |
+
if len(det):
|
127 |
+
# Rescale boxes from img_size to im0 size
|
128 |
+
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
|
129 |
+
|
130 |
+
# Write results
|
131 |
+
for *xyxy, conf, cls in reversed(det):
|
132 |
+
c = int(cls) # integer class
|
133 |
+
label = f"{names[c]} {conf:.2f}"
|
134 |
+
annotator.box_label(xyxy, label, color=colors(c, True))
|
135 |
+
# Stream results
|
136 |
+
im0 = annotator.result()
|
137 |
+
cv2.imwrite(new_path, im0)
|
requirements.txt
ADDED
@@ -0,0 +1,37 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# pip install -r requirements.txt
|
2 |
+
|
3 |
+
# Base ----------------------------------------
|
4 |
+
matplotlib>=3.2.2
|
5 |
+
numpy>=1.18.5
|
6 |
+
opencv-python>=4.1.2
|
7 |
+
Pillow>=7.1.2
|
8 |
+
PyYAML>=5.3.1
|
9 |
+
requests>=2.23.0
|
10 |
+
scipy>=1.4.1
|
11 |
+
tqdm>=4.41.0
|
12 |
+
torch==1.8.0
|
13 |
+
torchvision==0.9.0
|
14 |
+
|
15 |
+
# Logging -------------------------------------
|
16 |
+
tensorboard>=2.4.1
|
17 |
+
# wandb
|
18 |
+
|
19 |
+
# Plotting ------------------------------------
|
20 |
+
pandas>=1.1.4
|
21 |
+
seaborn>=0.11.0
|
22 |
+
|
23 |
+
# Export --------------------------------------
|
24 |
+
# coremltools>=4.1 # CoreML export
|
25 |
+
onnx>=1.9.0 # ONNX export
|
26 |
+
# onnxruntime
|
27 |
+
# onnx-simplifier>=0.3.6 # ONNX simplifier
|
28 |
+
# scikit-learn==0.19.2 # CoreML quantization
|
29 |
+
# tensorflow>=2.4.1 # TFLite export
|
30 |
+
# tensorflowjs>=3.9.0 # TF.js export
|
31 |
+
|
32 |
+
# Extras --------------------------------------
|
33 |
+
# albumentations>=1.0.3
|
34 |
+
# Cython # for pycocotools https://github.com/cocodataset/cocoapi/issues/172
|
35 |
+
pycocotools>=2.0 # COCO mAP
|
36 |
+
# roboflow
|
37 |
+
thop # FLOPs computation
|
utils.py
ADDED
@@ -0,0 +1,990 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
import cv2
|
3 |
+
from pathlib import Path
|
4 |
+
import torch
|
5 |
+
import time
|
6 |
+
import torchvision
|
7 |
+
import re
|
8 |
+
import glob
|
9 |
+
from torch.utils.data import Dataset
|
10 |
+
import yaml
|
11 |
+
import os
|
12 |
+
from multiprocessing.pool import ThreadPool, Pool
|
13 |
+
from tqdm import tqdm
|
14 |
+
from itertools import repeat
|
15 |
+
import logging
|
16 |
+
from PIL import Image, ExifTags
|
17 |
+
import hashlib
|
18 |
+
import sys
|
19 |
+
import pathlib
|
20 |
+
CURRENT_DIR = pathlib.Path(__file__).parent
|
21 |
+
sys.path.append(str(CURRENT_DIR))
|
22 |
+
# Parameters
|
23 |
+
IMG_FORMATS = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo']
|
24 |
+
NUM_THREADS = min(8, os.cpu_count())
|
25 |
+
img_formats = IMG_FORMATS # acceptable image suffixes
|
26 |
+
vid_formats = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'] # acceptable video suffixes
|
27 |
+
|
28 |
+
# Get orientation exif tag
|
29 |
+
for orientation in ExifTags.TAGS.keys():
|
30 |
+
if ExifTags.TAGS[orientation] == 'Orientation':
|
31 |
+
break
|
32 |
+
|
33 |
+
|
34 |
+
def make_dirs(dir='./datasets/coco'):
|
35 |
+
# Create folders
|
36 |
+
dir = Path(dir)
|
37 |
+
for p in [dir / 'labels']:
|
38 |
+
p.mkdir(parents=True, exist_ok=True) # make dir
|
39 |
+
return dir
|
40 |
+
|
41 |
+
|
42 |
+
def coco91_to_coco80_class(): # converts 80-index (val2014) to 91-index (paper)
|
43 |
+
# https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
|
44 |
+
x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, None, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, None, 24, 25, None,
|
45 |
+
None, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, None, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50,
|
46 |
+
51, 52, 53, 54, 55, 56, 57, 58, 59, None, 60, None, None, 61, None, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72,
|
47 |
+
None, 73, 74, 75, 76, 77, 78, 79, None]
|
48 |
+
return x
|
49 |
+
|
50 |
+
|
51 |
+
def is_ascii(s=""):
|
52 |
+
# Is string composed of all ASCII (no UTF) characters? (note str().isascii() introduced in python 3.7)
|
53 |
+
s = str(s) # convert list, tuple, None, etc. to str
|
54 |
+
return len(s.encode().decode("ascii", "ignore")) == len(s)
|
55 |
+
|
56 |
+
|
57 |
+
def is_chinese(s="人工智能"):
|
58 |
+
# Is string composed of any Chinese characters?
|
59 |
+
return re.search("[\u4e00-\u9fff]", s)
|
60 |
+
|
61 |
+
|
62 |
+
def letterbox(
|
63 |
+
im,
|
64 |
+
new_shape=(640, 640),
|
65 |
+
color=(114, 114, 114),
|
66 |
+
auto=True,
|
67 |
+
scaleFill=False,
|
68 |
+
scaleup=True,
|
69 |
+
stride=32,
|
70 |
+
):
|
71 |
+
# Resize and pad image while meeting stride-multiple constraints
|
72 |
+
shape = im.shape[:2] # current shape [height, width]
|
73 |
+
if isinstance(new_shape, int):
|
74 |
+
new_shape = (new_shape, new_shape)
|
75 |
+
|
76 |
+
# Scale ratio (new / old)
|
77 |
+
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
78 |
+
if not scaleup: # only scale down, do not scale up (for better val mAP)
|
79 |
+
r = min(r, 1.0)
|
80 |
+
|
81 |
+
# Compute padding
|
82 |
+
ratio = r, r # width, height ratios
|
83 |
+
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
|
84 |
+
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
|
85 |
+
if auto: # minimum rectangle
|
86 |
+
dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
|
87 |
+
elif scaleFill: # stretch
|
88 |
+
dw, dh = 0.0, 0.0
|
89 |
+
new_unpad = (new_shape[1], new_shape[0])
|
90 |
+
ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
|
91 |
+
|
92 |
+
dw /= 2 # divide padding into 2 sides
|
93 |
+
dh /= 2
|
94 |
+
|
95 |
+
if shape[::-1] != new_unpad: # resize
|
96 |
+
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
|
97 |
+
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
|
98 |
+
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
|
99 |
+
im = cv2.copyMakeBorder(
|
100 |
+
im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color
|
101 |
+
) # add border
|
102 |
+
return im, ratio, (dw, dh)
|
103 |
+
|
104 |
+
|
105 |
+
def xyxy2xywh(x):
|
106 |
+
# Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
|
107 |
+
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
108 |
+
y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center
|
109 |
+
y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center
|
110 |
+
y[:, 2] = x[:, 2] - x[:, 0] # width
|
111 |
+
y[:, 3] = x[:, 3] - x[:, 1] # height
|
112 |
+
return y
|
113 |
+
|
114 |
+
|
115 |
+
def xywh2xyxy(x):
|
116 |
+
# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
|
117 |
+
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
118 |
+
y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
|
119 |
+
y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y
|
120 |
+
y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x
|
121 |
+
y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y
|
122 |
+
return y
|
123 |
+
|
124 |
+
|
125 |
+
def non_max_suppression(
|
126 |
+
prediction,
|
127 |
+
conf_thres=0.25,
|
128 |
+
iou_thres=0.45,
|
129 |
+
classes=None,
|
130 |
+
agnostic=False,
|
131 |
+
multi_label=False,
|
132 |
+
labels=(),
|
133 |
+
max_det=300,
|
134 |
+
):
|
135 |
+
"""Runs Non-Maximum Suppression (NMS) on inference results
|
136 |
+
|
137 |
+
Returns:
|
138 |
+
list of detections, on (n,6) tensor per image [xyxy, conf, cls]
|
139 |
+
"""
|
140 |
+
|
141 |
+
nc = prediction.shape[2] - 5 # number of classes
|
142 |
+
xc = prediction[..., 4] > conf_thres # candidates
|
143 |
+
|
144 |
+
# Checks
|
145 |
+
assert (
|
146 |
+
0 <= conf_thres <= 1
|
147 |
+
), f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0"
|
148 |
+
assert (
|
149 |
+
0 <= iou_thres <= 1
|
150 |
+
), f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0"
|
151 |
+
|
152 |
+
# Settings
|
153 |
+
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
|
154 |
+
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
155 |
+
time_limit = 10.0 # seconds to quit after
|
156 |
+
redundant = True # require redundant detections
|
157 |
+
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
|
158 |
+
merge = False # use merge-NMS
|
159 |
+
|
160 |
+
t = time.time()
|
161 |
+
output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
|
162 |
+
for xi, x in enumerate(prediction): # image index, image inference
|
163 |
+
# Apply constraints
|
164 |
+
# x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
|
165 |
+
x = x[xc[xi]] # confidence
|
166 |
+
|
167 |
+
# Cat apriori labels if autolabelling
|
168 |
+
if labels and len(labels[xi]):
|
169 |
+
l = labels[xi]
|
170 |
+
v = torch.zeros((len(l), nc + 5), device=x.device)
|
171 |
+
v[:, :4] = l[:, 1:5] # box
|
172 |
+
v[:, 4] = 1.0 # conf
|
173 |
+
v[range(len(l)), l[:, 0].long() + 5] = 1.0 # cls
|
174 |
+
x = torch.cat((x, v), 0)
|
175 |
+
|
176 |
+
# If none remain process next image
|
177 |
+
if not x.shape[0]:
|
178 |
+
continue
|
179 |
+
|
180 |
+
# Compute conf
|
181 |
+
x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf
|
182 |
+
|
183 |
+
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
|
184 |
+
box = xywh2xyxy(x[:, :4])
|
185 |
+
|
186 |
+
# Detections matrix nx6 (xyxy, conf, cls)
|
187 |
+
if multi_label:
|
188 |
+
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
|
189 |
+
x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
|
190 |
+
else: # best class only
|
191 |
+
conf, j = x[:, 5:].max(1, keepdim=True)
|
192 |
+
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
|
193 |
+
|
194 |
+
# Filter by class
|
195 |
+
if classes is not None:
|
196 |
+
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
|
197 |
+
|
198 |
+
# Apply finite constraint
|
199 |
+
# if not torch.isfinite(x).all():
|
200 |
+
# x = x[torch.isfinite(x).all(1)]
|
201 |
+
|
202 |
+
# Check shape
|
203 |
+
n = x.shape[0] # number of boxes
|
204 |
+
if not n: # no boxes
|
205 |
+
continue
|
206 |
+
elif n > max_nms: # excess boxes
|
207 |
+
x = x[x[:, 4].argsort(descending=True)[:max_nms]] # sort by confidence
|
208 |
+
|
209 |
+
# Batched NMS
|
210 |
+
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
|
211 |
+
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
|
212 |
+
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
|
213 |
+
if i.shape[0] > max_det: # limit detections
|
214 |
+
i = i[:max_det]
|
215 |
+
if merge and (1 < n < 3e3): # Merge NMS (boxes merged using weighted mean)
|
216 |
+
# update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
|
217 |
+
iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
|
218 |
+
weights = iou * scores[None] # box weights
|
219 |
+
x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(
|
220 |
+
1, keepdim=True
|
221 |
+
) # merged boxes
|
222 |
+
if redundant:
|
223 |
+
i = i[iou.sum(1) > 1] # require redundancy
|
224 |
+
|
225 |
+
output[xi] = x[i]
|
226 |
+
if (time.time() - t) > time_limit:
|
227 |
+
print(f"WARNING: NMS time limit {time_limit}s exceeded")
|
228 |
+
break # time limit exceeded
|
229 |
+
|
230 |
+
return output
|
231 |
+
|
232 |
+
|
233 |
+
def clip_coords(boxes, shape):
|
234 |
+
# Clip bounding xyxy bounding boxes to image shape (height, width)
|
235 |
+
if isinstance(boxes, torch.Tensor): # faster individually
|
236 |
+
boxes[:, 0].clamp_(0, shape[1]) # x1
|
237 |
+
boxes[:, 1].clamp_(0, shape[0]) # y1
|
238 |
+
boxes[:, 2].clamp_(0, shape[1]) # x2
|
239 |
+
boxes[:, 3].clamp_(0, shape[0]) # y2
|
240 |
+
else: # np.array (faster grouped)
|
241 |
+
boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, shape[1]) # x1, x2
|
242 |
+
boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0]) # y1, y2
|
243 |
+
|
244 |
+
|
245 |
+
def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
|
246 |
+
# Rescale coords (xyxy) from img1_shape to img0_shape
|
247 |
+
if ratio_pad is None: # calculate from img0_shape
|
248 |
+
gain = min(
|
249 |
+
img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]
|
250 |
+
) # gain = old / new
|
251 |
+
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (
|
252 |
+
img1_shape[0] - img0_shape[0] * gain
|
253 |
+
) / 2 # wh padding
|
254 |
+
else:
|
255 |
+
gain = ratio_pad[0][0]
|
256 |
+
pad = ratio_pad[1]
|
257 |
+
|
258 |
+
coords[:, [0, 2]] -= pad[0] # x padding
|
259 |
+
coords[:, [1, 3]] -= pad[1] # y padding
|
260 |
+
coords[:, :4] /= gain
|
261 |
+
clip_coords(coords, img0_shape)
|
262 |
+
return coords
|
263 |
+
|
264 |
+
|
265 |
+
class Annotator:
|
266 |
+
# YOLOv5 Annotator for train/val mosaics and jpgs and detect/hub inference annotations
|
267 |
+
def __init__(
|
268 |
+
self,
|
269 |
+
im,
|
270 |
+
line_width=None,
|
271 |
+
font_size=None,
|
272 |
+
font="Arial.ttf",
|
273 |
+
pil=False,
|
274 |
+
example="abc",
|
275 |
+
):
|
276 |
+
assert (
|
277 |
+
im.data.contiguous
|
278 |
+
), "Image not contiguous. Apply np.ascontiguousarray(im) to Annotator() input images."
|
279 |
+
self.pil = pil or not is_ascii(example) or is_chinese(example)
|
280 |
+
self.im = im
|
281 |
+
self.lw = line_width or max(round(sum(im.shape) / 2 * 0.003), 2) # line width
|
282 |
+
|
283 |
+
def box_label(
|
284 |
+
self, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255)
|
285 |
+
):
|
286 |
+
# Add one xyxy box to image with label
|
287 |
+
p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
|
288 |
+
cv2.rectangle(
|
289 |
+
self.im, p1, p2, color, thickness=self.lw, lineType=cv2.LINE_AA
|
290 |
+
)
|
291 |
+
if label:
|
292 |
+
tf = max(self.lw - 1, 1) # font thickness
|
293 |
+
w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[
|
294 |
+
0
|
295 |
+
] # text width, height
|
296 |
+
outside = p1[1] - h - 3 >= 0 # label fits outside box
|
297 |
+
p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
|
298 |
+
cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled
|
299 |
+
cv2.putText(
|
300 |
+
self.im,
|
301 |
+
label,
|
302 |
+
(p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
|
303 |
+
0,
|
304 |
+
self.lw / 3,
|
305 |
+
txt_color,
|
306 |
+
thickness=tf,
|
307 |
+
lineType=cv2.LINE_AA,
|
308 |
+
)
|
309 |
+
|
310 |
+
def rectangle(self, xy, fill=None, outline=None, width=1):
|
311 |
+
# Add rectangle to image (PIL-only)
|
312 |
+
self.draw.rectangle(xy, fill, outline, width)
|
313 |
+
|
314 |
+
def result(self):
|
315 |
+
# Return annotated image as array
|
316 |
+
return np.asarray(self.im)
|
317 |
+
|
318 |
+
|
319 |
+
class Colors:
|
320 |
+
# Ultralytics color palette https://ultralytics.com/
|
321 |
+
def __init__(self):
|
322 |
+
# hex = matplotlib.colors.TABLEAU_COLORS.values()
|
323 |
+
hex = (
|
324 |
+
"FF3838",
|
325 |
+
"FF9D97",
|
326 |
+
"FF701F",
|
327 |
+
"FFB21D",
|
328 |
+
"CFD231",
|
329 |
+
"48F90A",
|
330 |
+
"92CC17",
|
331 |
+
"3DDB86",
|
332 |
+
"1A9334",
|
333 |
+
"00D4BB",
|
334 |
+
"2C99A8",
|
335 |
+
"00C2FF",
|
336 |
+
"344593",
|
337 |
+
"6473FF",
|
338 |
+
"0018EC",
|
339 |
+
"8438FF",
|
340 |
+
"520085",
|
341 |
+
"CB38FF",
|
342 |
+
"FF95C8",
|
343 |
+
"FF37C7",
|
344 |
+
)
|
345 |
+
self.palette = [self.hex2rgb("#" + c) for c in hex]
|
346 |
+
self.n = len(self.palette)
|
347 |
+
|
348 |
+
def __call__(self, i, bgr=False):
|
349 |
+
c = self.palette[int(i) % self.n]
|
350 |
+
return (c[2], c[1], c[0]) if bgr else c
|
351 |
+
|
352 |
+
@staticmethod
|
353 |
+
def hex2rgb(h): # rgb order (PIL)
|
354 |
+
return tuple(int(h[1 + i : 1 + i + 2], 16) for i in (0, 2, 4))
|
355 |
+
|
356 |
+
|
357 |
+
def create_dataloader(path, imgsz, batch_size, stride, single_cls=False, hyp=None, augment=False, cache=False, pad=0.0,
|
358 |
+
rect=False, rank=-1, workers=8, image_weights=False, quad=False, prefix=''):
|
359 |
+
|
360 |
+
dataset = LoadImagesAndLabels(path, imgsz, batch_size,
|
361 |
+
augment=augment, # augment images
|
362 |
+
hyp=hyp, # augmentation hyperparameters
|
363 |
+
rect=rect, # rectangular training
|
364 |
+
cache_images=cache,
|
365 |
+
single_cls=single_cls,
|
366 |
+
stride=int(stride),
|
367 |
+
pad=pad,
|
368 |
+
image_weights=image_weights,
|
369 |
+
prefix=prefix)
|
370 |
+
|
371 |
+
batch_size = min(batch_size, len(dataset))
|
372 |
+
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, workers]) # number of workers
|
373 |
+
sampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank != -1 else None
|
374 |
+
loader = torch.utils.data.DataLoader if image_weights else InfiniteDataLoader
|
375 |
+
# Use torch.utils.data.DataLoader() if dataset.properties will update during training else InfiniteDataLoader()
|
376 |
+
dataloader = loader(dataset,
|
377 |
+
batch_size=batch_size,
|
378 |
+
num_workers=nw,
|
379 |
+
sampler=sampler,
|
380 |
+
pin_memory=True,
|
381 |
+
collate_fn=LoadImagesAndLabels.collate_fn)
|
382 |
+
return dataloader, dataset
|
383 |
+
|
384 |
+
|
385 |
+
class LoadImagesAndLabels(Dataset):
|
386 |
+
# YOLOv5 train_loader/val_loader, loads images and labels for training and validation
|
387 |
+
cache_version = 0.5 # dataset labels *.cache version
|
388 |
+
|
389 |
+
def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
|
390 |
+
cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''):
|
391 |
+
self.img_size = img_size
|
392 |
+
self.augment = augment
|
393 |
+
self.hyp = hyp
|
394 |
+
self.image_weights = image_weights
|
395 |
+
self.rect = False if image_weights else rect
|
396 |
+
self.mosaic = False # load 4 images at a time into a mosaic (only during training)
|
397 |
+
self.mosaic_border = [-img_size // 2, -img_size // 2]
|
398 |
+
self.stride = stride
|
399 |
+
self.path = path
|
400 |
+
self.albumentations = None
|
401 |
+
|
402 |
+
f = [] # image files
|
403 |
+
for p in path if isinstance(path, list) else [path]:
|
404 |
+
p = Path(p) # os-agnostic
|
405 |
+
if p.is_dir(): # dir
|
406 |
+
f += glob.glob(str(p / '**' / '*.*'), recursive=True)
|
407 |
+
# f = list(p.rglob('**/*.*')) # pathlib
|
408 |
+
elif p.is_file(): # file
|
409 |
+
with open(p, 'r') as t:
|
410 |
+
t = t.read().strip().splitlines()
|
411 |
+
parent = str(p.parent) + os.sep
|
412 |
+
f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
|
413 |
+
# f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
|
414 |
+
else:
|
415 |
+
raise Exception(f'{prefix}{p} does not exist')
|
416 |
+
self.img_files = sorted([x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS])
|
417 |
+
# self.img_files = sorted([x for x in f if x.suffix[1:].lower() in img_formats]) # pathlib
|
418 |
+
assert self.img_files, f'{prefix}No images found'
|
419 |
+
|
420 |
+
# Check cache
|
421 |
+
self.label_files = img2label_paths(self.img_files) # labels
|
422 |
+
cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache')
|
423 |
+
try:
|
424 |
+
cache, exists = np.load(cache_path, allow_pickle=True).item(), True # load dict
|
425 |
+
assert cache['version'] == self.cache_version # same version
|
426 |
+
assert cache['hash'] == get_hash(self.label_files + self.img_files) # same hash
|
427 |
+
except:
|
428 |
+
cache, exists = self.cache_labels(cache_path, prefix), False # cache
|
429 |
+
|
430 |
+
# Display cache
|
431 |
+
nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupted, total
|
432 |
+
if exists:
|
433 |
+
d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted"
|
434 |
+
tqdm(None, desc=prefix + d, total=n, initial=n) # display cache results
|
435 |
+
if cache['msgs']:
|
436 |
+
logging.info('\n'.join(cache['msgs'])) # display warnings
|
437 |
+
|
438 |
+
# Read cache
|
439 |
+
[cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove items
|
440 |
+
labels, shapes, self.segments = zip(*cache.values())
|
441 |
+
self.labels = list(labels)
|
442 |
+
self.shapes = np.array(shapes, dtype=np.float64)
|
443 |
+
self.img_files = list(cache.keys()) # update
|
444 |
+
self.label_files = img2label_paths(cache.keys()) # update
|
445 |
+
if single_cls:
|
446 |
+
for x in self.labels:
|
447 |
+
x[:, 0] = 0
|
448 |
+
|
449 |
+
n = len(shapes) # number of images
|
450 |
+
bi = np.floor(np.arange(n) / batch_size).astype(int) # batch index
|
451 |
+
nb = bi[-1] + 1 # number of batches
|
452 |
+
self.batch = bi # batch index of image
|
453 |
+
self.n = n
|
454 |
+
self.indices = range(n)
|
455 |
+
|
456 |
+
# Rectangular Training
|
457 |
+
if self.rect:
|
458 |
+
# Sort by aspect ratio
|
459 |
+
s = self.shapes # wh
|
460 |
+
ar = s[:, 1] / s[:, 0] # aspect ratio
|
461 |
+
irect = ar.argsort()
|
462 |
+
self.img_files = [self.img_files[i] for i in irect]
|
463 |
+
self.label_files = [self.label_files[i] for i in irect]
|
464 |
+
self.labels = [self.labels[i] for i in irect]
|
465 |
+
self.shapes = s[irect] # wh
|
466 |
+
ar = ar[irect]
|
467 |
+
|
468 |
+
# Set training image shapes
|
469 |
+
shapes = [[1, 1]] * nb
|
470 |
+
for i in range(nb):
|
471 |
+
ari = ar[bi == i]
|
472 |
+
mini, maxi = ari.min(), ari.max()
|
473 |
+
if maxi < 1:
|
474 |
+
shapes[i] = [maxi, 1]
|
475 |
+
elif mini > 1:
|
476 |
+
shapes[i] = [1, 1 / mini]
|
477 |
+
|
478 |
+
self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(int) * stride
|
479 |
+
|
480 |
+
# Cache images into memory for faster training (WARNING: large datasets may exceed system RAM)
|
481 |
+
self.imgs, self.img_npy = [None] * n, [None] * n
|
482 |
+
if cache_images:
|
483 |
+
if cache_images == 'disk':
|
484 |
+
self.im_cache_dir = Path(Path(self.img_files[0]).parent.as_posix() + '_npy')
|
485 |
+
self.img_npy = [self.im_cache_dir / Path(f).with_suffix('.npy').name for f in self.img_files]
|
486 |
+
self.im_cache_dir.mkdir(parents=True, exist_ok=True)
|
487 |
+
gb = 0 # Gigabytes of cached images
|
488 |
+
self.img_hw0, self.img_hw = [None] * n, [None] * n
|
489 |
+
results = ThreadPool(NUM_THREADS).imap(lambda x: load_image(*x), zip(repeat(self), range(n)))
|
490 |
+
pbar = tqdm(enumerate(results), total=n)
|
491 |
+
for i, x in pbar:
|
492 |
+
if cache_images == 'disk':
|
493 |
+
if not self.img_npy[i].exists():
|
494 |
+
np.save(self.img_npy[i].as_posix(), x[0])
|
495 |
+
gb += self.img_npy[i].stat().st_size
|
496 |
+
else:
|
497 |
+
self.imgs[i], self.img_hw0[i], self.img_hw[i] = x # im, hw_orig, hw_resized = load_image(self, i)
|
498 |
+
gb += self.imgs[i].nbytes
|
499 |
+
pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB {cache_images})'
|
500 |
+
pbar.close()
|
501 |
+
|
502 |
+
def cache_labels(self, path=Path('./labels.cache'), prefix=''):
|
503 |
+
# Cache dataset labels, check images and read shapes
|
504 |
+
x = {} # dict
|
505 |
+
nm, nf, ne, nc, msgs = 0, 0, 0, 0, [] # number missing, found, empty, corrupt, messages
|
506 |
+
desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels..."
|
507 |
+
with Pool(NUM_THREADS) as pool:
|
508 |
+
pbar = tqdm(pool.imap(verify_image_label, zip(self.img_files, self.label_files, repeat(prefix))), desc=desc, total=len(self.img_files))
|
509 |
+
for im_file, l, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar:
|
510 |
+
nm += nm_f
|
511 |
+
nf += nf_f
|
512 |
+
ne += ne_f
|
513 |
+
nc += nc_f
|
514 |
+
if im_file:
|
515 |
+
x[im_file] = [l, shape, segments]
|
516 |
+
if msg:
|
517 |
+
msgs.append(msg)
|
518 |
+
pbar.desc = f"{desc}{nf} found, {nm} missing, {ne} empty, {nc} corrupted"
|
519 |
+
|
520 |
+
pbar.close()
|
521 |
+
if msgs:
|
522 |
+
logging.info('\n'.join(msgs))
|
523 |
+
x['hash'] = get_hash(self.label_files + self.img_files)
|
524 |
+
x['results'] = nf, nm, ne, nc, len(self.img_files)
|
525 |
+
x['msgs'] = msgs # warnings
|
526 |
+
x['version'] = self.cache_version # cache version
|
527 |
+
try:
|
528 |
+
np.save(path, x) # save cache for next time
|
529 |
+
path.with_suffix('.cache.npy').rename(path) # remove .npy suffix
|
530 |
+
logging.info(f'{prefix}New cache created: {path}')
|
531 |
+
except Exception as e:
|
532 |
+
logging.info(f'{prefix}WARNING: Cache directory {path.parent} is not writeable: {e}') # path not writeable
|
533 |
+
return x
|
534 |
+
|
535 |
+
def __len__(self):
|
536 |
+
return len(self.img_files)
|
537 |
+
|
538 |
+
# def __iter__(self):
|
539 |
+
# self.count = -1
|
540 |
+
# print('ran dataset iter')
|
541 |
+
# #self.shuffled_vector = np.random.permutation(self.nF) if self.augment else np.arange(self.nF)
|
542 |
+
# return self
|
543 |
+
|
544 |
+
def __getitem__(self, index):
|
545 |
+
index = self.indices[index] # linear, shuffled, or image_weights
|
546 |
+
|
547 |
+
hyp = self.hyp
|
548 |
+
mosaic = self.mosaic
|
549 |
+
|
550 |
+
# Load image
|
551 |
+
img, (h0, w0), (h, w) = load_image(self, index)
|
552 |
+
|
553 |
+
# Letterbox
|
554 |
+
shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape
|
555 |
+
img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
|
556 |
+
shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
|
557 |
+
|
558 |
+
labels = self.labels[index].copy()
|
559 |
+
if labels.size: # normalized xywh to pixel xyxy format
|
560 |
+
labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
|
561 |
+
|
562 |
+
nl = len(labels) # number of labels
|
563 |
+
if nl:
|
564 |
+
labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1E-3)
|
565 |
+
|
566 |
+
labels_out = torch.zeros((nl, 6))
|
567 |
+
if nl:
|
568 |
+
labels_out[:, 1:] = torch.from_numpy(labels)
|
569 |
+
|
570 |
+
# Convert
|
571 |
+
img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
|
572 |
+
img = np.ascontiguousarray(img)
|
573 |
+
|
574 |
+
return torch.from_numpy(img), labels_out, self.img_files[index], shapes
|
575 |
+
|
576 |
+
@staticmethod
|
577 |
+
def collate_fn(batch):
|
578 |
+
img, label, path, shapes = zip(*batch) # transposed
|
579 |
+
for i, l in enumerate(label):
|
580 |
+
l[:, 0] = i # add target image index for build_targets()
|
581 |
+
return torch.stack(img, 0), torch.cat(label, 0), path, shapes
|
582 |
+
|
583 |
+
|
584 |
+
def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper)
|
585 |
+
# https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
|
586 |
+
# a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n')
|
587 |
+
# b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
|
588 |
+
# x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco
|
589 |
+
# x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet
|
590 |
+
x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
|
591 |
+
35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
|
592 |
+
64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
|
593 |
+
return x
|
594 |
+
|
595 |
+
|
596 |
+
def check_dataset(data, autodownload=True):
|
597 |
+
# Download and/or unzip dataset if not found locally
|
598 |
+
# Usage: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128_with_yaml.zip
|
599 |
+
|
600 |
+
# Download (optional)
|
601 |
+
extract_dir = ''
|
602 |
+
|
603 |
+
# Read yaml (optional)
|
604 |
+
if isinstance(data, (str, Path)):
|
605 |
+
with open(data, errors='ignore') as f:
|
606 |
+
data = yaml.safe_load(f) # dictionary
|
607 |
+
|
608 |
+
# Parse yaml
|
609 |
+
path = extract_dir or Path(data.get('path') or '') # optional 'path' default to '.'
|
610 |
+
for k in 'train', 'val', 'test':
|
611 |
+
if data.get(k): # prepend path
|
612 |
+
data[k] = str(path / data[k]) if isinstance(data[k], str) else [str(path / x) for x in data[k]]
|
613 |
+
|
614 |
+
assert 'nc' in data, "Dataset 'nc' key missing."
|
615 |
+
if 'names' not in data:
|
616 |
+
data['names'] = [f'class{i}' for i in range(data['nc'])] # assign class names if missing
|
617 |
+
train, val, test, s = [data.get(x) for x in ('train', 'val', 'test', 'download')]
|
618 |
+
if val:
|
619 |
+
val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path
|
620 |
+
if not all(x.exists() for x in val):
|
621 |
+
print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()])
|
622 |
+
|
623 |
+
return data # dictionary
|
624 |
+
|
625 |
+
|
626 |
+
def box_iou(box1, box2):
|
627 |
+
# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
|
628 |
+
"""
|
629 |
+
Return intersection-over-union (Jaccard index) of boxes.
|
630 |
+
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
|
631 |
+
Arguments:
|
632 |
+
box1 (Tensor[N, 4])
|
633 |
+
box2 (Tensor[M, 4])
|
634 |
+
Returns:
|
635 |
+
iou (Tensor[N, M]): the NxM matrix containing the pairwise
|
636 |
+
IoU values for every element in boxes1 and boxes2
|
637 |
+
"""
|
638 |
+
|
639 |
+
def box_area(box):
|
640 |
+
# box = 4xn
|
641 |
+
return (box[2] - box[0]) * (box[3] - box[1])
|
642 |
+
|
643 |
+
area1 = box_area(box1.T)
|
644 |
+
area2 = box_area(box2.T)
|
645 |
+
|
646 |
+
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
|
647 |
+
inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
|
648 |
+
return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)
|
649 |
+
|
650 |
+
|
651 |
+
def increment_path(path, exist_ok=False, sep='', mkdir=False):
|
652 |
+
# Increment file or directory path, i.e. runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
|
653 |
+
path = Path(path) # os-agnostic
|
654 |
+
if path.exists() and not exist_ok:
|
655 |
+
suffix = path.suffix
|
656 |
+
path = path.with_suffix('')
|
657 |
+
dirs = glob.glob(f"{path}{sep}*") # similar paths
|
658 |
+
matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
|
659 |
+
i = [int(m.groups()[0]) for m in matches if m] # indices
|
660 |
+
n = max(i) + 1 if i else 2 # increment number
|
661 |
+
path = Path(f"{path}{sep}{n}{suffix}") # update path
|
662 |
+
dir = path if path.suffix == '' else path.parent # directory
|
663 |
+
if not dir.exists() and mkdir:
|
664 |
+
dir.mkdir(parents=True, exist_ok=True) # make directory
|
665 |
+
return path
|
666 |
+
|
667 |
+
|
668 |
+
def colorstr(*input):
|
669 |
+
# Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e. colorstr('blue', 'hello world')
|
670 |
+
*args, string = input if len(input) > 1 else ('blue', 'bold', input[0]) # color arguments, string
|
671 |
+
colors = {'black': '\033[30m', # basic colors
|
672 |
+
'red': '\033[31m',
|
673 |
+
'green': '\033[32m',
|
674 |
+
'yellow': '\033[33m',
|
675 |
+
'blue': '\033[34m',
|
676 |
+
'magenta': '\033[35m',
|
677 |
+
'cyan': '\033[36m',
|
678 |
+
'white': '\033[37m',
|
679 |
+
'bright_black': '\033[90m', # bright colors
|
680 |
+
'bright_red': '\033[91m',
|
681 |
+
'bright_green': '\033[92m',
|
682 |
+
'bright_yellow': '\033[93m',
|
683 |
+
'bright_blue': '\033[94m',
|
684 |
+
'bright_magenta': '\033[95m',
|
685 |
+
'bright_cyan': '\033[96m',
|
686 |
+
'bright_white': '\033[97m',
|
687 |
+
'end': '\033[0m', # misc
|
688 |
+
'bold': '\033[1m',
|
689 |
+
'underline': '\033[4m'}
|
690 |
+
return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
|
691 |
+
|
692 |
+
|
693 |
+
def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=()):
|
694 |
+
""" Compute the average precision, given the recall and precision curves.
|
695 |
+
Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
|
696 |
+
# Arguments
|
697 |
+
tp: True positives (nparray, nx1 or nx10).
|
698 |
+
conf: Objectness value from 0-1 (nparray).
|
699 |
+
pred_cls: Predicted object classes (nparray).
|
700 |
+
target_cls: True object classes (nparray).
|
701 |
+
plot: Plot precision-recall curve at mAP@0.5
|
702 |
+
save_dir: Plot save directory
|
703 |
+
# Returns
|
704 |
+
The average precision as computed in py-faster-rcnn.
|
705 |
+
"""
|
706 |
+
|
707 |
+
# Sort by objectness
|
708 |
+
i = np.argsort(-conf)
|
709 |
+
tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]
|
710 |
+
|
711 |
+
# Find unique classes
|
712 |
+
unique_classes = np.unique(target_cls)
|
713 |
+
nc = unique_classes.shape[0] # number of classes, number of detections
|
714 |
+
|
715 |
+
# Create Precision-Recall curve and compute AP for each class
|
716 |
+
px, py = np.linspace(0, 1, 1000), [] # for plotting
|
717 |
+
ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000))
|
718 |
+
for ci, c in enumerate(unique_classes):
|
719 |
+
i = pred_cls == c
|
720 |
+
n_l = (target_cls == c).sum() # number of labels
|
721 |
+
n_p = i.sum() # number of predictions
|
722 |
+
|
723 |
+
if n_p == 0 or n_l == 0:
|
724 |
+
continue
|
725 |
+
else:
|
726 |
+
# Accumulate FPs and TPs
|
727 |
+
fpc = (1 - tp[i]).cumsum(0)
|
728 |
+
tpc = tp[i].cumsum(0)
|
729 |
+
|
730 |
+
# Recall
|
731 |
+
recall = tpc / (n_l + 1e-16) # recall curve
|
732 |
+
r[ci] = np.interp(-px, -conf[i], recall[:, 0], left=0) # negative x, xp because xp decreases
|
733 |
+
|
734 |
+
# Precision
|
735 |
+
precision = tpc / (tpc + fpc) # precision curve
|
736 |
+
p[ci] = np.interp(-px, -conf[i], precision[:, 0], left=1) # p at pr_score
|
737 |
+
|
738 |
+
# AP from recall-precision curve
|
739 |
+
for j in range(tp.shape[1]):
|
740 |
+
ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j])
|
741 |
+
if plot and j == 0:
|
742 |
+
py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5
|
743 |
+
|
744 |
+
# Compute F1 (harmonic mean of precision and recall)
|
745 |
+
f1 = 2 * p * r / (p + r + 1e-16)
|
746 |
+
|
747 |
+
i = f1.mean(0).argmax() # max F1 index
|
748 |
+
return p[:, i], r[:, i], ap, f1[:, i], unique_classes.astype('int32')
|
749 |
+
|
750 |
+
|
751 |
+
def compute_ap(recall, precision):
|
752 |
+
""" Compute the average precision, given the recall and precision curves
|
753 |
+
# Arguments
|
754 |
+
recall: The recall curve (list)
|
755 |
+
precision: The precision curve (list)
|
756 |
+
# Returns
|
757 |
+
Average precision, precision curve, recall curve
|
758 |
+
"""
|
759 |
+
|
760 |
+
# Append sentinel values to beginning and end
|
761 |
+
mrec = np.concatenate(([0.0], recall, [1.0]))
|
762 |
+
mpre = np.concatenate(([1.0], precision, [0.0]))
|
763 |
+
|
764 |
+
# Compute the precision envelope
|
765 |
+
mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
|
766 |
+
|
767 |
+
# Integrate area under curve
|
768 |
+
method = 'interp' # methods: 'continuous', 'interp'
|
769 |
+
if method == 'interp':
|
770 |
+
x = np.linspace(0, 1, 101) # 101-point interp (COCO)
|
771 |
+
ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate
|
772 |
+
else: # 'continuous'
|
773 |
+
i = np.where(mrec[1:] != mrec[:-1])[0] # points where x axis (recall) changes
|
774 |
+
ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) # area under curve
|
775 |
+
|
776 |
+
return ap, mpre, mrec
|
777 |
+
|
778 |
+
|
779 |
+
def output_to_target(output):
|
780 |
+
# Convert model output to target format [batch_id, class_id, x, y, w, h, conf]
|
781 |
+
targets = []
|
782 |
+
for i, o in enumerate(output):
|
783 |
+
for *box, conf, cls in o.cpu().numpy():
|
784 |
+
targets.append([i, cls, *list(*xyxy2xywh(np.array(box)[None])), conf])
|
785 |
+
return np.array(targets)
|
786 |
+
|
787 |
+
|
788 |
+
def check_yaml(file, suffix=('.yaml', '.yml')):
|
789 |
+
# Search/download YAML file (if necessary) and return path, checking suffix
|
790 |
+
return check_file(file, suffix)
|
791 |
+
|
792 |
+
|
793 |
+
def check_file(file, suffix=''):
|
794 |
+
# Search/download file (if necessary) and return path
|
795 |
+
check_suffix(file, suffix) # optional
|
796 |
+
file = str(file) # convert to str()
|
797 |
+
return file
|
798 |
+
|
799 |
+
|
800 |
+
def check_suffix(file='yolov5s.pt', suffix=('.pt',), msg=''):
|
801 |
+
# Check file(s) for acceptable suffixes
|
802 |
+
if file and suffix:
|
803 |
+
if isinstance(suffix, str):
|
804 |
+
suffix = [suffix]
|
805 |
+
for f in file if isinstance(file, (list, tuple)) else [file]:
|
806 |
+
assert Path(f).suffix.lower() in suffix, f"{msg}{f} acceptable suffix is {suffix}"
|
807 |
+
|
808 |
+
|
809 |
+
def img2label_paths(img_paths):
|
810 |
+
# Define label paths as a function of image paths
|
811 |
+
sa, sb = os.sep + 'images' + os.sep, os.sep + 'labels' + os.sep # /images/, /labels/ substrings
|
812 |
+
return [sb.join(x.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt' for x in img_paths]
|
813 |
+
|
814 |
+
|
815 |
+
def exif_size(img):
|
816 |
+
# Returns exif-corrected PIL size
|
817 |
+
s = img.size # (width, height)
|
818 |
+
try:
|
819 |
+
rotation = dict(img._getexif().items())[orientation]
|
820 |
+
if rotation == 6: # rotation 270
|
821 |
+
s = (s[1], s[0])
|
822 |
+
elif rotation == 8: # rotation 90
|
823 |
+
s = (s[1], s[0])
|
824 |
+
except:
|
825 |
+
pass
|
826 |
+
|
827 |
+
return s
|
828 |
+
|
829 |
+
|
830 |
+
def verify_image_label(args):
|
831 |
+
# Verify one image-label pair
|
832 |
+
im_file, lb_file, prefix = args
|
833 |
+
nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, '', [] # number (missing, found, empty, corrupt), message, segments
|
834 |
+
try:
|
835 |
+
# verify images
|
836 |
+
im = Image.open(im_file)
|
837 |
+
im.verify() # PIL verify
|
838 |
+
shape = exif_size(im) # image size
|
839 |
+
assert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels'
|
840 |
+
assert im.format.lower() in IMG_FORMATS, f'invalid image format {im.format}'
|
841 |
+
if im.format.lower() in ('jpg', 'jpeg'):
|
842 |
+
with open(im_file, 'rb') as f:
|
843 |
+
f.seek(-2, 2)
|
844 |
+
if f.read() != b'\xff\xd9': # corrupt JPEG
|
845 |
+
Image.open(im_file).save(im_file, format='JPEG', subsampling=0, quality=100) # re-save image
|
846 |
+
msg = f'{prefix}WARNING: corrupt JPEG restored and saved {im_file}'
|
847 |
+
|
848 |
+
# verify labels
|
849 |
+
if os.path.isfile(lb_file):
|
850 |
+
nf = 1 # label found
|
851 |
+
with open(lb_file, 'r') as f:
|
852 |
+
l = [x.split() for x in f.read().strip().splitlines() if len(x)]
|
853 |
+
if any([len(x) > 8 for x in l]): # is segment
|
854 |
+
classes = np.array([x[0] for x in l], dtype=np.float32)
|
855 |
+
segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in l] # (cls, xy1...)
|
856 |
+
l = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1) # (cls, xywh)
|
857 |
+
l = np.array(l, dtype=np.float32)
|
858 |
+
if len(l):
|
859 |
+
assert l.shape[1] == 5, 'labels require 5 columns each'
|
860 |
+
assert (l >= 0).all(), 'negative labels'
|
861 |
+
assert (l[:, 1:] <= 1).all(), 'non-normalized or out of bounds coordinate labels'
|
862 |
+
assert np.unique(l, axis=0).shape[0] == l.shape[0], 'duplicate labels'
|
863 |
+
else:
|
864 |
+
ne = 1 # label empty
|
865 |
+
l = np.zeros((0, 5), dtype=np.float32)
|
866 |
+
else:
|
867 |
+
nm = 1 # label missing
|
868 |
+
l = np.zeros((0, 5), dtype=np.float32)
|
869 |
+
return im_file, l, shape, segments, nm, nf, ne, nc, msg
|
870 |
+
except Exception as e:
|
871 |
+
nc = 1
|
872 |
+
msg = f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}'
|
873 |
+
return [None, None, None, None, nm, nf, ne, nc, msg]
|
874 |
+
|
875 |
+
|
876 |
+
def segments2boxes(segments):
|
877 |
+
# Convert segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh)
|
878 |
+
boxes = []
|
879 |
+
for s in segments:
|
880 |
+
x, y = s.T # segment xy
|
881 |
+
boxes.append([x.min(), y.min(), x.max(), y.max()]) # cls, xyxy
|
882 |
+
return xyxy2xywh(np.array(boxes)) # cls, xywh
|
883 |
+
|
884 |
+
|
885 |
+
def get_hash(paths):
|
886 |
+
# Returns a single hash value of a list of paths (files or dirs)
|
887 |
+
size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) # sizes
|
888 |
+
h = hashlib.md5(str(size).encode()) # hash sizes
|
889 |
+
h.update(''.join(paths).encode()) # hash paths
|
890 |
+
return h.hexdigest() # return hash
|
891 |
+
|
892 |
+
|
893 |
+
class InfiniteDataLoader(torch.utils.data.dataloader.DataLoader):
|
894 |
+
""" Dataloader that reuses workers
|
895 |
+
|
896 |
+
Uses same syntax as vanilla DataLoader
|
897 |
+
"""
|
898 |
+
|
899 |
+
def __init__(self, *args, **kwargs):
|
900 |
+
super().__init__(*args, **kwargs)
|
901 |
+
object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))
|
902 |
+
self.iterator = super().__iter__()
|
903 |
+
|
904 |
+
def __len__(self):
|
905 |
+
return len(self.batch_sampler.sampler)
|
906 |
+
|
907 |
+
def __iter__(self):
|
908 |
+
for i in range(len(self)):
|
909 |
+
yield next(self.iterator)
|
910 |
+
|
911 |
+
|
912 |
+
class _RepeatSampler(object):
|
913 |
+
""" Sampler that repeats forever
|
914 |
+
|
915 |
+
Args:
|
916 |
+
sampler (Sampler)
|
917 |
+
"""
|
918 |
+
|
919 |
+
def __init__(self, sampler):
|
920 |
+
self.sampler = sampler
|
921 |
+
|
922 |
+
def __iter__(self):
|
923 |
+
while True:
|
924 |
+
yield from iter(self.sampler)
|
925 |
+
|
926 |
+
|
927 |
+
def load_image(self, i):
|
928 |
+
# loads 1 image from dataset index 'i', returns im, original hw, resized hw
|
929 |
+
im = self.imgs[i]
|
930 |
+
if im is None: # not cached in ram
|
931 |
+
npy = self.img_npy[i]
|
932 |
+
if npy and npy.exists(): # load npy
|
933 |
+
im = np.load(npy)
|
934 |
+
else: # read image
|
935 |
+
path = self.img_files[i]
|
936 |
+
im = cv2.imread(path) # BGR
|
937 |
+
assert im is not None, 'Image Not Found ' + path
|
938 |
+
h0, w0 = im.shape[:2] # orig hw
|
939 |
+
r = self.img_size / max(h0, w0) # ratio
|
940 |
+
if r != 1: # if sizes are not equal
|
941 |
+
im = cv2.resize(im, (int(w0 * r), int(h0 * r)),
|
942 |
+
interpolation=cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR)
|
943 |
+
return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized
|
944 |
+
else:
|
945 |
+
return self.imgs[i], self.img_hw0[i], self.img_hw[i] # im, hw_original, hw_resized
|
946 |
+
|
947 |
+
|
948 |
+
def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
|
949 |
+
# Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
|
950 |
+
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
951 |
+
y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw # top left x
|
952 |
+
y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh # top left y
|
953 |
+
y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw # bottom right x
|
954 |
+
y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh # bottom right y
|
955 |
+
return y
|
956 |
+
|
957 |
+
|
958 |
+
def xyxy2xywhn(x, w=640, h=640, clip=False, eps=0.0):
|
959 |
+
# Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] normalized where xy1=top-left, xy2=bottom-right
|
960 |
+
if clip:
|
961 |
+
clip_coords(x, (h - eps, w - eps)) # warning: inplace clip
|
962 |
+
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
963 |
+
y[:, 0] = ((x[:, 0] + x[:, 2]) / 2) / w # x center
|
964 |
+
y[:, 1] = ((x[:, 1] + x[:, 3]) / 2) / h # y center
|
965 |
+
y[:, 2] = (x[:, 2] - x[:, 0]) / w # width
|
966 |
+
y[:, 3] = (x[:, 3] - x[:, 1]) / h # height
|
967 |
+
return y
|
968 |
+
|
969 |
+
|
970 |
+
def post_process(x):
|
971 |
+
grid = np.load("./grid.npy", allow_pickle=True)
|
972 |
+
anchor_grid = np.load("./anchor_grid.npy", allow_pickle=True)
|
973 |
+
x = list(x)
|
974 |
+
z = [] # inference output
|
975 |
+
stride = [8, 16, 32]
|
976 |
+
for i in range(3):
|
977 |
+
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
|
978 |
+
x[i] = (
|
979 |
+
torch.tensor(x[i])
|
980 |
+
.view(bs, 3, 85, ny, nx)
|
981 |
+
.permute(0, 1, 3, 4, 2)
|
982 |
+
.contiguous()
|
983 |
+
)
|
984 |
+
y = x[i].sigmoid()
|
985 |
+
xy = (y[..., 0:2] * 2.0 - 0.5 + grid[i]) * stride[i]
|
986 |
+
wh = (y[..., 2:4] * 2) ** 2 * anchor_grid[i]
|
987 |
+
y = torch.cat((xy, wh, y[..., 4:]), -1)
|
988 |
+
z.append(y.view(bs, -1, 85))
|
989 |
+
|
990 |
+
return (torch.cat(z, 1), x)
|
yolov5s_qat.onnx
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:5ba00d5f170eab6130610bb543c1f4b1e8354f4944c127e61c28beb99beddf26
|
3 |
+
size 29141657
|