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.
Files changed (119) hide show
  1. yolo9-0.0.1/PKG-INFO +24 -0
  2. yolo9-0.0.1/pyproject.toml +37 -0
  3. yolo9-0.0.1/setup.cfg +4 -0
  4. yolo9-0.0.1/src/yolo9/__init__.py +75 -0
  5. yolo9-0.0.1/src/yolo9/benchmarks.py +142 -0
  6. yolo9-0.0.1/src/yolo9/classify/predict.py +224 -0
  7. yolo9-0.0.1/src/yolo9/classify/train.py +333 -0
  8. yolo9-0.0.1/src/yolo9/classify/val.py +170 -0
  9. yolo9-0.0.1/src/yolo9/data/coco.yaml +125 -0
  10. yolo9-0.0.1/src/yolo9/detect.py +231 -0
  11. yolo9-0.0.1/src/yolo9/detect_dual.py +232 -0
  12. yolo9-0.0.1/src/yolo9/export.py +686 -0
  13. yolo9-0.0.1/src/yolo9/hubconf.py +107 -0
  14. yolo9-0.0.1/src/yolo9/models/__init__.py +1 -0
  15. yolo9-0.0.1/src/yolo9/models/common.py +1233 -0
  16. yolo9-0.0.1/src/yolo9/models/detect/gelan-c.yaml +80 -0
  17. yolo9-0.0.1/src/yolo9/models/detect/gelan-e.yaml +121 -0
  18. yolo9-0.0.1/src/yolo9/models/detect/gelan-m.yaml +80 -0
  19. yolo9-0.0.1/src/yolo9/models/detect/gelan-s.yaml +80 -0
  20. yolo9-0.0.1/src/yolo9/models/detect/gelan-t.yaml +80 -0
  21. yolo9-0.0.1/src/yolo9/models/detect/gelan.yaml +80 -0
  22. yolo9-0.0.1/src/yolo9/models/detect/yolov7-af.yaml +137 -0
  23. yolo9-0.0.1/src/yolo9/models/detect/yolov9-c.yaml +124 -0
  24. yolo9-0.0.1/src/yolo9/models/detect/yolov9-cf.yaml +124 -0
  25. yolo9-0.0.1/src/yolo9/models/detect/yolov9-e.yaml +144 -0
  26. yolo9-0.0.1/src/yolo9/models/detect/yolov9-m.yaml +117 -0
  27. yolo9-0.0.1/src/yolo9/models/detect/yolov9-s.yaml +97 -0
  28. yolo9-0.0.1/src/yolo9/models/detect/yolov9-t.yaml +97 -0
  29. yolo9-0.0.1/src/yolo9/models/detect/yolov9.yaml +117 -0
  30. yolo9-0.0.1/src/yolo9/models/experimental.py +274 -0
  31. yolo9-0.0.1/src/yolo9/models/hub/anchors.yaml +59 -0
  32. yolo9-0.0.1/src/yolo9/models/hub/yolov3-spp.yaml +51 -0
  33. yolo9-0.0.1/src/yolo9/models/hub/yolov3-tiny.yaml +41 -0
  34. yolo9-0.0.1/src/yolo9/models/hub/yolov3.yaml +51 -0
  35. yolo9-0.0.1/src/yolo9/models/panoptic/gelan-c-pan.yaml +80 -0
  36. yolo9-0.0.1/src/yolo9/models/panoptic/yolov7-af-pan.yaml +137 -0
  37. yolo9-0.0.1/src/yolo9/models/segment/gelan-c-dseg.yaml +84 -0
  38. yolo9-0.0.1/src/yolo9/models/segment/gelan-c-seg.yaml +80 -0
  39. yolo9-0.0.1/src/yolo9/models/segment/yolov7-af-seg.yaml +136 -0
  40. yolo9-0.0.1/src/yolo9/models/segment/yolov9-c-dseg.yaml +130 -0
  41. yolo9-0.0.1/src/yolo9/models/tf.py +596 -0
  42. yolo9-0.0.1/src/yolo9/models/yolo.py +818 -0
  43. yolo9-0.0.1/src/yolo9/panoptic/predict.py +246 -0
  44. yolo9-0.0.1/src/yolo9/panoptic/train.py +662 -0
  45. yolo9-0.0.1/src/yolo9/panoptic/val.py +597 -0
  46. yolo9-0.0.1/src/yolo9/segment/predict.py +246 -0
  47. yolo9-0.0.1/src/yolo9/segment/train.py +646 -0
  48. yolo9-0.0.1/src/yolo9/segment/train_dual.py +647 -0
  49. yolo9-0.0.1/src/yolo9/segment/val.py +457 -0
  50. yolo9-0.0.1/src/yolo9/segment/val_dual.py +458 -0
  51. yolo9-0.0.1/src/yolo9/train.py +634 -0
  52. yolo9-0.0.1/src/yolo9/train_dual.py +644 -0
  53. yolo9-0.0.1/src/yolo9/train_triple.py +636 -0
  54. yolo9-0.0.1/src/yolo9/utils/__init__.py +75 -0
  55. yolo9-0.0.1/src/yolo9/utils/activations.py +98 -0
  56. yolo9-0.0.1/src/yolo9/utils/augmentations.py +395 -0
  57. yolo9-0.0.1/src/yolo9/utils/autoanchor.py +164 -0
  58. yolo9-0.0.1/src/yolo9/utils/autobatch.py +67 -0
  59. yolo9-0.0.1/src/yolo9/utils/callbacks.py +71 -0
  60. yolo9-0.0.1/src/yolo9/utils/coco_utils.py +108 -0
  61. yolo9-0.0.1/src/yolo9/utils/dataloaders.py +1217 -0
  62. yolo9-0.0.1/src/yolo9/utils/downloads.py +103 -0
  63. yolo9-0.0.1/src/yolo9/utils/general.py +1135 -0
  64. yolo9-0.0.1/src/yolo9/utils/lion.py +67 -0
  65. yolo9-0.0.1/src/yolo9/utils/loggers/__init__.py +399 -0
  66. yolo9-0.0.1/src/yolo9/utils/loggers/clearml/__init__.py +1 -0
  67. yolo9-0.0.1/src/yolo9/utils/loggers/clearml/clearml_utils.py +157 -0
  68. yolo9-0.0.1/src/yolo9/utils/loggers/clearml/hpo.py +84 -0
  69. yolo9-0.0.1/src/yolo9/utils/loggers/comet/__init__.py +508 -0
  70. yolo9-0.0.1/src/yolo9/utils/loggers/comet/comet_utils.py +150 -0
  71. yolo9-0.0.1/src/yolo9/utils/loggers/comet/hpo.py +118 -0
  72. yolo9-0.0.1/src/yolo9/utils/loggers/comet/optimizer_config.json +209 -0
  73. yolo9-0.0.1/src/yolo9/utils/loggers/wandb/__init__.py +1 -0
  74. yolo9-0.0.1/src/yolo9/utils/loggers/wandb/log_dataset.py +27 -0
  75. yolo9-0.0.1/src/yolo9/utils/loggers/wandb/sweep.py +41 -0
  76. yolo9-0.0.1/src/yolo9/utils/loggers/wandb/sweep.yaml +143 -0
  77. yolo9-0.0.1/src/yolo9/utils/loggers/wandb/wandb_utils.py +589 -0
  78. yolo9-0.0.1/src/yolo9/utils/loss.py +363 -0
  79. yolo9-0.0.1/src/yolo9/utils/loss_tal.py +215 -0
  80. yolo9-0.0.1/src/yolo9/utils/loss_tal_dual.py +385 -0
  81. yolo9-0.0.1/src/yolo9/utils/loss_tal_triple.py +282 -0
  82. yolo9-0.0.1/src/yolo9/utils/metrics.py +397 -0
  83. yolo9-0.0.1/src/yolo9/utils/panoptic/__init__.py +1 -0
  84. yolo9-0.0.1/src/yolo9/utils/panoptic/augmentations.py +183 -0
  85. yolo9-0.0.1/src/yolo9/utils/panoptic/dataloaders.py +478 -0
  86. yolo9-0.0.1/src/yolo9/utils/panoptic/general.py +137 -0
  87. yolo9-0.0.1/src/yolo9/utils/panoptic/loss.py +186 -0
  88. yolo9-0.0.1/src/yolo9/utils/panoptic/loss_tal.py +285 -0
  89. yolo9-0.0.1/src/yolo9/utils/panoptic/metrics.py +272 -0
  90. yolo9-0.0.1/src/yolo9/utils/panoptic/plots.py +164 -0
  91. yolo9-0.0.1/src/yolo9/utils/panoptic/tal/__init__.py +1 -0
  92. yolo9-0.0.1/src/yolo9/utils/panoptic/tal/anchor_generator.py +38 -0
  93. yolo9-0.0.1/src/yolo9/utils/panoptic/tal/assigner.py +181 -0
  94. yolo9-0.0.1/src/yolo9/utils/plots.py +570 -0
  95. yolo9-0.0.1/src/yolo9/utils/segment/__init__.py +1 -0
  96. yolo9-0.0.1/src/yolo9/utils/segment/augmentations.py +99 -0
  97. yolo9-0.0.1/src/yolo9/utils/segment/dataloaders.py +328 -0
  98. yolo9-0.0.1/src/yolo9/utils/segment/general.py +137 -0
  99. yolo9-0.0.1/src/yolo9/utils/segment/loss.py +186 -0
  100. yolo9-0.0.1/src/yolo9/utils/segment/loss_tal.py +261 -0
  101. yolo9-0.0.1/src/yolo9/utils/segment/loss_tal_dual.py +727 -0
  102. yolo9-0.0.1/src/yolo9/utils/segment/metrics.py +205 -0
  103. yolo9-0.0.1/src/yolo9/utils/segment/plots.py +143 -0
  104. yolo9-0.0.1/src/yolo9/utils/segment/tal/__init__.py +1 -0
  105. yolo9-0.0.1/src/yolo9/utils/segment/tal/anchor_generator.py +38 -0
  106. yolo9-0.0.1/src/yolo9/utils/segment/tal/assigner.py +180 -0
  107. yolo9-0.0.1/src/yolo9/utils/tal/__init__.py +1 -0
  108. yolo9-0.0.1/src/yolo9/utils/tal/anchor_generator.py +38 -0
  109. yolo9-0.0.1/src/yolo9/utils/tal/assigner.py +179 -0
  110. yolo9-0.0.1/src/yolo9/utils/torch_utils.py +529 -0
  111. yolo9-0.0.1/src/yolo9/utils/triton.py +81 -0
  112. yolo9-0.0.1/src/yolo9/val.py +389 -0
  113. yolo9-0.0.1/src/yolo9/val_dual.py +393 -0
  114. yolo9-0.0.1/src/yolo9/val_triple.py +391 -0
  115. yolo9-0.0.1/src/yolo9.egg-info/PKG-INFO +24 -0
  116. yolo9-0.0.1/src/yolo9.egg-info/SOURCES.txt +117 -0
  117. yolo9-0.0.1/src/yolo9.egg-info/dependency_links.txt +1 -0
  118. yolo9-0.0.1/src/yolo9.egg-info/requires.txt +18 -0
  119. 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,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -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)