dgenerate-ultralytics-headless 8.3.253__py3-none-any.whl

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 (299) hide show
  1. dgenerate_ultralytics_headless-8.3.253.dist-info/METADATA +405 -0
  2. dgenerate_ultralytics_headless-8.3.253.dist-info/RECORD +299 -0
  3. dgenerate_ultralytics_headless-8.3.253.dist-info/WHEEL +5 -0
  4. dgenerate_ultralytics_headless-8.3.253.dist-info/entry_points.txt +3 -0
  5. dgenerate_ultralytics_headless-8.3.253.dist-info/licenses/LICENSE +661 -0
  6. dgenerate_ultralytics_headless-8.3.253.dist-info/top_level.txt +1 -0
  7. tests/__init__.py +23 -0
  8. tests/conftest.py +59 -0
  9. tests/test_cli.py +131 -0
  10. tests/test_cuda.py +216 -0
  11. tests/test_engine.py +157 -0
  12. tests/test_exports.py +309 -0
  13. tests/test_integrations.py +151 -0
  14. tests/test_python.py +777 -0
  15. tests/test_solutions.py +371 -0
  16. ultralytics/__init__.py +48 -0
  17. ultralytics/assets/bus.jpg +0 -0
  18. ultralytics/assets/zidane.jpg +0 -0
  19. ultralytics/cfg/__init__.py +1028 -0
  20. ultralytics/cfg/datasets/Argoverse.yaml +78 -0
  21. ultralytics/cfg/datasets/DOTAv1.5.yaml +37 -0
  22. ultralytics/cfg/datasets/DOTAv1.yaml +36 -0
  23. ultralytics/cfg/datasets/GlobalWheat2020.yaml +68 -0
  24. ultralytics/cfg/datasets/HomeObjects-3K.yaml +32 -0
  25. ultralytics/cfg/datasets/ImageNet.yaml +2025 -0
  26. ultralytics/cfg/datasets/Objects365.yaml +447 -0
  27. ultralytics/cfg/datasets/SKU-110K.yaml +58 -0
  28. ultralytics/cfg/datasets/TT100K.yaml +346 -0
  29. ultralytics/cfg/datasets/VOC.yaml +102 -0
  30. ultralytics/cfg/datasets/VisDrone.yaml +87 -0
  31. ultralytics/cfg/datasets/african-wildlife.yaml +25 -0
  32. ultralytics/cfg/datasets/brain-tumor.yaml +22 -0
  33. ultralytics/cfg/datasets/carparts-seg.yaml +44 -0
  34. ultralytics/cfg/datasets/coco-pose.yaml +64 -0
  35. ultralytics/cfg/datasets/coco.yaml +118 -0
  36. ultralytics/cfg/datasets/coco128-seg.yaml +101 -0
  37. ultralytics/cfg/datasets/coco128.yaml +101 -0
  38. ultralytics/cfg/datasets/coco8-grayscale.yaml +103 -0
  39. ultralytics/cfg/datasets/coco8-multispectral.yaml +104 -0
  40. ultralytics/cfg/datasets/coco8-pose.yaml +47 -0
  41. ultralytics/cfg/datasets/coco8-seg.yaml +101 -0
  42. ultralytics/cfg/datasets/coco8.yaml +101 -0
  43. ultralytics/cfg/datasets/construction-ppe.yaml +32 -0
  44. ultralytics/cfg/datasets/crack-seg.yaml +22 -0
  45. ultralytics/cfg/datasets/dog-pose.yaml +52 -0
  46. ultralytics/cfg/datasets/dota8-multispectral.yaml +38 -0
  47. ultralytics/cfg/datasets/dota8.yaml +35 -0
  48. ultralytics/cfg/datasets/hand-keypoints.yaml +50 -0
  49. ultralytics/cfg/datasets/kitti.yaml +27 -0
  50. ultralytics/cfg/datasets/lvis.yaml +1240 -0
  51. ultralytics/cfg/datasets/medical-pills.yaml +21 -0
  52. ultralytics/cfg/datasets/open-images-v7.yaml +663 -0
  53. ultralytics/cfg/datasets/package-seg.yaml +22 -0
  54. ultralytics/cfg/datasets/signature.yaml +21 -0
  55. ultralytics/cfg/datasets/tiger-pose.yaml +41 -0
  56. ultralytics/cfg/datasets/xView.yaml +155 -0
  57. ultralytics/cfg/default.yaml +130 -0
  58. ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml +17 -0
  59. ultralytics/cfg/models/11/yolo11-cls.yaml +33 -0
  60. ultralytics/cfg/models/11/yolo11-obb.yaml +50 -0
  61. ultralytics/cfg/models/11/yolo11-pose.yaml +51 -0
  62. ultralytics/cfg/models/11/yolo11-seg.yaml +50 -0
  63. ultralytics/cfg/models/11/yolo11.yaml +50 -0
  64. ultralytics/cfg/models/11/yoloe-11-seg.yaml +48 -0
  65. ultralytics/cfg/models/11/yoloe-11.yaml +48 -0
  66. ultralytics/cfg/models/12/yolo12-cls.yaml +32 -0
  67. ultralytics/cfg/models/12/yolo12-obb.yaml +48 -0
  68. ultralytics/cfg/models/12/yolo12-pose.yaml +49 -0
  69. ultralytics/cfg/models/12/yolo12-seg.yaml +48 -0
  70. ultralytics/cfg/models/12/yolo12.yaml +48 -0
  71. ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +53 -0
  72. ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +45 -0
  73. ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +45 -0
  74. ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +57 -0
  75. ultralytics/cfg/models/v10/yolov10b.yaml +45 -0
  76. ultralytics/cfg/models/v10/yolov10l.yaml +45 -0
  77. ultralytics/cfg/models/v10/yolov10m.yaml +45 -0
  78. ultralytics/cfg/models/v10/yolov10n.yaml +45 -0
  79. ultralytics/cfg/models/v10/yolov10s.yaml +45 -0
  80. ultralytics/cfg/models/v10/yolov10x.yaml +45 -0
  81. ultralytics/cfg/models/v3/yolov3-spp.yaml +49 -0
  82. ultralytics/cfg/models/v3/yolov3-tiny.yaml +40 -0
  83. ultralytics/cfg/models/v3/yolov3.yaml +49 -0
  84. ultralytics/cfg/models/v5/yolov5-p6.yaml +62 -0
  85. ultralytics/cfg/models/v5/yolov5.yaml +51 -0
  86. ultralytics/cfg/models/v6/yolov6.yaml +56 -0
  87. ultralytics/cfg/models/v8/yoloe-v8-seg.yaml +48 -0
  88. ultralytics/cfg/models/v8/yoloe-v8.yaml +48 -0
  89. ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +28 -0
  90. ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +28 -0
  91. ultralytics/cfg/models/v8/yolov8-cls.yaml +32 -0
  92. ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +58 -0
  93. ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +60 -0
  94. ultralytics/cfg/models/v8/yolov8-ghost.yaml +50 -0
  95. ultralytics/cfg/models/v8/yolov8-obb.yaml +49 -0
  96. ultralytics/cfg/models/v8/yolov8-p2.yaml +57 -0
  97. ultralytics/cfg/models/v8/yolov8-p6.yaml +59 -0
  98. ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +60 -0
  99. ultralytics/cfg/models/v8/yolov8-pose.yaml +50 -0
  100. ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +49 -0
  101. ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +59 -0
  102. ultralytics/cfg/models/v8/yolov8-seg.yaml +49 -0
  103. ultralytics/cfg/models/v8/yolov8-world.yaml +51 -0
  104. ultralytics/cfg/models/v8/yolov8-worldv2.yaml +49 -0
  105. ultralytics/cfg/models/v8/yolov8.yaml +49 -0
  106. ultralytics/cfg/models/v9/yolov9c-seg.yaml +41 -0
  107. ultralytics/cfg/models/v9/yolov9c.yaml +41 -0
  108. ultralytics/cfg/models/v9/yolov9e-seg.yaml +64 -0
  109. ultralytics/cfg/models/v9/yolov9e.yaml +64 -0
  110. ultralytics/cfg/models/v9/yolov9m.yaml +41 -0
  111. ultralytics/cfg/models/v9/yolov9s.yaml +41 -0
  112. ultralytics/cfg/models/v9/yolov9t.yaml +41 -0
  113. ultralytics/cfg/trackers/botsort.yaml +21 -0
  114. ultralytics/cfg/trackers/bytetrack.yaml +12 -0
  115. ultralytics/data/__init__.py +26 -0
  116. ultralytics/data/annotator.py +66 -0
  117. ultralytics/data/augment.py +2801 -0
  118. ultralytics/data/base.py +435 -0
  119. ultralytics/data/build.py +437 -0
  120. ultralytics/data/converter.py +855 -0
  121. ultralytics/data/dataset.py +834 -0
  122. ultralytics/data/loaders.py +704 -0
  123. ultralytics/data/scripts/download_weights.sh +18 -0
  124. ultralytics/data/scripts/get_coco.sh +61 -0
  125. ultralytics/data/scripts/get_coco128.sh +18 -0
  126. ultralytics/data/scripts/get_imagenet.sh +52 -0
  127. ultralytics/data/split.py +138 -0
  128. ultralytics/data/split_dota.py +344 -0
  129. ultralytics/data/utils.py +798 -0
  130. ultralytics/engine/__init__.py +1 -0
  131. ultralytics/engine/exporter.py +1580 -0
  132. ultralytics/engine/model.py +1125 -0
  133. ultralytics/engine/predictor.py +508 -0
  134. ultralytics/engine/results.py +1522 -0
  135. ultralytics/engine/trainer.py +977 -0
  136. ultralytics/engine/tuner.py +449 -0
  137. ultralytics/engine/validator.py +387 -0
  138. ultralytics/hub/__init__.py +166 -0
  139. ultralytics/hub/auth.py +151 -0
  140. ultralytics/hub/google/__init__.py +174 -0
  141. ultralytics/hub/session.py +422 -0
  142. ultralytics/hub/utils.py +162 -0
  143. ultralytics/models/__init__.py +9 -0
  144. ultralytics/models/fastsam/__init__.py +7 -0
  145. ultralytics/models/fastsam/model.py +79 -0
  146. ultralytics/models/fastsam/predict.py +169 -0
  147. ultralytics/models/fastsam/utils.py +23 -0
  148. ultralytics/models/fastsam/val.py +38 -0
  149. ultralytics/models/nas/__init__.py +7 -0
  150. ultralytics/models/nas/model.py +98 -0
  151. ultralytics/models/nas/predict.py +56 -0
  152. ultralytics/models/nas/val.py +38 -0
  153. ultralytics/models/rtdetr/__init__.py +7 -0
  154. ultralytics/models/rtdetr/model.py +63 -0
  155. ultralytics/models/rtdetr/predict.py +88 -0
  156. ultralytics/models/rtdetr/train.py +89 -0
  157. ultralytics/models/rtdetr/val.py +216 -0
  158. ultralytics/models/sam/__init__.py +25 -0
  159. ultralytics/models/sam/amg.py +275 -0
  160. ultralytics/models/sam/build.py +365 -0
  161. ultralytics/models/sam/build_sam3.py +377 -0
  162. ultralytics/models/sam/model.py +169 -0
  163. ultralytics/models/sam/modules/__init__.py +1 -0
  164. ultralytics/models/sam/modules/blocks.py +1067 -0
  165. ultralytics/models/sam/modules/decoders.py +495 -0
  166. ultralytics/models/sam/modules/encoders.py +794 -0
  167. ultralytics/models/sam/modules/memory_attention.py +298 -0
  168. ultralytics/models/sam/modules/sam.py +1160 -0
  169. ultralytics/models/sam/modules/tiny_encoder.py +979 -0
  170. ultralytics/models/sam/modules/transformer.py +344 -0
  171. ultralytics/models/sam/modules/utils.py +512 -0
  172. ultralytics/models/sam/predict.py +3940 -0
  173. ultralytics/models/sam/sam3/__init__.py +3 -0
  174. ultralytics/models/sam/sam3/decoder.py +546 -0
  175. ultralytics/models/sam/sam3/encoder.py +529 -0
  176. ultralytics/models/sam/sam3/geometry_encoders.py +415 -0
  177. ultralytics/models/sam/sam3/maskformer_segmentation.py +286 -0
  178. ultralytics/models/sam/sam3/model_misc.py +199 -0
  179. ultralytics/models/sam/sam3/necks.py +129 -0
  180. ultralytics/models/sam/sam3/sam3_image.py +339 -0
  181. ultralytics/models/sam/sam3/text_encoder_ve.py +307 -0
  182. ultralytics/models/sam/sam3/vitdet.py +547 -0
  183. ultralytics/models/sam/sam3/vl_combiner.py +160 -0
  184. ultralytics/models/utils/__init__.py +1 -0
  185. ultralytics/models/utils/loss.py +466 -0
  186. ultralytics/models/utils/ops.py +315 -0
  187. ultralytics/models/yolo/__init__.py +7 -0
  188. ultralytics/models/yolo/classify/__init__.py +7 -0
  189. ultralytics/models/yolo/classify/predict.py +90 -0
  190. ultralytics/models/yolo/classify/train.py +202 -0
  191. ultralytics/models/yolo/classify/val.py +216 -0
  192. ultralytics/models/yolo/detect/__init__.py +7 -0
  193. ultralytics/models/yolo/detect/predict.py +122 -0
  194. ultralytics/models/yolo/detect/train.py +227 -0
  195. ultralytics/models/yolo/detect/val.py +507 -0
  196. ultralytics/models/yolo/model.py +430 -0
  197. ultralytics/models/yolo/obb/__init__.py +7 -0
  198. ultralytics/models/yolo/obb/predict.py +56 -0
  199. ultralytics/models/yolo/obb/train.py +79 -0
  200. ultralytics/models/yolo/obb/val.py +302 -0
  201. ultralytics/models/yolo/pose/__init__.py +7 -0
  202. ultralytics/models/yolo/pose/predict.py +65 -0
  203. ultralytics/models/yolo/pose/train.py +110 -0
  204. ultralytics/models/yolo/pose/val.py +248 -0
  205. ultralytics/models/yolo/segment/__init__.py +7 -0
  206. ultralytics/models/yolo/segment/predict.py +109 -0
  207. ultralytics/models/yolo/segment/train.py +69 -0
  208. ultralytics/models/yolo/segment/val.py +307 -0
  209. ultralytics/models/yolo/world/__init__.py +5 -0
  210. ultralytics/models/yolo/world/train.py +173 -0
  211. ultralytics/models/yolo/world/train_world.py +178 -0
  212. ultralytics/models/yolo/yoloe/__init__.py +22 -0
  213. ultralytics/models/yolo/yoloe/predict.py +162 -0
  214. ultralytics/models/yolo/yoloe/train.py +287 -0
  215. ultralytics/models/yolo/yoloe/train_seg.py +122 -0
  216. ultralytics/models/yolo/yoloe/val.py +206 -0
  217. ultralytics/nn/__init__.py +27 -0
  218. ultralytics/nn/autobackend.py +964 -0
  219. ultralytics/nn/modules/__init__.py +182 -0
  220. ultralytics/nn/modules/activation.py +54 -0
  221. ultralytics/nn/modules/block.py +1947 -0
  222. ultralytics/nn/modules/conv.py +669 -0
  223. ultralytics/nn/modules/head.py +1183 -0
  224. ultralytics/nn/modules/transformer.py +793 -0
  225. ultralytics/nn/modules/utils.py +159 -0
  226. ultralytics/nn/tasks.py +1768 -0
  227. ultralytics/nn/text_model.py +356 -0
  228. ultralytics/py.typed +1 -0
  229. ultralytics/solutions/__init__.py +41 -0
  230. ultralytics/solutions/ai_gym.py +108 -0
  231. ultralytics/solutions/analytics.py +264 -0
  232. ultralytics/solutions/config.py +107 -0
  233. ultralytics/solutions/distance_calculation.py +123 -0
  234. ultralytics/solutions/heatmap.py +125 -0
  235. ultralytics/solutions/instance_segmentation.py +86 -0
  236. ultralytics/solutions/object_blurrer.py +89 -0
  237. ultralytics/solutions/object_counter.py +190 -0
  238. ultralytics/solutions/object_cropper.py +87 -0
  239. ultralytics/solutions/parking_management.py +280 -0
  240. ultralytics/solutions/queue_management.py +93 -0
  241. ultralytics/solutions/region_counter.py +133 -0
  242. ultralytics/solutions/security_alarm.py +151 -0
  243. ultralytics/solutions/similarity_search.py +219 -0
  244. ultralytics/solutions/solutions.py +828 -0
  245. ultralytics/solutions/speed_estimation.py +114 -0
  246. ultralytics/solutions/streamlit_inference.py +260 -0
  247. ultralytics/solutions/templates/similarity-search.html +156 -0
  248. ultralytics/solutions/trackzone.py +88 -0
  249. ultralytics/solutions/vision_eye.py +67 -0
  250. ultralytics/trackers/__init__.py +7 -0
  251. ultralytics/trackers/basetrack.py +115 -0
  252. ultralytics/trackers/bot_sort.py +257 -0
  253. ultralytics/trackers/byte_tracker.py +469 -0
  254. ultralytics/trackers/track.py +116 -0
  255. ultralytics/trackers/utils/__init__.py +1 -0
  256. ultralytics/trackers/utils/gmc.py +339 -0
  257. ultralytics/trackers/utils/kalman_filter.py +482 -0
  258. ultralytics/trackers/utils/matching.py +154 -0
  259. ultralytics/utils/__init__.py +1450 -0
  260. ultralytics/utils/autobatch.py +118 -0
  261. ultralytics/utils/autodevice.py +205 -0
  262. ultralytics/utils/benchmarks.py +728 -0
  263. ultralytics/utils/callbacks/__init__.py +5 -0
  264. ultralytics/utils/callbacks/base.py +233 -0
  265. ultralytics/utils/callbacks/clearml.py +146 -0
  266. ultralytics/utils/callbacks/comet.py +625 -0
  267. ultralytics/utils/callbacks/dvc.py +197 -0
  268. ultralytics/utils/callbacks/hub.py +110 -0
  269. ultralytics/utils/callbacks/mlflow.py +134 -0
  270. ultralytics/utils/callbacks/neptune.py +126 -0
  271. ultralytics/utils/callbacks/platform.py +453 -0
  272. ultralytics/utils/callbacks/raytune.py +42 -0
  273. ultralytics/utils/callbacks/tensorboard.py +123 -0
  274. ultralytics/utils/callbacks/wb.py +188 -0
  275. ultralytics/utils/checks.py +1020 -0
  276. ultralytics/utils/cpu.py +85 -0
  277. ultralytics/utils/dist.py +123 -0
  278. ultralytics/utils/downloads.py +529 -0
  279. ultralytics/utils/errors.py +35 -0
  280. ultralytics/utils/events.py +113 -0
  281. ultralytics/utils/export/__init__.py +7 -0
  282. ultralytics/utils/export/engine.py +237 -0
  283. ultralytics/utils/export/imx.py +325 -0
  284. ultralytics/utils/export/tensorflow.py +231 -0
  285. ultralytics/utils/files.py +219 -0
  286. ultralytics/utils/git.py +137 -0
  287. ultralytics/utils/instance.py +484 -0
  288. ultralytics/utils/logger.py +506 -0
  289. ultralytics/utils/loss.py +849 -0
  290. ultralytics/utils/metrics.py +1563 -0
  291. ultralytics/utils/nms.py +337 -0
  292. ultralytics/utils/ops.py +664 -0
  293. ultralytics/utils/patches.py +201 -0
  294. ultralytics/utils/plotting.py +1047 -0
  295. ultralytics/utils/tal.py +404 -0
  296. ultralytics/utils/torch_utils.py +984 -0
  297. ultralytics/utils/tqdm.py +443 -0
  298. ultralytics/utils/triton.py +112 -0
  299. ultralytics/utils/tuner.py +168 -0
