yolo9 0.0.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- yolo9-0.0.1/PKG-INFO +24 -0
- yolo9-0.0.1/pyproject.toml +37 -0
- yolo9-0.0.1/setup.cfg +4 -0
- yolo9-0.0.1/src/yolo9/__init__.py +75 -0
- yolo9-0.0.1/src/yolo9/benchmarks.py +142 -0
- yolo9-0.0.1/src/yolo9/classify/predict.py +224 -0
- yolo9-0.0.1/src/yolo9/classify/train.py +333 -0
- yolo9-0.0.1/src/yolo9/classify/val.py +170 -0
- yolo9-0.0.1/src/yolo9/data/coco.yaml +125 -0
- yolo9-0.0.1/src/yolo9/detect.py +231 -0
- yolo9-0.0.1/src/yolo9/detect_dual.py +232 -0
- yolo9-0.0.1/src/yolo9/export.py +686 -0
- yolo9-0.0.1/src/yolo9/hubconf.py +107 -0
- yolo9-0.0.1/src/yolo9/models/__init__.py +1 -0
- yolo9-0.0.1/src/yolo9/models/common.py +1233 -0
- yolo9-0.0.1/src/yolo9/models/detect/gelan-c.yaml +80 -0
- yolo9-0.0.1/src/yolo9/models/detect/gelan-e.yaml +121 -0
- yolo9-0.0.1/src/yolo9/models/detect/gelan-m.yaml +80 -0
- yolo9-0.0.1/src/yolo9/models/detect/gelan-s.yaml +80 -0
- yolo9-0.0.1/src/yolo9/models/detect/gelan-t.yaml +80 -0
- yolo9-0.0.1/src/yolo9/models/detect/gelan.yaml +80 -0
- yolo9-0.0.1/src/yolo9/models/detect/yolov7-af.yaml +137 -0
- yolo9-0.0.1/src/yolo9/models/detect/yolov9-c.yaml +124 -0
- yolo9-0.0.1/src/yolo9/models/detect/yolov9-cf.yaml +124 -0
- yolo9-0.0.1/src/yolo9/models/detect/yolov9-e.yaml +144 -0
- yolo9-0.0.1/src/yolo9/models/detect/yolov9-m.yaml +117 -0
- yolo9-0.0.1/src/yolo9/models/detect/yolov9-s.yaml +97 -0
- yolo9-0.0.1/src/yolo9/models/detect/yolov9-t.yaml +97 -0
- yolo9-0.0.1/src/yolo9/models/detect/yolov9.yaml +117 -0
- yolo9-0.0.1/src/yolo9/models/experimental.py +274 -0
- yolo9-0.0.1/src/yolo9/models/hub/anchors.yaml +59 -0
- yolo9-0.0.1/src/yolo9/models/hub/yolov3-spp.yaml +51 -0
- yolo9-0.0.1/src/yolo9/models/hub/yolov3-tiny.yaml +41 -0
- yolo9-0.0.1/src/yolo9/models/hub/yolov3.yaml +51 -0
- yolo9-0.0.1/src/yolo9/models/panoptic/gelan-c-pan.yaml +80 -0
- yolo9-0.0.1/src/yolo9/models/panoptic/yolov7-af-pan.yaml +137 -0
- yolo9-0.0.1/src/yolo9/models/segment/gelan-c-dseg.yaml +84 -0
- yolo9-0.0.1/src/yolo9/models/segment/gelan-c-seg.yaml +80 -0
- yolo9-0.0.1/src/yolo9/models/segment/yolov7-af-seg.yaml +136 -0
- yolo9-0.0.1/src/yolo9/models/segment/yolov9-c-dseg.yaml +130 -0
- yolo9-0.0.1/src/yolo9/models/tf.py +596 -0
- yolo9-0.0.1/src/yolo9/models/yolo.py +818 -0
- yolo9-0.0.1/src/yolo9/panoptic/predict.py +246 -0
- yolo9-0.0.1/src/yolo9/panoptic/train.py +662 -0
- yolo9-0.0.1/src/yolo9/panoptic/val.py +597 -0
- yolo9-0.0.1/src/yolo9/segment/predict.py +246 -0
- yolo9-0.0.1/src/yolo9/segment/train.py +646 -0
- yolo9-0.0.1/src/yolo9/segment/train_dual.py +647 -0
- yolo9-0.0.1/src/yolo9/segment/val.py +457 -0
- yolo9-0.0.1/src/yolo9/segment/val_dual.py +458 -0
- yolo9-0.0.1/src/yolo9/train.py +634 -0
- yolo9-0.0.1/src/yolo9/train_dual.py +644 -0
- yolo9-0.0.1/src/yolo9/train_triple.py +636 -0
- yolo9-0.0.1/src/yolo9/utils/__init__.py +75 -0
- yolo9-0.0.1/src/yolo9/utils/activations.py +98 -0
- yolo9-0.0.1/src/yolo9/utils/augmentations.py +395 -0
- yolo9-0.0.1/src/yolo9/utils/autoanchor.py +164 -0
- yolo9-0.0.1/src/yolo9/utils/autobatch.py +67 -0
- yolo9-0.0.1/src/yolo9/utils/callbacks.py +71 -0
- yolo9-0.0.1/src/yolo9/utils/coco_utils.py +108 -0
- yolo9-0.0.1/src/yolo9/utils/dataloaders.py +1217 -0
- yolo9-0.0.1/src/yolo9/utils/downloads.py +103 -0
- yolo9-0.0.1/src/yolo9/utils/general.py +1135 -0
- yolo9-0.0.1/src/yolo9/utils/lion.py +67 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/__init__.py +399 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/clearml/__init__.py +1 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/clearml/clearml_utils.py +157 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/clearml/hpo.py +84 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/comet/__init__.py +508 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/comet/comet_utils.py +150 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/comet/hpo.py +118 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/comet/optimizer_config.json +209 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/wandb/__init__.py +1 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/wandb/log_dataset.py +27 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/wandb/sweep.py +41 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/wandb/sweep.yaml +143 -0
- yolo9-0.0.1/src/yolo9/utils/loggers/wandb/wandb_utils.py +589 -0
- yolo9-0.0.1/src/yolo9/utils/loss.py +363 -0
- yolo9-0.0.1/src/yolo9/utils/loss_tal.py +215 -0
- yolo9-0.0.1/src/yolo9/utils/loss_tal_dual.py +385 -0
- yolo9-0.0.1/src/yolo9/utils/loss_tal_triple.py +282 -0
- yolo9-0.0.1/src/yolo9/utils/metrics.py +397 -0
- yolo9-0.0.1/src/yolo9/utils/panoptic/__init__.py +1 -0
- yolo9-0.0.1/src/yolo9/utils/panoptic/augmentations.py +183 -0
- yolo9-0.0.1/src/yolo9/utils/panoptic/dataloaders.py +478 -0
- yolo9-0.0.1/src/yolo9/utils/panoptic/general.py +137 -0
- yolo9-0.0.1/src/yolo9/utils/panoptic/loss.py +186 -0
- yolo9-0.0.1/src/yolo9/utils/panoptic/loss_tal.py +285 -0
- yolo9-0.0.1/src/yolo9/utils/panoptic/metrics.py +272 -0
- yolo9-0.0.1/src/yolo9/utils/panoptic/plots.py +164 -0
- yolo9-0.0.1/src/yolo9/utils/panoptic/tal/__init__.py +1 -0
- yolo9-0.0.1/src/yolo9/utils/panoptic/tal/anchor_generator.py +38 -0
- yolo9-0.0.1/src/yolo9/utils/panoptic/tal/assigner.py +181 -0
- yolo9-0.0.1/src/yolo9/utils/plots.py +570 -0
- yolo9-0.0.1/src/yolo9/utils/segment/__init__.py +1 -0
- yolo9-0.0.1/src/yolo9/utils/segment/augmentations.py +99 -0
- yolo9-0.0.1/src/yolo9/utils/segment/dataloaders.py +328 -0
- yolo9-0.0.1/src/yolo9/utils/segment/general.py +137 -0
- yolo9-0.0.1/src/yolo9/utils/segment/loss.py +186 -0
- yolo9-0.0.1/src/yolo9/utils/segment/loss_tal.py +261 -0
- yolo9-0.0.1/src/yolo9/utils/segment/loss_tal_dual.py +727 -0
- yolo9-0.0.1/src/yolo9/utils/segment/metrics.py +205 -0
- yolo9-0.0.1/src/yolo9/utils/segment/plots.py +143 -0
- yolo9-0.0.1/src/yolo9/utils/segment/tal/__init__.py +1 -0
- yolo9-0.0.1/src/yolo9/utils/segment/tal/anchor_generator.py +38 -0
- yolo9-0.0.1/src/yolo9/utils/segment/tal/assigner.py +180 -0
- yolo9-0.0.1/src/yolo9/utils/tal/__init__.py +1 -0
- yolo9-0.0.1/src/yolo9/utils/tal/anchor_generator.py +38 -0
- yolo9-0.0.1/src/yolo9/utils/tal/assigner.py +179 -0
- yolo9-0.0.1/src/yolo9/utils/torch_utils.py +529 -0
- yolo9-0.0.1/src/yolo9/utils/triton.py +81 -0
- yolo9-0.0.1/src/yolo9/val.py +389 -0
- yolo9-0.0.1/src/yolo9/val_dual.py +393 -0
- yolo9-0.0.1/src/yolo9/val_triple.py +391 -0
- yolo9-0.0.1/src/yolo9.egg-info/PKG-INFO +24 -0
- yolo9-0.0.1/src/yolo9.egg-info/SOURCES.txt +117 -0
- yolo9-0.0.1/src/yolo9.egg-info/dependency_links.txt +1 -0
- yolo9-0.0.1/src/yolo9.egg-info/requires.txt +18 -0
- yolo9-0.0.1/src/yolo9.egg-info/top_level.txt +1 -0
yolo9-0.0.1/PKG-INFO
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: yolo9
|
|
3
|
+
Version: 0.0.1
|
|
4
|
+
Summary: Package version of WongKinYiu/yolov9
|
|
5
|
+
Requires-Python: >=3.11
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
Requires-Dist: gitpython
|
|
8
|
+
Requires-Dist: ipython
|
|
9
|
+
Requires-Dist: matplotlib>=3.2.2
|
|
10
|
+
Requires-Dist: numpy>=1.18.5
|
|
11
|
+
Requires-Dist: opencv-python>=4.1.1
|
|
12
|
+
Requires-Dist: Pillow>=7.1.2
|
|
13
|
+
Requires-Dist: psutil
|
|
14
|
+
Requires-Dist: PyYAML>=5.3.1
|
|
15
|
+
Requires-Dist: requests>=2.23.0
|
|
16
|
+
Requires-Dist: scipy>=1.4.1
|
|
17
|
+
Requires-Dist: thop>=0.1.1
|
|
18
|
+
Requires-Dist: torch>=1.7.0
|
|
19
|
+
Requires-Dist: torchvision>=0.8.1
|
|
20
|
+
Requires-Dist: tqdm>=4.64.0
|
|
21
|
+
Requires-Dist: pandas>=1.1.4
|
|
22
|
+
Requires-Dist: seaborn>=0.11.0
|
|
23
|
+
Requires-Dist: albumentations>=1.0.3
|
|
24
|
+
Requires-Dist: pycocotools>=2.0
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "yolo9"
|
|
7
|
+
version = "0.0.1"
|
|
8
|
+
description = "Package version of WongKinYiu/yolov9"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.11"
|
|
11
|
+
dependencies = [
|
|
12
|
+
"gitpython",
|
|
13
|
+
"ipython",
|
|
14
|
+
"matplotlib>=3.2.2",
|
|
15
|
+
"numpy>=1.18.5",
|
|
16
|
+
"opencv-python>=4.1.1",
|
|
17
|
+
"Pillow>=7.1.2",
|
|
18
|
+
"psutil",
|
|
19
|
+
"PyYAML>=5.3.1",
|
|
20
|
+
"requests>=2.23.0",
|
|
21
|
+
"scipy>=1.4.1",
|
|
22
|
+
"thop>=0.1.1",
|
|
23
|
+
"torch>=1.7.0",
|
|
24
|
+
"torchvision>=0.8.1",
|
|
25
|
+
"tqdm>=4.64.0",
|
|
26
|
+
"pandas>=1.1.4",
|
|
27
|
+
"seaborn>=0.11.0",
|
|
28
|
+
"albumentations>=1.0.3",
|
|
29
|
+
"pycocotools>=2.0",
|
|
30
|
+
]
|
|
31
|
+
|
|
32
|
+
[tool.setuptools.packages.find]
|
|
33
|
+
where = ["src"]
|
|
34
|
+
include = ["*"]
|
|
35
|
+
|
|
36
|
+
[tool.setuptools.package-data]
|
|
37
|
+
"*" = ["*"]
|
yolo9-0.0.1/setup.cfg
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
from pathlib import Path
|
|
2
|
+
from typing import Dict, Optional
|
|
3
|
+
|
|
4
|
+
import numpy as np
|
|
5
|
+
import torch
|
|
6
|
+
import cv2
|
|
7
|
+
from yolo9.models.common import DetectMultiBackend
|
|
8
|
+
from yolo9.utils.general import LOGGER, non_max_suppression
|
|
9
|
+
from yolo9.utils.torch_utils import select_device
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class YOLO9:
|
|
13
|
+
def __init__(
|
|
14
|
+
self,
|
|
15
|
+
model_name: str,
|
|
16
|
+
device: str,
|
|
17
|
+
dnn: bool = False,
|
|
18
|
+
half: bool = False,
|
|
19
|
+
batch_size: int = 1, # batch size
|
|
20
|
+
conf_thres: float = 0.25,
|
|
21
|
+
iou_thres: float = 0.45,
|
|
22
|
+
max_det: int = 1000,
|
|
23
|
+
classes: Optional[Dict[int, float]] = None, # class id -> confidence threshold
|
|
24
|
+
):
|
|
25
|
+
weights_dir = Path(__file__).parent / 'weights'
|
|
26
|
+
weights_dir.mkdir(exist_ok=True)
|
|
27
|
+
data = Path(__file__).parent / 'data' / 'coco.yaml'
|
|
28
|
+
|
|
29
|
+
self.weights_path = weights_dir / f'{model_name}.pt'
|
|
30
|
+
self.device = select_device(device)
|
|
31
|
+
self.conf_thres = conf_thres
|
|
32
|
+
self.iou_thres = iou_thres
|
|
33
|
+
self.max_det = max_det
|
|
34
|
+
self.classes = classes
|
|
35
|
+
if not self.weights_path.exists():
|
|
36
|
+
LOGGER.info(f"Downloading {self.weights_path.name} from GitHub releases...")
|
|
37
|
+
import requests
|
|
38
|
+
url = f"https://github.com/WongKinYiu/yolov9/releases/download/v0.1/{model_name}.pt"
|
|
39
|
+
response = requests.get(url, stream=True)
|
|
40
|
+
response.raise_for_status()
|
|
41
|
+
|
|
42
|
+
with open(self.weights_path, 'wb') as f:
|
|
43
|
+
for chunk in response.iter_content(chunk_size=8192):
|
|
44
|
+
f.write(chunk)
|
|
45
|
+
LOGGER.info(f"Downloaded {self.weights_path.name} successfully")
|
|
46
|
+
|
|
47
|
+
self.model = DetectMultiBackend(self.weights_path, device=self.device, dnn=dnn, data=data, fp16=half)
|
|
48
|
+
self.img_size = 640
|
|
49
|
+
self.model.warmup(imgsz=(1 if self.model.pt or self.model.triton else batch_size, 3, self.img_size, self.img_size))
|
|
50
|
+
|
|
51
|
+
def detect(self, image: np.ndarray) -> list:
|
|
52
|
+
im = cv2.resize(image, (self.img_size, self.img_size))
|
|
53
|
+
im = torch.from_numpy(im).to(self.device)
|
|
54
|
+
im = im.half() if self.model.fp16 else im.float() # uint8 to fp16/32
|
|
55
|
+
im /= 255 # 0-255 to 0-1
|
|
56
|
+
if len(im.shape) == 3:
|
|
57
|
+
im = im[None] # expand for batch dim
|
|
58
|
+
im = im.permute(0, 3, 1, 2) # BHWC to BCHW
|
|
59
|
+
|
|
60
|
+
pred = self.model(im)
|
|
61
|
+
pred = non_max_suppression(pred, conf_thres=self.conf_thres, iou_thres=self.iou_thres, max_det=self.max_det)[0]
|
|
62
|
+
|
|
63
|
+
detections = []
|
|
64
|
+
for det in pred:
|
|
65
|
+
confidence, class_id = det[4].tolist(), int(det[5].tolist())
|
|
66
|
+
if isinstance(self.classes, dict):
|
|
67
|
+
if class_id not in self.classes:
|
|
68
|
+
continue
|
|
69
|
+
if confidence < self.classes[class_id]:
|
|
70
|
+
continue
|
|
71
|
+
x1, y1, x2, y2 = (det[0:4] / self.img_size).tolist()
|
|
72
|
+
polygon = [(x1, y1), (x2, y1), (x2, y2), (x1, y2)]
|
|
73
|
+
detections.append((polygon, confidence, class_id))
|
|
74
|
+
|
|
75
|
+
return detections
|
|
@@ -0,0 +1,142 @@
|
|
|
1
|
+
import argparse
|
|
2
|
+
import platform
|
|
3
|
+
import sys
|
|
4
|
+
import time
|
|
5
|
+
from pathlib import Path
|
|
6
|
+
|
|
7
|
+
import pandas as pd
|
|
8
|
+
|
|
9
|
+
FILE = Path(__file__).resolve()
|
|
10
|
+
ROOT = FILE.parents[0] # YOLO root directory
|
|
11
|
+
if str(ROOT) not in sys.path:
|
|
12
|
+
sys.path.append(str(ROOT)) # add ROOT to PATH
|
|
13
|
+
# ROOT = ROOT.relative_to(Path.cwd()) # relative
|
|
14
|
+
|
|
15
|
+
import export
|
|
16
|
+
from yolo9.models.experimental import attempt_load
|
|
17
|
+
from yolo9.models.yolo import SegmentationModel
|
|
18
|
+
from yolo9.segment.val import run as val_seg
|
|
19
|
+
from yolo9.utils import notebook_init
|
|
20
|
+
from yolo9.utils.general import LOGGER, check_yaml, file_size, print_args
|
|
21
|
+
from yolo9.utils.torch_utils import select_device
|
|
22
|
+
from yolo9.val import run as val_det
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def run(
|
|
26
|
+
weights=ROOT / 'yolo.pt', # weights path
|
|
27
|
+
imgsz=640, # inference size (pixels)
|
|
28
|
+
batch_size=1, # batch size
|
|
29
|
+
data=ROOT / 'data/coco.yaml', # dataset.yaml path
|
|
30
|
+
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
|
31
|
+
half=False, # use FP16 half-precision inference
|
|
32
|
+
test=False, # test exports only
|
|
33
|
+
pt_only=False, # test PyTorch only
|
|
34
|
+
hard_fail=False, # throw error on benchmark failure
|
|
35
|
+
):
|
|
36
|
+
y, t = [], time.time()
|
|
37
|
+
device = select_device(device)
|
|
38
|
+
model_type = type(attempt_load(weights, fuse=False)) # DetectionModel, SegmentationModel, etc.
|
|
39
|
+
for i, (name, f, suffix, cpu, gpu) in export.export_formats().iterrows(): # index, (name, file, suffix, CPU, GPU)
|
|
40
|
+
try:
|
|
41
|
+
assert i not in (9, 10), 'inference not supported' # Edge TPU and TF.js are unsupported
|
|
42
|
+
assert i != 5 or platform.system() == 'Darwin', 'inference only supported on macOS>=10.13' # CoreML
|
|
43
|
+
if 'cpu' in device.type:
|
|
44
|
+
assert cpu, 'inference not supported on CPU'
|
|
45
|
+
if 'cuda' in device.type:
|
|
46
|
+
assert gpu, 'inference not supported on GPU'
|
|
47
|
+
|
|
48
|
+
# Export
|
|
49
|
+
if f == '-':
|
|
50
|
+
w = weights # PyTorch format
|
|
51
|
+
else:
|
|
52
|
+
w = export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1] # all others
|
|
53
|
+
assert suffix in str(w), 'export failed'
|
|
54
|
+
|
|
55
|
+
# Validate
|
|
56
|
+
if model_type == SegmentationModel:
|
|
57
|
+
result = val_seg(data, w, batch_size, imgsz, plots=False, device=device, task='speed', half=half)
|
|
58
|
+
metric = result[0][7] # (box(p, r, map50, map), mask(p, r, map50, map), *loss(box, obj, cls))
|
|
59
|
+
else: # DetectionModel:
|
|
60
|
+
result = val_det(data, w, batch_size, imgsz, plots=False, device=device, task='speed', half=half)
|
|
61
|
+
metric = result[0][3] # (p, r, map50, map, *loss(box, obj, cls))
|
|
62
|
+
speed = result[2][1] # times (preprocess, inference, postprocess)
|
|
63
|
+
y.append([name, round(file_size(w), 1), round(metric, 4), round(speed, 2)]) # MB, mAP, t_inference
|
|
64
|
+
except Exception as e:
|
|
65
|
+
if hard_fail:
|
|
66
|
+
assert type(e) is AssertionError, f'Benchmark --hard-fail for {name}: {e}'
|
|
67
|
+
LOGGER.warning(f'WARNING ⚠️ Benchmark failure for {name}: {e}')
|
|
68
|
+
y.append([name, None, None, None]) # mAP, t_inference
|
|
69
|
+
if pt_only and i == 0:
|
|
70
|
+
break # break after PyTorch
|
|
71
|
+
|
|
72
|
+
# Print results
|
|
73
|
+
LOGGER.info('\n')
|
|
74
|
+
parse_opt()
|
|
75
|
+
notebook_init() # print system info
|
|
76
|
+
c = ['Format', 'Size (MB)', 'mAP50-95', 'Inference time (ms)'] if map else ['Format', 'Export', '', '']
|
|
77
|
+
py = pd.DataFrame(y, columns=c)
|
|
78
|
+
LOGGER.info(f'\nBenchmarks complete ({time.time() - t:.2f}s)')
|
|
79
|
+
LOGGER.info(str(py if map else py.iloc[:, :2]))
|
|
80
|
+
if hard_fail and isinstance(hard_fail, str):
|
|
81
|
+
metrics = py['mAP50-95'].array # values to compare to floor
|
|
82
|
+
floor = eval(hard_fail) # minimum metric floor to pass
|
|
83
|
+
assert all(x > floor for x in metrics if pd.notna(x)), f'HARD FAIL: mAP50-95 < floor {floor}'
|
|
84
|
+
return py
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def test(
|
|
88
|
+
weights=ROOT / 'yolo.pt', # weights path
|
|
89
|
+
imgsz=640, # inference size (pixels)
|
|
90
|
+
batch_size=1, # batch size
|
|
91
|
+
data=ROOT / 'data/coco128.yaml', # dataset.yaml path
|
|
92
|
+
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
|
93
|
+
half=False, # use FP16 half-precision inference
|
|
94
|
+
test=False, # test exports only
|
|
95
|
+
pt_only=False, # test PyTorch only
|
|
96
|
+
hard_fail=False, # throw error on benchmark failure
|
|
97
|
+
):
|
|
98
|
+
y, t = [], time.time()
|
|
99
|
+
device = select_device(device)
|
|
100
|
+
for i, (name, f, suffix, gpu) in export.export_formats().iterrows(): # index, (name, file, suffix, gpu-capable)
|
|
101
|
+
try:
|
|
102
|
+
w = weights if f == '-' else \
|
|
103
|
+
export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1] # weights
|
|
104
|
+
assert suffix in str(w), 'export failed'
|
|
105
|
+
y.append([name, True])
|
|
106
|
+
except Exception:
|
|
107
|
+
y.append([name, False]) # mAP, t_inference
|
|
108
|
+
|
|
109
|
+
# Print results
|
|
110
|
+
LOGGER.info('\n')
|
|
111
|
+
parse_opt()
|
|
112
|
+
notebook_init() # print system info
|
|
113
|
+
py = pd.DataFrame(y, columns=['Format', 'Export'])
|
|
114
|
+
LOGGER.info(f'\nExports complete ({time.time() - t:.2f}s)')
|
|
115
|
+
LOGGER.info(str(py))
|
|
116
|
+
return py
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
def parse_opt():
|
|
120
|
+
parser = argparse.ArgumentParser()
|
|
121
|
+
parser.add_argument('--weights', type=str, default=ROOT / 'yolo.pt', help='weights path')
|
|
122
|
+
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
|
|
123
|
+
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
|
|
124
|
+
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
|
|
125
|
+
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
|
126
|
+
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
|
|
127
|
+
parser.add_argument('--test', action='store_true', help='test exports only')
|
|
128
|
+
parser.add_argument('--pt-only', action='store_true', help='test PyTorch only')
|
|
129
|
+
parser.add_argument('--hard-fail', nargs='?', const=True, default=False, help='Exception on error or < min metric')
|
|
130
|
+
opt = parser.parse_args()
|
|
131
|
+
opt.data = check_yaml(opt.data) # check YAML
|
|
132
|
+
print_args(vars(opt))
|
|
133
|
+
return opt
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
def main(opt):
|
|
137
|
+
test(**vars(opt)) if opt.test else run(**vars(opt))
|
|
138
|
+
|
|
139
|
+
|
|
140
|
+
if __name__ == "__main__":
|
|
141
|
+
opt = parse_opt()
|
|
142
|
+
main(opt)
|
|
@@ -0,0 +1,224 @@
|
|
|
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
|
+
screen # screenshot
|
|
10
|
+
path/ # directory
|
|
11
|
+
'path/*.jpg' # glob
|
|
12
|
+
'https://youtu.be/Zgi9g1ksQHc' # YouTube
|
|
13
|
+
'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
|
|
14
|
+
|
|
15
|
+
Usage - formats:
|
|
16
|
+
$ python classify/predict.py --weights yolov5s-cls.pt # PyTorch
|
|
17
|
+
yolov5s-cls.torchscript # TorchScript
|
|
18
|
+
yolov5s-cls.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
|
19
|
+
yolov5s-cls_openvino_model # OpenVINO
|
|
20
|
+
yolov5s-cls.engine # TensorRT
|
|
21
|
+
yolov5s-cls.mlmodel # CoreML (macOS-only)
|
|
22
|
+
yolov5s-cls_saved_model # TensorFlow SavedModel
|
|
23
|
+
yolov5s-cls.pb # TensorFlow GraphDef
|
|
24
|
+
yolov5s-cls.tflite # TensorFlow Lite
|
|
25
|
+
yolov5s-cls_edgetpu.tflite # TensorFlow Edge TPU
|
|
26
|
+
yolov5s-cls_paddle_model # PaddlePaddle
|
|
27
|
+
"""
|
|
28
|
+
|
|
29
|
+
import argparse
|
|
30
|
+
import os
|
|
31
|
+
import platform
|
|
32
|
+
import sys
|
|
33
|
+
from pathlib import Path
|
|
34
|
+
|
|
35
|
+
import torch
|
|
36
|
+
import torch.nn.functional as F
|
|
37
|
+
|
|
38
|
+
FILE = Path(__file__).resolve()
|
|
39
|
+
ROOT = FILE.parents[1] # YOLOv5 root directory
|
|
40
|
+
if str(ROOT) not in sys.path:
|
|
41
|
+
sys.path.append(str(ROOT)) # add ROOT to PATH
|
|
42
|
+
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
|
43
|
+
|
|
44
|
+
from models.common import DetectMultiBackend
|
|
45
|
+
from utils.augmentations import classify_transforms
|
|
46
|
+
from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
|
|
47
|
+
from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
|
|
48
|
+
increment_path, print_args, strip_optimizer)
|
|
49
|
+
from utils.plots import Annotator
|
|
50
|
+
from utils.torch_utils import select_device, smart_inference_mode
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
@smart_inference_mode()
|
|
54
|
+
def run(
|
|
55
|
+
weights=ROOT / 'yolov5s-cls.pt', # model.pt path(s)
|
|
56
|
+
source=ROOT / 'data/images', # file/dir/URL/glob/screen/0(webcam)
|
|
57
|
+
data=ROOT / 'data/coco128.yaml', # dataset.yaml path
|
|
58
|
+
imgsz=(224, 224), # inference size (height, width)
|
|
59
|
+
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
|
60
|
+
view_img=False, # show results
|
|
61
|
+
save_txt=False, # save results to *.txt
|
|
62
|
+
nosave=False, # do not save images/videos
|
|
63
|
+
augment=False, # augmented inference
|
|
64
|
+
visualize=False, # visualize features
|
|
65
|
+
update=False, # update all models
|
|
66
|
+
project=ROOT / 'runs/predict-cls', # save results to project/name
|
|
67
|
+
name='exp', # save results to project/name
|
|
68
|
+
exist_ok=False, # existing project/name ok, do not increment
|
|
69
|
+
half=False, # use FP16 half-precision inference
|
|
70
|
+
dnn=False, # use OpenCV DNN for ONNX inference
|
|
71
|
+
vid_stride=1, # video frame-rate stride
|
|
72
|
+
):
|
|
73
|
+
source = str(source)
|
|
74
|
+
save_img = not nosave and not source.endswith('.txt') # save inference images
|
|
75
|
+
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
|
|
76
|
+
is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
|
|
77
|
+
webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
|
|
78
|
+
screenshot = source.lower().startswith('screen')
|
|
79
|
+
if is_url and is_file:
|
|
80
|
+
source = check_file(source) # download
|
|
81
|
+
|
|
82
|
+
# Directories
|
|
83
|
+
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
|
|
84
|
+
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
|
|
85
|
+
|
|
86
|
+
# Load model
|
|
87
|
+
device = select_device(device)
|
|
88
|
+
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
|
|
89
|
+
stride, names, pt = model.stride, model.names, model.pt
|
|
90
|
+
imgsz = check_img_size(imgsz, s=stride) # check image size
|
|
91
|
+
|
|
92
|
+
# Dataloader
|
|
93
|
+
bs = 1 # batch_size
|
|
94
|
+
if webcam:
|
|
95
|
+
view_img = check_imshow(warn=True)
|
|
96
|
+
dataset = LoadStreams(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride)
|
|
97
|
+
bs = len(dataset)
|
|
98
|
+
elif screenshot:
|
|
99
|
+
dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
|
|
100
|
+
else:
|
|
101
|
+
dataset = LoadImages(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride)
|
|
102
|
+
vid_path, vid_writer = [None] * bs, [None] * bs
|
|
103
|
+
|
|
104
|
+
# Run inference
|
|
105
|
+
model.warmup(imgsz=(1 if pt else bs, 3, *imgsz)) # warmup
|
|
106
|
+
seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
|
|
107
|
+
for path, im, im0s, vid_cap, s in dataset:
|
|
108
|
+
with dt[0]:
|
|
109
|
+
im = torch.Tensor(im).to(model.device)
|
|
110
|
+
im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
|
|
111
|
+
if len(im.shape) == 3:
|
|
112
|
+
im = im[None] # expand for batch dim
|
|
113
|
+
|
|
114
|
+
# Inference
|
|
115
|
+
with dt[1]:
|
|
116
|
+
results = model(im)
|
|
117
|
+
|
|
118
|
+
# Post-process
|
|
119
|
+
with dt[2]:
|
|
120
|
+
pred = F.softmax(results, dim=1) # probabilities
|
|
121
|
+
|
|
122
|
+
# Process predictions
|
|
123
|
+
for i, prob in enumerate(pred): # per image
|
|
124
|
+
seen += 1
|
|
125
|
+
if webcam: # batch_size >= 1
|
|
126
|
+
p, im0, frame = path[i], im0s[i].copy(), dataset.count
|
|
127
|
+
s += f'{i}: '
|
|
128
|
+
else:
|
|
129
|
+
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
|
|
130
|
+
|
|
131
|
+
p = Path(p) # to Path
|
|
132
|
+
save_path = str(save_dir / p.name) # im.jpg
|
|
133
|
+
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # im.txt
|
|
134
|
+
|
|
135
|
+
s += '%gx%g ' % im.shape[2:] # print string
|
|
136
|
+
annotator = Annotator(im0, example=str(names), pil=True)
|
|
137
|
+
|
|
138
|
+
# Print results
|
|
139
|
+
top5i = prob.argsort(0, descending=True)[:5].tolist() # top 5 indices
|
|
140
|
+
s += f"{', '.join(f'{names[j]} {prob[j]:.2f}' for j in top5i)}, "
|
|
141
|
+
|
|
142
|
+
# Write results
|
|
143
|
+
text = '\n'.join(f'{prob[j]:.2f} {names[j]}' for j in top5i)
|
|
144
|
+
if save_img or view_img: # Add bbox to image
|
|
145
|
+
annotator.text((32, 32), text, txt_color=(255, 255, 255))
|
|
146
|
+
if save_txt: # Write to file
|
|
147
|
+
with open(f'{txt_path}.txt', 'a') as f:
|
|
148
|
+
f.write(text + '\n')
|
|
149
|
+
|
|
150
|
+
# Stream results
|
|
151
|
+
im0 = annotator.result()
|
|
152
|
+
if view_img:
|
|
153
|
+
if platform.system() == 'Linux' and p not in windows:
|
|
154
|
+
windows.append(p)
|
|
155
|
+
cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO) # allow window resize (Linux)
|
|
156
|
+
cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
|
|
157
|
+
cv2.imshow(str(p), im0)
|
|
158
|
+
cv2.waitKey(1) # 1 millisecond
|
|
159
|
+
|
|
160
|
+
# Save results (image with detections)
|
|
161
|
+
if save_img:
|
|
162
|
+
if dataset.mode == 'image':
|
|
163
|
+
cv2.imwrite(save_path, im0)
|
|
164
|
+
else: # 'video' or 'stream'
|
|
165
|
+
if vid_path[i] != save_path: # new video
|
|
166
|
+
vid_path[i] = save_path
|
|
167
|
+
if isinstance(vid_writer[i], cv2.VideoWriter):
|
|
168
|
+
vid_writer[i].release() # release previous video writer
|
|
169
|
+
if vid_cap: # video
|
|
170
|
+
fps = vid_cap.get(cv2.CAP_PROP_FPS)
|
|
171
|
+
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
172
|
+
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
173
|
+
else: # stream
|
|
174
|
+
fps, w, h = 30, im0.shape[1], im0.shape[0]
|
|
175
|
+
save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos
|
|
176
|
+
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
|
|
177
|
+
vid_writer[i].write(im0)
|
|
178
|
+
|
|
179
|
+
# Print time (inference-only)
|
|
180
|
+
LOGGER.info(f"{s}{dt[1].dt * 1E3:.1f}ms")
|
|
181
|
+
|
|
182
|
+
# Print results
|
|
183
|
+
t = tuple(x.t / seen * 1E3 for x in dt) # speeds per image
|
|
184
|
+
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
|
|
185
|
+
if save_txt or save_img:
|
|
186
|
+
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
|
|
187
|
+
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
|
|
188
|
+
if update:
|
|
189
|
+
strip_optimizer(weights[0]) # update model (to fix SourceChangeWarning)
|
|
190
|
+
|
|
191
|
+
|
|
192
|
+
def parse_opt():
|
|
193
|
+
parser = argparse.ArgumentParser()
|
|
194
|
+
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-cls.pt', help='model path(s)')
|
|
195
|
+
parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob/screen/0(webcam)')
|
|
196
|
+
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
|
|
197
|
+
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[224], help='inference size h,w')
|
|
198
|
+
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
|
199
|
+
parser.add_argument('--view-img', action='store_true', help='show results')
|
|
200
|
+
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
|
|
201
|
+
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
|
|
202
|
+
parser.add_argument('--augment', action='store_true', help='augmented inference')
|
|
203
|
+
parser.add_argument('--visualize', action='store_true', help='visualize features')
|
|
204
|
+
parser.add_argument('--update', action='store_true', help='update all models')
|
|
205
|
+
parser.add_argument('--project', default=ROOT / 'runs/predict-cls', help='save results to project/name')
|
|
206
|
+
parser.add_argument('--name', default='exp', help='save results to project/name')
|
|
207
|
+
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
|
|
208
|
+
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
|
|
209
|
+
parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
|
|
210
|
+
parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
|
|
211
|
+
opt = parser.parse_args()
|
|
212
|
+
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
|
|
213
|
+
print_args(vars(opt))
|
|
214
|
+
return opt
|
|
215
|
+
|
|
216
|
+
|
|
217
|
+
def main(opt):
|
|
218
|
+
check_requirements(exclude=('tensorboard', 'thop'))
|
|
219
|
+
run(**vars(opt))
|
|
220
|
+
|
|
221
|
+
|
|
222
|
+
if __name__ == "__main__":
|
|
223
|
+
opt = parse_opt()
|
|
224
|
+
main(opt)
|