@@ -0,0 +1,704 @@
1
+ # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
2
+
3
+ from __future__ import annotations
4
+
5
+ import glob
6
+ import math
7
+ import os
8
+ import time
9
+ import urllib
10
+ from dataclasses import dataclass
11
+ from pathlib import Path
12
+ from threading import Thread
13
+ from typing import Any
14
+
15
+ import cv2
16
+ import numpy as np
17
+ import torch
18
+ from PIL import Image
19
+
20
+ from ultralytics.data.utils import FORMATS_HELP_MSG, IMG_FORMATS, VID_FORMATS
21
+ from ultralytics.utils import IS_COLAB, IS_KAGGLE, LOGGER, ops
22
+ from ultralytics.utils.checks import check_requirements
23
+ from ultralytics.utils.patches import imread
24
+
25
+
26
+ @dataclass
27
+ class SourceTypes:
28
+ """Class to represent various types of input sources for predictions.
29
+
30
+ This class uses dataclass to define boolean flags for different types of input sources that can be used for making
31
+ predictions with YOLO models.
32
+
33
+ Attributes:
34
+ stream (bool): Flag indicating if the input source is a video stream.
35
+ screenshot (bool): Flag indicating if the input source is a screenshot.
36
+ from_img (bool): Flag indicating if the input source is an image file.
37
+ tensor (bool): Flag indicating if the input source is a tensor.
38
+
39
+ Examples:
40
+ >>> source_types = SourceTypes(stream=True, screenshot=False, from_img=False)
41
+ >>> print(source_types.stream)
42
+ True
43
+ >>> print(source_types.from_img)
44
+ False
45
+ """
46
+
47
+ stream: bool = False
48
+ screenshot: bool = False
49
+ from_img: bool = False
50
+ tensor: bool = False
51
+
52
+
53
+ class LoadStreams:
54
+ """Stream Loader for various types of video streams.
55
+
56
+ Supports RTSP, RTMP, HTTP, and TCP streams. This class handles the loading and processing of multiple video streams
57
+ simultaneously, making it suitable for real-time video analysis tasks.
58
+
59
+ Attributes:
60
+ sources (list[str]): The source input paths or URLs for the video streams.
61
+ vid_stride (int): Video frame-rate stride.
62
+ buffer (bool): Whether to buffer input streams.
63
+ running (bool): Flag to indicate if the streaming thread is running.
64
+ mode (str): Set to 'stream' indicating real-time capture.
65
+ imgs (list[list[np.ndarray]]): List of image frames for each stream.
66
+ fps (list[float]): List of FPS for each stream.
67
+ frames (list[int]): List of total frames for each stream.
68
+ threads (list[Thread]): List of threads for each stream.
69
+ shape (list[tuple[int, int, int]]): List of shapes for each stream.
70
+ caps (list[cv2.VideoCapture]): List of cv2.VideoCapture objects for each stream.
71
+ bs (int): Batch size for processing.
72
+ cv2_flag (int): OpenCV flag for image reading (grayscale or color/BGR).
73
+
74
+ Methods:
75
+ update: Read stream frames in daemon thread.
76
+ close: Close stream loader and release resources.
77
+ __iter__: Returns an iterator object for the class.
78
+ __next__: Returns source paths, transformed, and original images for processing.
79
+ __len__: Return the length of the sources object.
80
+
81
+ Examples:
82
+ >>> stream_loader = LoadStreams("rtsp://example.com/stream1.mp4")
83
+ >>> for sources, imgs, _ in stream_loader:
84
+ ... # Process the images
85
+ ... pass
86
+ >>> stream_loader.close()
87
+
88
+ Notes:
89
+ - The class uses threading to efficiently load frames from multiple streams simultaneously.
90
+ - It automatically handles YouTube links, converting them to the best available stream URL.
91
+ - The class implements a buffer system to manage frame storage and retrieval.
92
+ """
93
+
94
+ def __init__(self, sources: str = "file.streams", vid_stride: int = 1, buffer: bool = False, channels: int = 3):
95
+ """Initialize stream loader for multiple video sources, supporting various stream types.
96
+
97
+ Args:
98
+ sources (str): Path to streams file or single stream URL.
99
+ vid_stride (int): Video frame-rate stride.
100
+ buffer (bool): Whether to buffer input streams.
101
+ channels (int): Number of image channels (1 for grayscale, 3 for color).
102
+ """
103
+ torch.backends.cudnn.benchmark = True # faster for fixed-size inference
104
+ self.buffer = buffer # buffer input streams
105
+ self.running = True # running flag for Thread
106
+ self.mode = "stream"
107
+ self.vid_stride = vid_stride # video frame-rate stride
108
+ self.cv2_flag = cv2.IMREAD_GRAYSCALE if channels == 1 else cv2.IMREAD_COLOR # grayscale or color (BGR)
109
+
110
+ sources = Path(sources).read_text().rsplit() if os.path.isfile(sources) else [sources]
111
+ n = len(sources)
112
+ self.bs = n
113
+ self.fps = [0] * n # frames per second
114
+ self.frames = [0] * n
115
+ self.threads = [None] * n
116
+ self.caps = [None] * n # video capture objects
117
+ self.imgs = [[] for _ in range(n)] # images
118
+ self.shape = [[] for _ in range(n)] # image shapes
119
+ self.sources = [ops.clean_str(x).replace(os.sep, "_") for x in sources] # clean source names for later
120
+ for i, s in enumerate(sources): # index, source
121
+ # Start thread to read frames from video stream
122
+ st = f"{i + 1}/{n}: {s}... "
123
+ if urllib.parse.urlparse(s).hostname in {"www.youtube.com", "youtube.com", "youtu.be"}: # YouTube video
124
+ # YouTube format i.e. 'https://www.youtube.com/watch?v=Jsn8D3aC840' or 'https://youtu.be/Jsn8D3aC840'
125
+ s = get_best_youtube_url(s)
126
+ s = int(s) if s.isnumeric() else s # i.e. s = '0' local webcam
127
+ if s == 0 and (IS_COLAB or IS_KAGGLE):
128
+ raise NotImplementedError(
129
+ "'source=0' webcam not supported in Colab and Kaggle notebooks. "
130
+ "Try running 'source=0' in a local environment."
131
+ )
132
+ self.caps[i] = cv2.VideoCapture(s) # store video capture object
133
+ if not self.caps[i].isOpened():
134
+ raise ConnectionError(f"{st}Failed to open {s}")
135
+ w = int(self.caps[i].get(cv2.CAP_PROP_FRAME_WIDTH))
136
+ h = int(self.caps[i].get(cv2.CAP_PROP_FRAME_HEIGHT))
137
+ fps = self.caps[i].get(cv2.CAP_PROP_FPS) # warning: may return 0 or nan
138
+ self.frames[i] = max(int(self.caps[i].get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float(
139
+ "inf"
140
+ ) # infinite stream fallback
141
+ self.fps[i] = max((fps if math.isfinite(fps) else 0) % 100, 0) or 30 # 30 FPS fallback
142
+
143
+ success, im = self.caps[i].read() # guarantee first frame
144
+ im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)[..., None] if self.cv2_flag == cv2.IMREAD_GRAYSCALE else im
145
+ if not success or im is None:
146
+ raise ConnectionError(f"{st}Failed to read images from {s}")
147
+ self.imgs[i].append(im)
148
+ self.shape[i] = im.shape
149
+ self.threads[i] = Thread(target=self.update, args=([i, self.caps[i], s]), daemon=True)
150
+ LOGGER.info(f"{st}Success ✅ ({self.frames[i]} frames of shape {w}x{h} at {self.fps[i]:.2f} FPS)")
151
+ self.threads[i].start()
152
+ LOGGER.info("") # newline
153
+
154
+ def update(self, i: int, cap: cv2.VideoCapture, stream: str):
155
+ """Read stream frames in daemon thread and update image buffer."""
156
+ n, f = 0, self.frames[i] # frame number, frame array
157
+ while self.running and cap.isOpened() and n < (f - 1):
158
+ if len(self.imgs[i]) < 30: # keep a <=30-image buffer
159
+ n += 1
160
+ cap.grab() # .read() = .grab() followed by .retrieve()
161
+ if n % self.vid_stride == 0:
162
+ success, im = cap.retrieve()
163
+ im = (
164
+ cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)[..., None] if self.cv2_flag == cv2.IMREAD_GRAYSCALE else im
165
+ )
166
+ if not success:
167
+ im = np.zeros(self.shape[i], dtype=np.uint8)
168
+ LOGGER.warning("Video stream unresponsive, please check your IP camera connection.")
169
+ cap.open(stream) # re-open stream if signal was lost
170
+ if self.buffer:
171
+ self.imgs[i].append(im)
172
+ else:
173
+ self.imgs[i] = [im]
174
+ else:
175
+ time.sleep(0.01) # wait until the buffer is empty
176
+
177
+ def close(self):
178
+ """Terminate stream loader, stop threads, and release video capture resources."""
179
+ self.running = False # stop flag for Thread
180
+ for thread in self.threads:
181
+ if thread.is_alive():
182
+ thread.join(timeout=5) # Add timeout
183
+ for cap in self.caps: # Iterate through the stored VideoCapture objects
184
+ try:
185
+ cap.release() # release video capture
186
+ except Exception as e:
187
+ LOGGER.warning(f"Could not release VideoCapture object: {e}")
188
+
189
+ def __iter__(self):
190
+ """Iterate through YOLO image feed and re-open unresponsive streams."""
191
+ self.count = -1
192
+ return self
193
+
194
+ def __next__(self) -> tuple[list[str], list[np.ndarray], list[str]]:
195
+ """Return the next batch of frames from multiple video streams for processing."""
196
+ self.count += 1
197
+
198
+ images = []
199
+ for i, x in enumerate(self.imgs):
200
+ # Wait until a frame is available in each buffer
201
+ while not x:
202
+ if not self.threads[i].is_alive():
203
+ self.close()
204
+ raise StopIteration
205
+ time.sleep(1 / min(self.fps))
206
+ x = self.imgs[i]
207
+ if not x:
208
+ LOGGER.warning(f"Waiting for stream {i}")
209
+
210
+ # Get and remove the first frame from imgs buffer
211
+ if self.buffer:
212
+ images.append(x.pop(0))
213
+
214
+ # Get the last frame, and clear the rest from the imgs buffer
215
+ else:
216
+ images.append(x.pop(-1) if x else np.zeros(self.shape[i], dtype=np.uint8))
217
+ x.clear()
218
+
219
+ return self.sources, images, [""] * self.bs
220
+
221
+ def __len__(self) -> int:
222
+ """Return the number of video streams in the LoadStreams object."""
223
+ return self.bs # 1E12 frames = 32 streams at 30 FPS for 30 years
224
+
225
+
226
+ class LoadScreenshots:
227
+ """Ultralytics screenshot dataloader for capturing and processing screen images.
228
+
229
+ This class manages the loading of screenshot images for processing with YOLO. It is suitable for use with `yolo
230
+ predict source=screen`.
231
+
232
+ Attributes:
233
+ source (str): The source input indicating which screen to capture.
234
+ screen (int): The screen number to capture.
235
+ left (int): The left coordinate for screen capture area.
236
+ top (int): The top coordinate for screen capture area.
237
+ width (int): The width of the screen capture area.
238
+ height (int): The height of the screen capture area.
239
+ mode (str): Set to 'stream' indicating real-time capture.
240
+ frame (int): Counter for captured frames.
241
+ sct (mss.mss): Screen capture object from `mss` library.
242
+ bs (int): Batch size, set to 1.
243
+ fps (int): Frames per second, set to 30.
244
+ monitor (dict[str, int]): Monitor configuration details.
245
+ cv2_flag (int): OpenCV flag for image reading (grayscale or color/BGR).
246
+
247
+ Methods:
248
+ __iter__: Returns an iterator object.
249
+ __next__: Captures the next screenshot and returns it.
250
+
251
+ Examples:
252
+ >>> loader = LoadScreenshots("0 100 100 640 480") # screen 0, top-left (100,100), 640x480
253
+ >>> for source, im, im0s, vid_cap, s in loader:
254
+ ... print(f"Captured frame: {im.shape}")
255
+ """
256
+
257
+ def __init__(self, source: str, channels: int = 3):
258
+ """Initialize screenshot capture with specified screen and region parameters.
259
+
260
+ Args:
261
+ source (str): Screen capture source string in format "screen_num left top width height".
262
+ channels (int): Number of image channels (1 for grayscale, 3 for color).
263
+ """
264
+ check_requirements("mss")
265
+ import mss
266
+
267
+ source, *params = source.split()
268
+ self.screen, left, top, width, height = 0, None, None, None, None # default to full screen 0
269
+ if len(params) == 1:
270
+ self.screen = int(params[0])
271
+ elif len(params) == 4:
272
+ left, top, width, height = (int(x) for x in params)
273
+ elif len(params) == 5:
274
+ self.screen, left, top, width, height = (int(x) for x in params)
275
+ self.mode = "stream"
276
+ self.frame = 0
277
+ self.sct = mss.mss()
278
+ self.bs = 1
279
+ self.fps = 30
280
+ self.cv2_flag = cv2.IMREAD_GRAYSCALE if channels == 1 else cv2.IMREAD_COLOR # grayscale or color (BGR)
281
+
282
+ # Parse monitor shape
283
+ monitor = self.sct.monitors[self.screen]
284
+ self.top = monitor["top"] if top is None else (monitor["top"] + top)
285
+ self.left = monitor["left"] if left is None else (monitor["left"] + left)
286
+ self.width = width or monitor["width"]
287
+ self.height = height or monitor["height"]
288
+ self.monitor = {"left": self.left, "top": self.top, "width": self.width, "height": self.height}
289
+
290
+ def __iter__(self):
291
+ """Yield the next screenshot image from the specified screen or region for processing."""
292
+ return self
293
+
294
+ def __next__(self) -> tuple[list[str], list[np.ndarray], list[str]]:
295
+ """Capture and return the next screenshot as a numpy array using the mss library."""
296
+ im0 = np.asarray(self.sct.grab(self.monitor))[:, :, :3] # BGRA to BGR
297
+ im0 = cv2.cvtColor(im0, cv2.COLOR_BGR2GRAY)[..., None] if self.cv2_flag == cv2.IMREAD_GRAYSCALE else im0
298
+ s = f"screen {self.screen} (LTWH): {self.left},{self.top},{self.width},{self.height}: "
299
+
300
+ self.frame += 1
301
+ return [str(self.screen)], [im0], [s] # screen, img, string
302
+
303
+
304
+ class LoadImagesAndVideos:
305
+ """A class for loading and processing images and videos for YOLO object detection.
306
+
307
+ This class manages the loading and pre-processing of image and video data from various sources, including single
308
+ image files, video files, and lists of image and video paths.
309
+
310
+ Attributes:
311
+ files (list[str]): List of image and video file paths.
312
+ nf (int): Total number of files (images and videos).
313
+ video_flag (list[bool]): Flags indicating whether a file is a video (True) or an image (False).
314
+ mode (str): Current mode, 'image' or 'video'.
315
+ vid_stride (int): Stride for video frame-rate.
316
+ bs (int): Batch size.
317
+ cap (cv2.VideoCapture): Video capture object for OpenCV.
318
+ frame (int): Frame counter for video.
319
+ frames (int): Total number of frames in the video.
320
+ count (int): Counter for iteration, initialized at 0 during __iter__().
321
+ ni (int): Number of images.
322
+ cv2_flag (int): OpenCV flag for image reading (grayscale or color/BGR).
323
+
324
+ Methods:
325
+ __init__: Initialize the LoadImagesAndVideos object.
326
+ __iter__: Returns an iterator object for VideoStream or ImageFolder.
327
+ __next__: Returns the next batch of images or video frames along with their paths and metadata.
328
+ _new_video: Creates a new video capture object for the given path.
329
+ __len__: Returns the number of batches in the object.
330
+
331
+ Examples:
332
+ >>> loader = LoadImagesAndVideos("path/to/data", batch=32, vid_stride=1)
333
+ >>> for paths, imgs, info in loader:
334
+ ... # Process batch of images or video frames
335
+ ... pass
336
+
337
+ Notes:
338
+ - Supports various image formats including HEIC.
339
+ - Handles both local files and directories.
340
+ - Can read from a text file containing paths to images and videos.
341
+ """
342
+
343
+ def __init__(self, path: str | Path | list, batch: int = 1, vid_stride: int = 1, channels: int = 3):
344
+ """Initialize dataloader for images and videos, supporting various input formats.
345
+
346
+ Args:
347
+ path (str | Path | list): Path to images/videos, directory, or list of paths.
348
+ batch (int): Batch size for processing.
349
+ vid_stride (int): Video frame-rate stride.
350
+ channels (int): Number of image channels (1 for grayscale, 3 for color).
351
+ """
352
+ parent = None
353
+ if isinstance(path, str) and Path(path).suffix in {".txt", ".csv"}: # txt/csv file with source paths
354
+ parent, content = Path(path).parent, Path(path).read_text()
355
+ path = content.splitlines() if Path(path).suffix == ".txt" else content.split(",") # list of sources
356
+ path = [p.strip() for p in path]
357
+ files = []
358
+ for p in sorted(path) if isinstance(path, (list, tuple)) else [path]:
359
+ a = str(Path(p).absolute()) # do not use .resolve() https://github.com/ultralytics/ultralytics/issues/2912
360
+ if "*" in a:
361
+ files.extend(sorted(glob.glob(a, recursive=True))) # glob
362
+ elif os.path.isdir(a):
363
+ files.extend(sorted(glob.glob(os.path.join(a, "*.*")))) # dir
364
+ elif os.path.isfile(a):
365
+ files.append(a) # files (absolute or relative to CWD)
366
+ elif parent and (parent / p).is_file():
367
+ files.append(str((parent / p).absolute())) # files (relative to *.txt file parent)
368
+ else:
369
+ raise FileNotFoundError(f"{p} does not exist")
370
+
371
+ # Define files as images or videos
372
+ images, videos = [], []
373
+ for f in files:
374
+ suffix = f.rpartition(".")[-1].lower() # Get file extension without the dot and lowercase
375
+ if suffix in IMG_FORMATS:
376
+ images.append(f)
377
+ elif suffix in VID_FORMATS:
378
+ videos.append(f)
379
+ ni, nv = len(images), len(videos)
380
+
381
+ self.files = images + videos
382
+ self.nf = ni + nv # number of files
383
+ self.ni = ni # number of images
384
+ self.video_flag = [False] * ni + [True] * nv
385
+ self.mode = "video" if ni == 0 else "image" # default to video if no images
386
+ self.vid_stride = vid_stride # video frame-rate stride
387
+ self.bs = batch
388
+ self.cv2_flag = cv2.IMREAD_GRAYSCALE if channels == 1 else cv2.IMREAD_COLOR # grayscale or color (BGR)
389
+ if any(videos):
390
+ self._new_video(videos[0]) # new video
391
+ else:
392
+ self.cap = None
393
+ if self.nf == 0:
394
+ raise FileNotFoundError(f"No images or videos found in {p}. {FORMATS_HELP_MSG}")
395
+
396
+ def __iter__(self):
397
+ """Iterate through image/video files, yielding source paths, images, and metadata."""
398
+ self.count = 0
399
+ return self
400
+
401
+ def __next__(self) -> tuple[list[str], list[np.ndarray], list[str]]:
402
+ """Return the next batch of images or video frames with their paths and metadata."""
403
+ paths, imgs, info = [], [], []
404
+ while len(imgs) < self.bs:
405
+ if self.count >= self.nf: # end of file list
406
+ if imgs:
407
+ return paths, imgs, info # return last partial batch
408
+ else:
409
+ raise StopIteration
410
+
411
+ path = self.files[self.count]
412
+ if self.video_flag[self.count]:
413
+ self.mode = "video"
414
+ if not self.cap or not self.cap.isOpened():
415
+ self._new_video(path)
416
+
417
+ success = False
418
+ for _ in range(self.vid_stride):
419
+ success = self.cap.grab()
420
+ if not success:
421
+ break # end of video or failure
422
+
423
+ if success:
424
+ success, im0 = self.cap.retrieve()
425
+ im0 = (
426
+ cv2.cvtColor(im0, cv2.COLOR_BGR2GRAY)[..., None]
427
+ if self.cv2_flag == cv2.IMREAD_GRAYSCALE
428
+ else im0
429
+ )
430
+ if success:
431
+ self.frame += 1
432
+ paths.append(path)
433
+ imgs.append(im0)
434
+ info.append(f"video {self.count + 1}/{self.nf} (frame {self.frame}/{self.frames}) {path}: ")
435
+ if self.frame == self.frames: # end of video
436
+ self.count += 1
437
+ self.cap.release()
438
+ else:
439
+ # Move to the next file if the current video ended or failed to open
440
+ self.count += 1
441
+ if self.cap:
442
+ self.cap.release()
443
+ if self.count < self.nf:
444
+ self._new_video(self.files[self.count])
445
+ else:
446
+ # Handle image files (including HEIC)
447
+ self.mode = "image"
448
+ if path.rpartition(".")[-1].lower() == "heic":
449
+ # Load HEIC image using Pillow with pillow-heif
450
+ check_requirements("pi-heif")
451
+
452
+ from pi_heif import register_heif_opener
453
+
454
+ register_heif_opener() # Register HEIF opener with Pillow
455
+ with Image.open(path) as img:
456
+ im0 = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR) # convert image to BGR nparray
457
+ else:
458
+ im0 = imread(path, flags=self.cv2_flag) # BGR
459
+ if im0 is None:
460
+ LOGGER.warning(f"Image Read Error {path}")
461
+ else:
462
+ paths.append(path)
463
+ imgs.append(im0)
464
+ info.append(f"image {self.count + 1}/{self.nf} {path}: ")
465
+ self.count += 1 # move to the next file
466
+ if self.count >= self.ni: # end of image list
467
+ break
468
+
469
+ return paths, imgs, info
470
+
471
+ def _new_video(self, path: str):
472
+ """Create a new video capture object for the given path and initialize video-related attributes."""
473
+ self.frame = 0
474
+ self.cap = cv2.VideoCapture(path)
475
+ self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))
476
+ if not self.cap.isOpened():
477
+ raise FileNotFoundError(f"Failed to open video {path}")
478
+ self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT) / self.vid_stride)
479
+
480
+ def __len__(self) -> int:
481
+ """Return the number of files (images and videos) in the dataset."""
482
+ return math.ceil(self.nf / self.bs) # number of batches
483
+
484
+
485
+ class LoadPilAndNumpy:
486
+ """Load images from PIL and Numpy arrays for batch processing.
487
+
488
+ This class manages loading and pre-processing of image data from both PIL and Numpy formats. It performs basic
489
+ validation and format conversion to ensure that the images are in the required format for downstream processing.
490
+
491
+ Attributes:
492
+ paths (list[str]): List of image paths or autogenerated filenames.
493
+ im0 (list[np.ndarray]): List of images stored as Numpy arrays.
494
+ mode (str): Type of data being processed, set to 'image'.
495
+ bs (int): Batch size, equivalent to the length of `im0`.
496
+
497
+ Methods:
498
+ _single_check: Validate and format a single image to a Numpy array.
499
+
500
+ Examples:
501
+ >>> from PIL import Image
502
+ >>> import numpy as np
503
+ >>> pil_img = Image.new("RGB", (100, 100))
504
+ >>> np_img = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
505
+ >>> loader = LoadPilAndNumpy([pil_img, np_img])
506
+ >>> paths, images, _ = next(iter(loader))
507
+ >>> print(f"Loaded {len(images)} images")
508
+ Loaded 2 images
509
+ """
510
+
511
+ def __init__(self, im0: Image.Image | np.ndarray | list, channels: int = 3):
512
+ """Initialize a loader for PIL and Numpy images, converting inputs to a standardized format.
513
+
514
+ Args:
515
+ im0 (PIL.Image.Image | np.ndarray | list): Single image or list of images in PIL or numpy format.
516
+ channels (int): Number of image channels (1 for grayscale, 3 for color).
517
+ """
518
+ if not isinstance(im0, list):
519
+ im0 = [im0]
520
+ # use `image{i}.jpg` when Image.filename returns an empty path.
521
+ self.paths = [getattr(im, "filename", "") or f"image{i}.jpg" for i, im in enumerate(im0)]
522
+ pil_flag = "L" if channels == 1 else "RGB" # grayscale or RGB
523
+ self.im0 = [self._single_check(im, pil_flag) for im in im0]
524
+ self.mode = "image"
525
+ self.bs = len(self.im0)
526
+
527
+ @staticmethod
528
+ def _single_check(im: Image.Image | np.ndarray, flag: str = "RGB") -> np.ndarray:
529
+ """Validate and format an image to a NumPy array.
530
+
531
+ Notes:
532
+ - PIL inputs are converted to NumPy and returned in OpenCV-compatible BGR order for color images.
533
+ - NumPy inputs are returned as-is (no channel-order conversion is applied).
534
+ """
535
+ assert isinstance(im, (Image.Image, np.ndarray)), f"Expected PIL/np.ndarray image type, but got {type(im)}"
536
+ if isinstance(im, Image.Image):
537
+ im = np.asarray(im.convert(flag))
538
+ # Add a new axis if grayscale; convert RGB -> BGR for OpenCV compatibility.
539
+ im = im[..., None] if flag == "L" else im[..., ::-1]
540
+ im = np.ascontiguousarray(im) # contiguous
541
+ elif im.ndim == 2: # grayscale in numpy form
542
+ im = im[..., None]
543
+ return im
544
+
545
+ def __len__(self) -> int:
546
+ """Return the length of the 'im0' attribute, representing the number of loaded images."""
547
+ return len(self.im0)
548
+
549
+ def __next__(self) -> tuple[list[str], list[np.ndarray], list[str]]:
550
+ """Return the next batch of images, paths, and metadata for processing."""
551
+ if self.count == 1: # loop only once as it's batch inference
552
+ raise StopIteration
553
+ self.count += 1
554
+ return self.paths, self.im0, [""] * self.bs
555
+
556
+ def __iter__(self):
557
+ """Iterate through PIL/numpy images, yielding paths, raw images, and metadata for processing."""
558
+ self.count = 0
559
+ return self
560
+
561
+
562
+ class LoadTensor:
563
+ """A class for loading and processing tensor data for object detection tasks.
564
+
565
+ This class handles the loading and pre-processing of image data from PyTorch tensors, preparing them for further
566
+ processing in object detection pipelines.
567
+
568
+ Attributes:
569
+ im0 (torch.Tensor): The input tensor containing the image(s) with shape (B, C, H, W).
570
+ bs (int): Batch size, inferred from the shape of `im0`.
571
+ mode (str): Current processing mode, set to 'image'.
572
+ paths (list[str]): List of image paths or auto-generated filenames.
573
+
574
+ Methods:
575
+ _single_check: Validates and formats an input tensor.
576
+
577
+ Examples:
578
+ >>> import torch
579
+ >>> tensor = torch.rand(1, 3, 640, 640)
580
+ >>> loader = LoadTensor(tensor)
581
+ >>> paths, images, info = next(iter(loader))
582
+ >>> print(f"Processed {len(images)} images")
583
+ """
584
+
585
+ def __init__(self, im0: torch.Tensor) -> None:
586
+ """Initialize LoadTensor object for processing torch.Tensor image data.
587
+
588
+ Args:
589
+ im0 (torch.Tensor): Input tensor with shape (B, C, H, W).
590
+ """
591
+ self.im0 = self._single_check(im0)
592
+ self.bs = self.im0.shape[0]
593
+ self.mode = "image"
594
+ self.paths = [getattr(im, "filename", f"image{i}.jpg") for i, im in enumerate(im0)]
595
+
596
+ @staticmethod
597
+ def _single_check(im: torch.Tensor, stride: int = 32) -> torch.Tensor:
598
+ """Validate and format a single image tensor, ensuring correct shape and normalization."""
599
+ s = (
600
+ f"torch.Tensor inputs should be BCHW i.e. shape(1, 3, 640, 640) "
601
+ f"divisible by stride {stride}. Input shape{tuple(im.shape)} is incompatible."
602
+ )
603
+ if len(im.shape) != 4:
604
+ if len(im.shape) != 3:
605
+ raise ValueError(s)
606
+ LOGGER.warning(s)
607
+ im = im.unsqueeze(0)
608
+ if im.shape[2] % stride or im.shape[3] % stride:
609
+ raise ValueError(s)
610
+ if im.max() > 1.0 + torch.finfo(im.dtype).eps: # torch.float32 eps is 1.2e-07
611
+ LOGGER.warning(
612
+ f"torch.Tensor inputs should be normalized 0.0-1.0 but max value is {im.max()}. Dividing input by 255."
613
+ )
614
+ im = im.float() / 255.0
615
+
616
+ return im
617
+
618
+ def __iter__(self):
619
+ """Yield an iterator object for iterating through tensor image data."""
620
+ self.count = 0
621
+ return self
622
+
623
+ def __next__(self) -> tuple[list[str], torch.Tensor, list[str]]:
624
+ """Yield the next batch of tensor images and metadata for processing."""
625
+ if self.count == 1:
626
+ raise StopIteration
627
+ self.count += 1
628
+ return self.paths, self.im0, [""] * self.bs
629
+
630
+ def __len__(self) -> int:
631
+ """Return the batch size of the tensor input."""
632
+ return self.bs
633
+
634
+
635
+ def autocast_list(source: list[Any]) -> list[Image.Image | np.ndarray]:
636
+ """Merge a list of sources into a list of numpy arrays or PIL images for Ultralytics prediction."""
637
+ files = []
638
+ for im in source:
639
+ if isinstance(im, (str, Path)): # filename or uri
640
+ files.append(Image.open(urllib.request.urlopen(im) if str(im).startswith("http") else im))
641
+ elif isinstance(im, (Image.Image, np.ndarray)): # PIL or np Image
642
+ files.append(im)
643
+ else:
644
+ raise TypeError(
645
+ f"type {type(im).__name__} is not a supported Ultralytics prediction source type. \n"
646
+ f"See https://docs.ultralytics.com/modes/predict for supported source types."
647
+ )
648
+
649
+ return files
650
+
651
+
652
+ def get_best_youtube_url(url: str, method: str = "pytube") -> str | None:
653
+ """Retrieve the URL of the best quality MP4 video stream from a given YouTube video.
654
+
655
+ Args:
656
+ url (str): The URL of the YouTube video.
657
+ method (str): The method to use for extracting video info. Options are "pytube", "pafy", and "yt-dlp".
658
+
659
+ Returns:
660
+ (str | None): The URL of the best quality MP4 video stream, or None if no suitable stream is found.
661
+
662
+ Examples:
663
+ >>> url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ"
664
+ >>> best_url = get_best_youtube_url(url)
665
+ >>> print(best_url)
666
+ https://rr4---sn-q4flrnek.googlevideo.com/videoplayback?expire=...
667
+
668
+ Notes:
669
+ - Requires additional libraries based on the chosen method: pytubefix, pafy, or yt-dlp.
670
+ - The function prioritizes streams with at least 1080p resolution when available.
671
+ - For the "yt-dlp" method, it looks for formats with video codec, no audio, and *.mp4 extension.
672
+ """
673
+ if method == "pytube":
674
+ # Switched from pytube to pytubefix to resolve https://github.com/pytube/pytube/issues/1954
675
+ check_requirements("pytubefix>=6.5.2")
676
+ from pytubefix import YouTube
677
+
678
+ streams = YouTube(url).streams.filter(file_extension="mp4", only_video=True)
679
+ streams = sorted(streams, key=lambda s: s.resolution, reverse=True) # sort streams by resolution
680
+ for stream in streams:
681
+ if stream.resolution and int(stream.resolution[:-1]) >= 1080: # check if resolution is at least 1080p
682
+ return stream.url
683
+
684
+ elif method == "pafy":
685
+ check_requirements(("pafy", "youtube_dl==2020.12.2"))
686
+ import pafy
687
+
688
+ return pafy.new(url).getbestvideo(preftype="mp4").url
689
+
690
+ elif method == "yt-dlp":
691
+ check_requirements("yt-dlp")
692
+ import yt_dlp
693
+
694
+ with yt_dlp.YoutubeDL({"quiet": True}) as ydl:
695
+ info_dict = ydl.extract_info(url, download=False) # extract info
696
+ for f in reversed(info_dict.get("formats", [])): # reversed because best is usually last
697
+ # Find a format with video codec, no audio, *.mp4 extension at least 1920x1080 size
698
+ good_size = (f.get("width") or 0) >= 1920 or (f.get("height") or 0) >= 1080
699
+ if good_size and f["vcodec"] != "none" and f["acodec"] == "none" and f["ext"] == "mp4":
700
+ return f.get("url")
701
+
702
+
703
+ # Define constants
704
+ LOADERS = (LoadStreams, LoadPilAndNumpy, LoadImagesAndVideos, LoadScreenshots)