Merge branch 'master' into Maj-Pierre-1

pull/11869/head
Ultralytics Assistant 2024-07-08 22:06:22 +02:00 committed by GitHub
commit c81d5bf5d6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 275 additions and 250 deletions

View File

@ -68,7 +68,7 @@ def run(
imgsz (int): Inference size in pixels (default: 640). imgsz (int): Inference size in pixels (default: 640).
batch_size (int): Batch size for inference (default: 1). batch_size (int): Batch size for inference (default: 1).
data (Path | str): Path to the dataset.yaml file (default: ROOT / "data/coco128.yaml"). data (Path | str): Path to the dataset.yaml file (default: ROOT / "data/coco128.yaml").
device (str): CUDA device, e.g., '0' or '0,1,2,3' or 'cpu' (default: None). device (str): CUDA device, e.g., '0' or '0,1,2,3' or 'cpu' (default: "").
half (bool): Use FP16 half-precision inference (default: False). half (bool): Use FP16 half-precision inference (default: False).
test (bool): Test export formats only (default: False). test (bool): Test export formats only (default: False).
pt_only (bool): Test PyTorch format only (default: False). pt_only (bool): Test PyTorch format only (default: False).
@ -175,6 +175,24 @@ def test(
Returns: Returns:
pd.DataFrame: DataFrame containing the results of the export tests, including format names and export statuses. pd.DataFrame: DataFrame containing the results of the export tests, including format names and export statuses.
Examples:
```python
$ python benchmarks.py --weights yolov5s.pt --img 640
```
Notes:
Supported export formats and models include PyTorch, TorchScript, ONNX, OpenVINO, TensorRT, CoreML, TensorFlow
SavedModel, TensorFlow GraphDef, TensorFlow Lite, and TensorFlow Edge TPU. Edge TPU and TF.js are unsupported.
Usage:
Install required packages:
$ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU support
$ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU support
$ pip install -U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com # TensorRT
Run export tests:
$ python benchmarks.py --weights yolov5s.pt --img 640
""" """
y, t = [], time.time() y, t = [], time.time()
device = select_device(device) device = select_device(device)
@ -213,8 +231,8 @@ def parse_opt():
half (bool): Use FP16 half-precision inference. This is a flag and defaults to False. half (bool): Use FP16 half-precision inference. This is a flag and defaults to False.
test (bool): Test exports only. This is a flag and defaults to False. test (bool): Test exports only. This is a flag and defaults to False.
pt_only (bool): Test PyTorch only. This is a flag and defaults to False. pt_only (bool): Test PyTorch only. This is a flag and defaults to False.
hard_fail (bool|str): Throw an error on benchmark failure. Can be a boolean or a string representing a minimum metric hard_fail (bool | str): Throw an error on benchmark failure. Can be a boolean or a string representing a minimum
floor, i.e., '0.29'. Defaults to False. metric floor, e.g., '0.29'. Defaults to False.
Returns: Returns:
argparse.Namespace: Parsed command-line arguments encapsulated in an argparse Namespace object. argparse.Namespace: Parsed command-line arguments encapsulated in an argparse Namespace object.

View File

@ -101,40 +101,40 @@ def run(
Runs YOLOv5 detection inference on various sources like images, videos, directories, streams, etc. Runs YOLOv5 detection inference on various sources like images, videos, directories, streams, etc.
Args: Args:
weights (str | Path): Path to the model weights file or a Triton URL. Default is 'yolov5s.pt'. weights (str | Path): Path to the model weights file or a Triton URL. Default is 'yolov5s.pt'.
source (str | Path): Input source, which can be a file, directory, URL, glob pattern, screen capture, or webcam index. source (str | Path): Input source, which can be a file, directory, URL, glob pattern, screen capture, or webcam index.
Default is 'data/images'. Default is 'data/images'.
data (str | Path): Path to the dataset YAML file. Default is 'data/coco128.yaml'. data (str | Path): Path to the dataset YAML file. Default is 'data/coco128.yaml'.
imgsz (tuple[int, int]): Inference image size as a tuple (height, width). Default is (640, 640). imgsz (tuple[int, int]): Inference image size as a tuple (height, width). Default is (640, 640).
conf_thres (float): Confidence threshold for detections. Default is 0.25. conf_thres (float): Confidence threshold for detections. Default is 0.25.
iou_thres (float): Intersection Over Union (IOU) threshold for non-max suppression. Default is 0.45. iou_thres (float): Intersection Over Union (IOU) threshold for non-max suppression. Default is 0.45.
max_det (int): Maximum number of detections per image. Default is 1000. max_det (int): Maximum number of detections per image. Default is 1000.
device (str): CUDA device identifier (e.g., '0' or '0,1,2,3') or 'cpu'. Default is an empty string, which device (str): CUDA device identifier (e.g., '0' or '0,1,2,3') or 'cpu'. Default is an empty string, which
uses the best available device. uses the best available device.
view_img (bool): If True, display inference results using OpenCV. Default is False. view_img (bool): If True, display inference results using OpenCV. Default is False.
save_txt (bool): If True, save results in a text file. Default is False. save_txt (bool): If True, save results in a text file. Default is False.
save_csv (bool): If True, save results in a CSV file. Default is False. save_csv (bool): If True, save results in a CSV file. Default is False.
save_conf (bool): If True, include confidence scores in the saved results. Default is False. save_conf (bool): If True, include confidence scores in the saved results. Default is False.
save_crop (bool): If True, save cropped prediction boxes. Default is False. save_crop (bool): If True, save cropped prediction boxes. Default is False.
nosave (bool): If True, do not save inference images or videos. Default is False. nosave (bool): If True, do not save inference images or videos. Default is False.
classes (list[int]): List of class indices to filter detections by. Default is None. classes (list[int]): List of class indices to filter detections by. Default is None.
agnostic_nms (bool): If True, perform class-agnostic non-max suppression. Default is False. agnostic_nms (bool): If True, perform class-agnostic non-max suppression. Default is False.
augment (bool): If True, use augmented inference. Default is False. augment (bool): If True, use augmented inference. Default is False.
visualize (bool): If True, visualize feature maps. Default is False. visualize (bool): If True, visualize feature maps. Default is False.
update (bool): If True, update all models' weights. Default is False. update (bool): If True, update all models' weights. Default is False.
project (str | Path): Directory to save results. Default is 'runs/detect'. project (str | Path): Directory to save results. Default is 'runs/detect'.
name (str): Name of the current experiment; used to create a subdirectory within 'project'. Default is 'exp'. name (str): Name of the current experiment; used to create a subdirectory within 'project'. Default is 'exp'.
exist_ok (bool): If True, existing directories with the same name are reused instead of being incremented. Default is exist_ok (bool): If True, existing directories with the same name are reused instead of being incremented. Default is
False. False.
line_thickness (int): Thickness of bounding box lines in pixels. Default is 3. line_thickness (int): Thickness of bounding box lines in pixels. Default is 3.
hide_labels (bool): If True, do not display labels on bounding boxes. Default is False. hide_labels (bool): If True, do not display labels on bounding boxes. Default is False.
hide_conf (bool): If True, do not display confidence scores on bounding boxes. Default is False. hide_conf (bool): If True, do not display confidence scores on bounding boxes. Default is False.
half (bool): If True, use FP16 half-precision inference. Default is False. half (bool): If True, use FP16 half-precision inference. Default is False.
dnn (bool): If True, use OpenCV DNN backend for ONNX inference. Default is False. dnn (bool): If True, use OpenCV DNN backend for ONNX inference. Default is False.
vid_stride (int): Stride for processing video frames, to skip frames between processing. Default is 1. vid_stride (int): Stride for processing video frames, to skip frames between processing. Default is 1.
Returns: Returns:
None None
Examples: Examples:
```python ```python
@ -397,15 +397,15 @@ def main(opt):
Executes YOLOv5 model inference based on provided command-line arguments, validating dependencies before running. Executes YOLOv5 model inference based on provided command-line arguments, validating dependencies before running.
Args: Args:
opt (argparse.Namespace): Command-line arguments for YOLOv5 detection. See function `parse_opt` for details. opt (argparse.Namespace): Command-line arguments for YOLOv5 detection. See function `parse_opt` for details.
Returns: Returns:
None None
Note: Note:
This function performs essential pre-execution checks and initiates the YOLOv5 detection process based on user-specified options. This function performs essential pre-execution checks and initiates the YOLOv5 detection process based on user-specified options.
Refer to the usage guide and examples for more information about different sources and formats at: Refer to the usage guide and examples for more information about different sources and formats at:
https://github.com/ultralytics/ultralytics https://github.com/ultralytics/ultralytics
Example usage: Example usage:

202
export.py
View File

@ -103,10 +103,9 @@ class iOSModel(torch.nn.Module):
None: This method does not return any value. None: This method does not return any value.
Notes: Notes:
This initializer configures normalization based on the input image dimensions, which is critical for This initializer configures normalization based on the input image dimensions, which is critical for ensuring the
ensuring the model's compatibility and proper functionality on iOS devices. The normalization step model's compatibility and proper functionality on iOS devices. The normalization step involves dividing by the image
involves dividing by the image width if the image is square; otherwise, additional conditions might width if the image is square; otherwise, additional conditions might apply (trimmed for brevity).
apply (trimmed for brevity).
""" """
super().__init__() super().__init__()
b, c, h, w = im.shape # batch, channel, height, width b, c, h, w = im.shape # batch, channel, height, width
@ -124,16 +123,17 @@ class iOSModel(torch.nn.Module):
Runs a forward pass on the input tensor, returning class confidences and normalized coordinates. Runs a forward pass on the input tensor, returning class confidences and normalized coordinates.
Args: Args:
x (torch.Tensor): Input tensor containing the image data. x (torch.Tensor): Input tensor containing the image data.
Returns: Returns:
torch.Tensor: Concatenated tensor with normalized coordinates (xywh), confidence scores (conf), and class probabilities (cls). torch.Tensor: Concatenated tensor with normalized coordinates (xywh), confidence scores (conf),
and class probabilities (cls).
Examples: Examples:
```python ```python
model = iOSModel(pretrained_model, input_image) model = iOSModel(pretrained_model, input_image)
output = model.forward(torch_input_tensor) output = model.forward(torch_input_tensor)
``` ```
""" """
xywh, conf, cls = self.model(x)[0].squeeze().split((4, 1, self.nc), 1) xywh, conf, cls = self.model(x)[0].squeeze().split((4, 1, self.nc), 1)
return cls * conf, xywh * self.normalize # confidence (3780, 80), coordinates (3780, 4) return cls * conf, xywh * self.normalize # confidence (3780, 80), coordinates (3780, 4)
@ -151,7 +151,7 @@ def export_formats():
Examples: Examples:
```python ```python
formats = export_formats() formats = export_formats()
print(formats) print(f"Supported export formats:\n{formats}")
``` ```
Notes: Notes:
@ -188,16 +188,18 @@ def try_export(inner_func):
Returns: Returns:
Callable: The wrapped function that logs execution details. When executed, this wrapper function returns either: Callable: The wrapped function that logs execution details. When executed, this wrapper function returns either:
- Tuple (str, torch.nn.Module): On success the file path of the exported model and the model instance. - Tuple (str | torch.nn.Module): On success the file path of the exported model and the model instance.
- Tuple (None, None): On failure None values indicating export failed. - Tuple (None, None): On failure None values indicating export failure.
Examples: Examples:
```python
@try_export @try_export
def export_onnx(model, filepath): def export_onnx(model, filepath):
# implementation here # implementation here
pass pass
exported_file, exported_model = export_onnx(yolo_model, 'path/to/save/model.onnx') exported_file, exported_model = export_onnx(yolo_model, 'path/to/save/model.onnx')
```
Notes: Notes:
For additional requirements and model export formats, refer to the [Ultralytics YOLOv5 GitHub repository](https://github.com/ultralytics/ultralytics). For additional requirements and model export formats, refer to the [Ultralytics YOLOv5 GitHub repository](https://github.com/ultralytics/ultralytics).
@ -239,7 +241,7 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:"
Notes: Notes:
- This function uses tracing to create the TorchScript model. - This function uses tracing to create the TorchScript model.
- Metadata, including the input image shape, model stride, and class names, is saved in an extra file (`config.txt`) - Metadata, including the input image shape, model stride, and class names, is saved in an extra file (`config.txt`)
within the TorchScript model package. within the TorchScript model package.
- For mobile optimization, refer to the PyTorch tutorial: https://pytorch.org/tutorials/recipes/mobile_interpreter.html - For mobile optimization, refer to the PyTorch tutorial: https://pytorch.org/tutorials/recipes/mobile_interpreter.html
Example: Example:
@ -359,8 +361,7 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX
@try_export @try_export
def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")): def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")):
""" """
Exports a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization; see Exports a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization.
https://pypi.org/project/openvino-dev/.
Args: Args:
file (Path): The path to the output file where the OpenVINO model will be saved. file (Path): The path to the output file where the OpenVINO model will be saved.
@ -450,7 +451,7 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:
@try_export @try_export
def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")): def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
""" """
Exports a YOLOv5 model to PaddlePaddle format using X2Paddle, saving the converted model and metadata. Export a YOLOv5 model to PaddlePaddle format using X2Paddle, saving the converted model and metadata.
Args: Args:
model (torch.nn.Module): The YOLOv5 model to be exported. model (torch.nn.Module): The YOLOv5 model to be exported.
@ -476,6 +477,7 @@ def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
export_paddle(model=model, im=im, file=file, metadata=metadata) export_paddle(model=model, im=im, file=file, metadata=metadata)
``` ```
Notes: Notes:
Ensure that `paddlepaddle` and `x2paddle` are installed, as these are required for the export function. You can Ensure that `paddlepaddle` and `x2paddle` are installed, as these are required for the export function. You can
install them via pip: install them via pip:
@ -556,7 +558,7 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose
Args: Args:
model (torch.nn.Module): YOLOv5 model to be exported. model (torch.nn.Module): YOLOv5 model to be exported.
im (torch.Tensor): Input tensor of shape (B,C,H,W). im (torch.Tensor): Input tensor of shape (B, C, H, W).
file (Path): Path to save the exported model. file (Path): Path to save the exported model.
half (bool): Set to True to export with FP16 precision. half (bool): Set to True to export with FP16 precision.
dynamic (bool): Set to True to enable dynamic input shapes. dynamic (bool): Set to True to enable dynamic input shapes.
@ -664,7 +666,8 @@ def export_saved_model(
prefix=colorstr("TensorFlow SavedModel:"), prefix=colorstr("TensorFlow SavedModel:"),
): ):
""" """
Exports a YOLOv5 model to TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression (NMS). Exports a YOLOv5 model to the TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression
(NMS).
Args: Args:
model (torch.nn.Module): The PyTorch model to convert. model (torch.nn.Module): The PyTorch model to convert.
@ -681,7 +684,8 @@ def export_saved_model(
prefix (str, optional): Prefix for logging messages. Default is "TensorFlow SavedModel:". prefix (str, optional): Prefix for logging messages. Default is "TensorFlow SavedModel:".
Returns: Returns:
tuple: A tuple containing the path to the saved model folder (str) and the Keras model instance (tf.keras.Model | None). tuple[str, tf.keras.Model | None]: A tuple containing the path to the saved model folder and the Keras model instance,
or None if TensorFlow export fails.
Notes: Notes:
- The method supports TensorFlow versions up to 2.15.1. - The method supports TensorFlow versions up to 2.15.1.
@ -689,9 +693,6 @@ def export_saved_model(
- If the TensorFlow version exceeds 2.13.1, it might cause issues when exporting to TFLite. - If the TensorFlow version exceeds 2.13.1, it might cause issues when exporting to TFLite.
Refer to: https://github.com/ultralytics/yolov5/issues/12489 Refer to: https://github.com/ultralytics/yolov5/issues/12489
Raises:
Exception: If TensorFlow is not installed.
Example: Example:
```python ```python
model, im = ... # Initialize your PyTorch model and input tensor model, im = ... # Initialize your PyTorch model and input tensor
@ -749,7 +750,7 @@ def export_saved_model(
@try_export @try_export
def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")): def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")):
""" """
Exports YOLOv5 model to TensorFlow GraphDef (*.pb) format. Export YOLOv5 model to TensorFlow GraphDef (*.pb) format.
Args: Args:
keras_model (tf.keras.Model): The Keras model to be converted. keras_model (tf.keras.Model): The Keras model to be converted.
@ -790,44 +791,43 @@ def export_tflite(
): ):
# YOLOv5 TensorFlow Lite export # YOLOv5 TensorFlow Lite export
""" """
Exports YOLOv5 model to TensorFlow Lite format with optional FP16, INT8, and NMS support. Exports a YOLOv5 model to TensorFlow Lite format with optional INT8 quantization and NMS support.
Args: Args:
keras_model (tf.keras.Model): The Keras model to be exported. keras_model (tf.keras.Model): The Keras model to be exported.
im (torch.Tensor): Image tensor for normalization and model tracing. im (torch.Tensor): An input image tensor for normalization and model tracing.
file (Path): The file path to save the exported TensorFlow Lite model. file (Path): The file path to save the TensorFlow Lite model.
int8 (bool): Enables INT8 quantization if True. int8 (bool): Enables INT8 quantization if True.
per_tensor (bool): If True, disable per-channel quantization (applicable when int8 is True). per_tensor (bool): If True, disables per-channel quantization.
data (str): Path to dataset for representative dataset generation in INT8 quantization. data (str): Path to the dataset for representative dataset generation in INT8 quantization.
nms (bool): Enables Non-Maximum Suppression (NMS) support if True. nms (bool): Enables Non-Maximum Suppression (NMS) if True.
agnostic_nms (bool): Enables class-agnostic NMS support if True. agnostic_nms (bool): Enables class-agnostic NMS if True.
prefix (str): Prefix for logging messages. prefix (str): Prefix for log messages.
Returns: Returns:
(str | None, tf.lite.Model | None): The file path of the saved TFLite model, and the TFLite model instance if successful. (str | None, tflite.Model | None): The file path of the exported TFLite model and the TFLite model instance, or None if export failed.
Example: Example:
```python ```python
from pathlib import Path from pathlib import Path
import torch import torch
import tensorflow as tf import tensorflow as tf
from torchvision import models
# Load a pre-trained model from torchvision # Load a Keras model wrapping a YOLOv5 model
model = models.yolov5() # Placeholder for actual YOLOv5 model loading
im = torch.zeros(1, 3, 640, 640) # Example image tensor
# Provide the Keras model wrapping the PyTorch YOLOv5 model
keras_model = tf.keras.models.load_model('path/to/keras_model.h5') keras_model = tf.keras.models.load_model('path/to/keras_model.h5')
# Export the model to TensorFlow Lite format # Example input tensor
file_path = export_tflite(keras_model, im, Path('model.tflite'), int8=False, per_tensor=False, im = torch.zeros(1, 3, 640, 640)
data='path/to/dataset.yaml', nms=False, agnostic_nms=False)
# Export the model
export_tflite(keras_model, im, Path('model.tflite'), int8=True, per_tensor=False, data='data/coco.yaml',
nms=True, agnostic_nms=False)
``` ```
Notes: Notes:
Ensure the TensorFlow and TensorFlow Lite dependencies are installed. The exported TFLite model can be used for - Ensure TensorFlow and TensorFlow Lite dependencies are installed.
efficient inference on mobile and edge devices. - INT8 quantization requires a representative dataset to achieve optimal accuracy.
- TensorFlow Lite models are suitable for efficient inference on mobile and edge devices.
""" """
import tensorflow as tf import tensorflow as tf
@ -935,18 +935,18 @@ def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
prefix (str): Prefix for logging information (default: colorstr("TensorFlow.js:")). prefix (str): Prefix for logging information (default: colorstr("TensorFlow.js:")).
Returns: Returns:
tuple: Output directory path (str), None (str, None): The output directory path as a string and None.
Notes: Notes:
This function requires `tensorflowjs` to be installed. You can install it using: This function requires `tensorflowjs` to be installed. You can install it using:
```shell ```shell
pip install tensorflowjs pip install tensorflowjs
``` ```
Example usage: Example usage:
```python ```python
export_tfjs(Path('yolov5s.onnx'), int8=False) export_tfjs(Path('yolov5s.onnx'), int8=False)
``` ```
The TensorFlow.js converted model is saved in the directory specified by `file` with "_web_model" suffix. The TensorFlow.js converted model is saved in the directory specified by `file` with "_web_model" suffix.
""" """
@ -987,12 +987,12 @@ def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
def add_tflite_metadata(file, metadata, num_outputs): def add_tflite_metadata(file, metadata, num_outputs):
""" """
Adds TFLite metadata to a model file, supporting multiple outputs, as specified by TensorFlow guidelines. Adds metadata to a TensorFlow Lite (TFLite) model file, supporting multiple outputs, based on TensorFlow guidelines.
Args: Args:
file (str): The path to the TensorFlow Lite model file to which metadata will be added. file (str): Path to the TFLite model file to which metadata will be added.
metadata (dict): Metadata information to be added to the model, structured as required by TFLite metadata schema. metadata (dict): Metadata information to be added to the model, structured as required by the TFLite metadata schema.
num_outputs (int): Number of output tensors the model has, to properly configure the metadata. num_outputs (int): Number of output tensors the model has, used to configure the metadata properly.
Returns: Returns:
None None
@ -1011,7 +1011,7 @@ def add_tflite_metadata(file, metadata, num_outputs):
Note: Note:
TFLite metadata can include information such as model name, version, author, and other relevant details. TFLite metadata can include information such as model name, version, author, and other relevant details.
For more details and structure of the metadata, refer to the TensorFlow Lite For more details on the structure of the metadata, refer to TensorFlow Lite
[metadata guidelines](https://www.tensorflow.org/lite/models/convert/metadata). [metadata guidelines](https://www.tensorflow.org/lite/models/convert/metadata).
""" """
with contextlib.suppress(ImportError): with contextlib.suppress(ImportError):
@ -1071,24 +1071,24 @@ def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:
- Flexible input shapes and additional NMS options can be customized within the function. - Flexible input shapes and additional NMS options can be customized within the function.
Examples: Examples:
```python ```python
from pathlib import Path from pathlib import Path
import torch import torch
# Load YOLOv5 model and an example input tensor # Load YOLOv5 model and an example input tensor
model = torch.load("yolov5s.pt") model = torch.load("yolov5s.pt")
im = torch.zeros(1, 3, 640, 640) # Example input tensor im = torch.zeros(1, 3, 640, 640) # Example input tensor
# Define class names # Define class names
names = {0: "person", 1: "bicycle", 2: "car", ...} names = {0: "person", 1: "bicycle", 2: "car", ...}
# Perform forward pass to get model output # Perform forward pass to get model output
y = model(im) y = model(im)
# Convert to CoreML # Convert to CoreML
output_file = Path("yolov5s.mlmodel") output_file = Path("yolov5s.mlmodel")
pipeline_coreml(model, im, output_file, names, y) pipeline_coreml(model, im, output_file, names, y)
``` ```
""" """
import coremltools as ct import coremltools as ct
from PIL import Image from PIL import Image
@ -1246,6 +1246,8 @@ def run(
conf_thres=0.25, # TF.js NMS: confidence threshold conf_thres=0.25, # TF.js NMS: confidence threshold
): ):
""" """
Clear and concise summary line describing the function's purpose:
Exports a YOLOv5 model to specified formats including ONNX, TensorRT, CoreML, and TensorFlow. Exports a YOLOv5 model to specified formats including ONNX, TensorRT, CoreML, and TensorFlow.
Args: Args:
@ -1474,42 +1476,42 @@ def parse_opt(known=False):
def main(opt): def main(opt):
""" """
```python Exports the YOLOv5 model to specified formats, including ONNX, TensorRT, CoreML, and TensorFlow. Exports the YOLOv5 model to specified formats, including ONNX, TensorRT, CoreML, and TensorFlow.
Args: Args:
opt (argparse.Namespace): Parsed command-line arguments containing the export configurations. opt (argparse.Namespace): Parsed command-line arguments containing the export configurations.
- data (str): Path to the dataset.yaml. - data (str): Path to the dataset YAML configuration file (e.g., 'data/coco128.yaml').
- weights (list[str]): Paths to model (.pt) file(s). - weights (list[str] | str): Paths to the pretrained model weights file(s) (e.g., 'yolov5s.pt').
- imgsz (list[int]): Image size (height, width). - imgsz (list[int]): Image size as a list [height, width].
- batch_size (int): Batch size. - batch_size (int): Batch size for exporting the model.
- device (str): CUDA device, e.g., '0' or '0,1,2,3' or 'cpu'. - device (str): Device to run the export on, such as '0' for GPU, or 'cpu' for CPU.
- half (bool): FP16 half-precision export flag. - half (bool): Flag to export the model with FP16 half-precision.
- inplace (bool): Set YOLOv5 Detect() inplace to True. - inplace (bool): Set the YOLOv5 Detect() module inplace mode to True.
- keras (bool): Use Keras for TensorFlow models. - keras (bool): Flag to use Keras for TensorFlow SavedModel export.
- optimize (bool): Optimize TorchScript model for mobile. - optimize (bool): Optimize TorchScript model for mobile deployment.
- int8 (bool): INT8 quantization flag. - int8 (bool): Apply INT8 quantization for CoreML or TensorFlow models.
- per_tensor (bool): Per tensor quantization for TensorFlow. - per_tensor (bool): Apply per-tensor quantization for TensorFlow models.
- dynamic (bool): Dynamic axes for ONNX/TF/TensorRT. - dynamic (bool): Enable dynamic axes for ONNX, TensorFlow, or TensorRT exports.
- simplify (bool): Simplify ONNX model. - simplify (bool): Simplify ONNX model during export.
- opset (int): ONNX opset version. - opset (int): ONNX opset version.
- verbose (bool): Verbose logging for TensorRT. - verbose (bool): Enable verbose logging for TensorRT export.
- workspace (int): Workspace size for TensorRT (in GB). - workspace (int): TensorRT workspace size in GB.
- nms (bool): Add NMS to TensorFlow model. - nms (bool): Add non-maximum suppression (NMS) to the TensorFlow model.
- agnostic_nms (bool): Add agnostic NMS to TensorFlow model. - agnostic_nms (bool): Add class-agnostic NMS to the TensorFlow model.
- topk_per_class (int): Top-k per class for TensorFlow.js NMS. - topk_per_class (int): Top-K boxes per class to keep for TensorFlow.js NMS.
- topk_all (int): Top-k for all classes for TensorFlow.js NMS. - topk_all (int): Top-K boxes for all classes to keep for TensorFlow.js NMS.
- iou_thres (float): IoU threshold for TensorFlow.js NMS. - iou_thres (float): IoU threshold for NMS.
- conf_thres (float): Confidence threshold for TensorFlow.js NMS. - conf_thres (float): Confidence threshold for NMS.
- include (list[str]): List of formats to include in export, e.g., ['torchscript', 'onnx']. - include (list[str]): List of formats to include in export (e.g., ['torchscript', 'onnx']).
Returns: Returns:
list[str]: List of exported file paths. list[str]: List of paths to the exported model files.
```python
# Example usage: Example:
# opt = parse_opt() ```python
# main(opt) opt = parse_opt()
``` main(opt)
```
""" """
for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]: for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
run(**vars(opt)) run(**vars(opt))

View File

@ -188,8 +188,7 @@ def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
auto-select the best available device. Defaults to None. auto-select the best available device. Defaults to None.
Returns: Returns:
YOLOv5 model (torch.nn.Module): The YOLOv5-small model loaded with specified configurations and optionally torch.nn.Module: The YOLOv5-small model loaded with specified configurations and optionally pretrained weights.
pretrained weights.
Usage: Usage:
```python ```python
@ -245,10 +244,12 @@ def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
classes (int): Number of model classes. Default is 80. classes (int): Number of model classes. Default is 80.
autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model. Default is True. autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model. Default is True.
_verbose (bool): Print all information to screen. Default is True. _verbose (bool): Print all information to screen. Default is True.
device (str | torch.device | None): Device to use for model parameters, e.g., 'cpu', 'cuda', or a torch.device instance. Default is None. device (str | torch.device | None): Device to use for model parameters, e.g., 'cpu', 'cuda', or a torch.device instance.
Default is None.
Returns: Returns:
YOLOv5 model (torch.nn.Module). YOLOv5 model (torch.nn.Module): The YOLOv5-large model instantiated with specified configurations and possibly
pretrained weights.
Example: Example:
```python ```python
@ -368,9 +369,11 @@ def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
pretrained (bool): If True, loads pretrained weights. Default is True. pretrained (bool): If True, loads pretrained weights. Default is True.
channels (int): Number of input channels. Default is 3. channels (int): Number of input channels. Default is 3.
classes (int): Number of model classes. Default is 80. classes (int): Number of model classes. Default is 80.
autoshape (bool): Apply YOLOv5 .autoshape() wrapper to the model for file/URI/PIL/cv2/np inputs and NMS. Default is True. autoshape (bool): Apply YOLOv5 .autoshape() wrapper to the model for file/URI/PIL/cv2/np inputs and NMS.
Default is True.
_verbose (bool): If True, prints detailed information to the screen. Default is True. _verbose (bool): If True, prints detailed information to the screen. Default is True.
device (str | torch.device | None): Device to use for model parameters. Default is None, which uses the best available device. device (str | torch.device | None): Device to use for model parameters. Default is None, which uses the
best available device.
Returns: Returns:
torch.nn.Module: The YOLOv5-medium-P6 model. torch.nn.Module: The YOLOv5-medium-P6 model.
@ -378,12 +381,12 @@ def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
Refer to the PyTorch Hub models documentation: https://pytorch.org/hub/ultralytics_yolov5 for additional details. Refer to the PyTorch Hub models documentation: https://pytorch.org/hub/ultralytics_yolov5 for additional details.
Example: Example:
```python ```python
import torch import torch
# Load YOLOv5-medium-P6 model # Load YOLOv5-medium-P6 model
model = torch.hub.load('ultralytics/yolov5', 'yolov5m6') model = torch.hub.load('ultralytics/yolov5', 'yolov5m6')
``` ```
Notes: Notes:
- The model can be loaded with pre-trained weights for better performance on specific tasks. - The model can be loaded with pre-trained weights for better performance on specific tasks.
@ -401,12 +404,10 @@ def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
pretrained (bool, optional): If True, load pretrained weights into the model. Default is True. pretrained (bool, optional): If True, load pretrained weights into the model. Default is True.
channels (int, optional): Number of input channels. Default is 3. channels (int, optional): Number of input channels. Default is 3.
classes (int, optional): Number of model classes. Default is 80. classes (int, optional): Number of model classes. Default is 80.
autoshape (bool, optional): If True, apply YOLOv5 .autoshape() wrapper to the model for input flexibility. autoshape (bool, optional): If True, apply YOLOv5 .autoshape() wrapper to the model for input flexibility. Default is True.
Default is True.
_verbose (bool, optional): If True, print all information to the screen. Default is True. _verbose (bool, optional): If True, print all information to the screen. Default is True.
device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda', or device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda', or torch.device.
torch.device. If None, automatically selects the best available If None, automatically selects the best available device. Default is None.
device. Default is None.
Returns: Returns:
torch.nn.Module: The instantiated YOLOv5-large-P6 model. torch.nn.Module: The instantiated YOLOv5-large-P6 model.
@ -444,10 +445,10 @@ def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
torch.nn.Module: The instantiated YOLOv5-xlarge-P6 model. torch.nn.Module: The instantiated YOLOv5-xlarge-P6 model.
Example: Example:
```python ```python
import torch import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5x6') # load the YOLOv5-xlarge-P6 model model = torch.hub.load('ultralytics/yolov5', 'yolov5x6') # load the YOLOv5-xlarge-P6 model
``` ```
Note: Note:
For more information on YOLOv5 models, visit the official documentation: For more information on YOLOv5 models, visit the official documentation:

107
train.py
View File

@ -546,10 +546,10 @@ def parse_opt(known=False):
Parses command-line arguments for YOLOv5 training, validation, and testing. Parses command-line arguments for YOLOv5 training, validation, and testing.
Args: Args:
known (bool, optional): If True, parses known arguments, ignoring the unknown. Defaults to False. known (bool, optional): If True, parses known arguments, ignoring the unknown. Defaults to False.
Returns: Returns:
argparse.Namespace: Parsed command-line arguments. argparse.Namespace: Parsed command-line arguments.
Example: Example:
```python ```python
@ -559,9 +559,9 @@ def parse_opt(known=False):
``` ```
Links: Links:
Models: https://github.com/ultralytics/yolov5/tree/master/models Models: https://github.com/ultralytics/yolov5/tree/master/models
Datasets: https://github.com/ultralytics/yolov5/tree/master/data Datasets: https://github.com/ultralytics/yolov5/tree/master/data
Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data
""" """
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path") parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")
@ -621,15 +621,15 @@ def main(opt, callbacks=Callbacks()):
Runs training or hyperparameter evolution with specified options and optional callbacks. Runs training or hyperparameter evolution with specified options and optional callbacks.
Args: Args:
opt (argparse.Namespace): The command-line arguments parsed for YOLOv5 training and evolution. opt (argparse.Namespace): The command-line arguments parsed for YOLOv5 training and evolution.
callbacks (ultralytics.utils.callbacks.Callbacks, optional): Callback functions for various training stages. callbacks (ultralytics.utils.callbacks.Callbacks, optional): Callback functions for various training stages.
Defaults to Callbacks(). Defaults to Callbacks().
Returns: Returns:
None None
Note: Note:
For detailed usage, visit: For detailed usage, visit:
https://github.com/ultralytics/yolov5/tree/master/models https://github.com/ultralytics/yolov5/tree/master/models
""" """
if RANK in {-1, 0}: if RANK in {-1, 0}:
@ -918,51 +918,56 @@ def run(**kwargs):
Executes YOLOv5 training with given options, allowing optional overrides through keyword arguments. Executes YOLOv5 training with given options, allowing optional overrides through keyword arguments.
Args: Args:
weights (str): Path to initial weights. Defaults to ROOT / 'yolov5s.pt'. weights (str, optional): Path to initial weights. Defaults to ROOT / 'yolov5s.pt'.
cfg (str): Path to model YAML configuration. Defaults to an empty string. cfg (str, optional): Path to model YAML configuration. Defaults to an empty string.
data (str): Path to dataset YAML configuration. Defaults to ROOT / 'data/coco128.yaml'. data (str, optional): Path to dataset YAML configuration. Defaults to ROOT / 'data/coco128.yaml'.
hyp (str): Path to hyperparameters YAML configuration. Defaults to ROOT / 'data/hyps/hyp.scratch-low.yaml'. hyp (str, optional): Path to hyperparameters YAML configuration. Defaults to ROOT / 'data/hyps/hyp.scratch-low.yaml'.
epochs (int): Total number of training epochs. Defaults to 100. epochs (int, optional): Total number of training epochs. Defaults to 100.
batch_size (int): Total batch size for all GPUs. Use -1 for automatic batch size determination. Defaults to 16. batch_size (int, optional): Total batch size for all GPUs. Use -1 for automatic batch size determination. Defaults to 16.
imgsz (int): Image size (pixels) for training and validation. Defaults to 640. imgsz (int, optional): Image size (pixels) for training and validation. Defaults to 640.
rect (bool): Use rectangular training. Defaults to False. rect (bool, optional): Use rectangular training. Defaults to False.
resume (bool | str): Resume most recent training with an optional path. Defaults to False. resume (bool | str, optional): Resume most recent training with an optional path. Defaults to False.
nosave (bool): Only save final checkpoint. Defaults to False. nosave (bool, optional): Only save the final checkpoint. Defaults to False.
noval (bool): Only validate at the final epoch. Defaults to False. noval (bool, optional): Only validate at the final epoch. Defaults to False.
noautoanchor (bool): Disable AutoAnchor. Defaults to False. noautoanchor (bool, optional): Disable AutoAnchor. Defaults to False.
noplots (bool): Do not save plot files. Defaults to False. noplots (bool, optional): Do not save plot files. Defaults to False.
evolve (int): Evolve hyperparameters for a specified number of generations. Use 300 if provided without a value. evolve (int, optional): Evolve hyperparameters for a specified number of generations. Use 300 if provided without a value.
evolve_population (str): Directory for loading population during evolution. Defaults to ROOT / 'data/hyps'. evolve_population (str, optional): Directory for loading population during evolution. Defaults to ROOT / 'data/hyps'.
resume_evolve (str): Resume hyperparameter evolution from the last generation. Defaults to None. resume_evolve (str, optional): Resume hyperparameter evolution from the last generation. Defaults to None.
bucket (str): gsutil bucket for saving checkpoints. Defaults to an empty string. bucket (str, optional): gsutil bucket for saving checkpoints. Defaults to an empty string.
cache (str): Cache image data in 'ram' or 'disk'. Defaults to None. cache (str, optional): Cache image data in 'ram' or 'disk'. Defaults to None.
image_weights (bool): Use weighted image selection for training. Defaults to False. image_weights (bool, optional): Use weighted image selection for training. Defaults to False.
device (str): CUDA device identifier, e.g., '0', '0,1,2,3', or 'cpu'. Defaults to an empty string. device (str, optional): CUDA device identifier, e.g., '0', '0,1,2,3', or 'cpu'. Defaults to an empty string.
multi_scale (bool): Use multi-scale training, varying image size by ±50%. Defaults to False. multi_scale (bool, optional): Use multi-scale training, varying image size by ±50%. Defaults to False.
single_cls (bool): Train with multi-class data as single-class. Defaults to False. single_cls (bool, optional): Train with multi-class data as single-class. Defaults to False.
optimizer (str): Optimizer type, choices are ['SGD', 'Adam', 'AdamW']. Defaults to 'SGD'. optimizer (str, optional): Optimizer type, choices are ['SGD', 'Adam', 'AdamW']. Defaults to 'SGD'.
sync_bn (bool): Use synchronized BatchNorm, only available in DDP mode. Defaults to False. sync_bn (bool, optional): Use synchronized BatchNorm, only available in DDP mode. Defaults to False.
workers (int): Maximum dataloader workers per rank in DDP mode. Defaults to 8. workers (int, optional): Maximum dataloader workers per rank in DDP mode. Defaults to 8.
project (str): Directory for saving training runs. Defaults to ROOT / 'runs/train'. project (str, optional): Directory for saving training runs. Defaults to ROOT / 'runs/train'.
name (str): Name for saving the training run. Defaults to 'exp'. name (str, optional): Name for saving the training run. Defaults to 'exp'.
exist_ok (bool): Allow existing project/name without incrementing. Defaults to False. exist_ok (bool, optional): Allow existing project/name without incrementing. Defaults to False.
quad (bool): Use quad dataloader. Defaults to False. quad (bool, optional): Use quad dataloader. Defaults to False.
cos_lr (bool): Use cosine learning rate scheduler. Defaults to False. cos_lr (bool, optional): Use cosine learning rate scheduler. Defaults to False.
label_smoothing (float): Label smoothing epsilon value. Defaults to 0.0. label_smoothing (float, optional): Label smoothing epsilon value. Defaults to 0.0.
patience (int): Patience for early stopping, measured in epochs without improvement. Defaults to 100. patience (int, optional): Patience for early stopping, measured in epochs without improvement. Defaults to 100.
freeze (list): Layers to freeze, e.g., backbone=10, first 3 layers = [0, 1, 2]. Defaults to [0]. freeze (list, optional): Layers to freeze, e.g., backbone=10, first 3 layers = [0, 1, 2]. Defaults to [0].
save_period (int): Frequency in epochs to save checkpoints. Disabled if < 1. Defaults to -1. save_period (int, optional): Frequency in epochs to save checkpoints. Disabled if < 1. Defaults to -1.
seed (int): Global training random seed. Defaults to 0. seed (int, optional): Global training random seed. Defaults to 0.
local_rank (int): Automatic DDP Multi-GPU argument. Do not modify. Defaults to -1. local_rank (int, optional): Automatic DDP Multi-GPU argument. Do not modify. Defaults to -1.
Returns: Returns:
None: The function initiates YOLOv5 training or hyperparameter evolution based on the provided options. None: The function initiates YOLOv5 training or hyperparameter evolution based on the provided options.
Examples: Examples:
```python ```python
import train import train
train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt') train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
``` ```
Notes:
- Models: https://github.com/ultralytics/yolov5/tree/master/models
- Datasets: https://github.com/ultralytics/yolov5/tree/master/data
- Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data
""" """
opt = parse_opt(True) opt = parse_opt(True)
for k, v in kwargs.items(): for k, v in kwargs.items():

77
val.py
View File

@ -66,8 +66,8 @@ def save_one_txt(predn, save_conf, shape, file):
Saves one detection result to a txt file in normalized xywh format, optionally including confidence. Saves one detection result to a txt file in normalized xywh format, optionally including confidence.
Args: Args:
predn (torch.Tensor): Predicted bounding boxes and associated confidence scores and classes predn (torch.Tensor): Predicted bounding boxes and associated confidence scores and classes in xyxy format,
in xyxy format, tensor of shape (N, 6) where N is the number of detections. tensor of shape (N, 6) where N is the number of detections.
save_conf (bool): If True, saves the confidence scores along with the bounding box coordinates. save_conf (bool): If True, saves the confidence scores along with the bounding box coordinates.
shape (tuple): Shape of the original image as (height, width). shape (tuple): Shape of the original image as (height, width).
file (str | Path): File path where the result will be saved. file (str | Path): File path where the result will be saved.
@ -77,8 +77,8 @@ def save_one_txt(predn, save_conf, shape, file):
Notes: Notes:
The xyxy bounding box format represents the coordinates (xmin, ymin, xmax, ymax). The xyxy bounding box format represents the coordinates (xmin, ymin, xmax, ymax).
The xywh format represents the coordinates (center_x, center_y, width, height) and is The xywh format represents the coordinates (center_x, center_y, width, height) and is normalized by the width and
normalized by the width and height of the image. height of the image.
Example: Example:
```python ```python
@ -99,33 +99,33 @@ def save_one_json(predn, jdict, path, class_map):
Saves a single JSON detection result, including image ID, category ID, bounding box, and confidence score. Saves a single JSON detection result, including image ID, category ID, bounding box, and confidence score.
Args: Args:
predn (torch.Tensor): Predicted detections in xyxy format with shape (n, 6) where n is the number of detections. predn (torch.Tensor): Predicted detections in xyxy format with shape (n, 6) where n is the number of detections.
The tensor should contain [x_min, y_min, x_max, y_max, confidence, class_id] for each detection. The tensor should contain [x_min, y_min, x_max, y_max, confidence, class_id] for each detection.
jdict (list[dict]): List to collect JSON formatted detection results. jdict (list[dict]): List to collect JSON formatted detection results.
path (pathlib.Path): Path object of the image file, used to extract image_id. path (pathlib.Path): Path object of the image file, used to extract image_id.
class_map (dict[int, int]): Mapping from model class indices to dataset-specific category IDs. class_map (dict[int, int]): Mapping from model class indices to dataset-specific category IDs.
Returns: Returns:
None: Appends detection results as dictionaries to `jdict` list in-place. None: Appends detection results as dictionaries to `jdict` list in-place.
Example: Example:
```python ```python
predn = torch.tensor([[100, 50, 200, 150, 0.9, 0], [50, 30, 100, 80, 0.8, 1]]) predn = torch.tensor([[100, 50, 200, 150, 0.9, 0], [50, 30, 100, 80, 0.8, 1]])
jdict = [] jdict = []
path = Path("42.jpg") path = Path("42.jpg")
class_map = {0: 18, 1: 19} class_map = {0: 18, 1: 19}
save_one_json(predn, jdict, path, class_map) save_one_json(predn, jdict, path, class_map)
``` ```
This will append to `jdict`: This will append to `jdict`:
``` ```
[ [
{'image_id': 42, 'category_id': 18, 'bbox': [125.0, 75.0, 100.0, 100.0], 'score': 0.9}, {'image_id': 42, 'category_id': 18, 'bbox': [125.0, 75.0, 100.0, 100.0], 'score': 0.9},
{'image_id': 42, 'category_id': 19, 'bbox': [75.0, 55.0, 50.0, 50.0], 'score': 0.8} {'image_id': 42, 'category_id': 19, 'bbox': [75.0, 55.0, 50.0, 50.0], 'score': 0.8}
] ]
``` ```
Notes: Notes:
The `bbox` values are formatted as [x, y, width, height], where x and y represent the top-left corner of the box. The `bbox` values are formatted as [x, y, width, height], where x and y represent the top-left corner of the box.
""" """
image_id = int(path.stem) if path.stem.isnumeric() else path.stem image_id = int(path.stem) if path.stem.isnumeric() else path.stem
box = xyxy2xywh(predn[:, :4]) # xywh box = xyxy2xywh(predn[:, :4]) # xywh
@ -146,15 +146,15 @@ def process_batch(detections, labels, iouv):
Return a correct prediction matrix given detections and labels at various IoU thresholds. Return a correct prediction matrix given detections and labels at various IoU thresholds.
Args: Args:
detections (np.ndarray): Array of shape (N, 6) where each row corresponds to a detection with detections (np.ndarray): Array of shape (N, 6) where each row corresponds to a detection with format
format [x1, y1, x2, y2, conf, class]. [x1, y1, x2, y2, conf, class].
labels (np.ndarray): Array of shape (M, 5) where each row corresponds to a ground truth label with labels (np.ndarray): Array of shape (M, 5) where each row corresponds to a ground truth label with format
format [class, x1, y1, x2, y2]. [class, x1, y1, x2, y2].
iouv (np.ndarray): Array of IoU thresholds to evaluate at. iouv (np.ndarray): Array of IoU thresholds to evaluate at.
Returns: Returns:
correct (np.ndarray): A binary array of shape (N, len(iouv)) indicating whether each detection correct (np.ndarray): A binary array of shape (N, len(iouv)) indicating whether each detection is a true positive
is a true positive for each IoU threshold. There are 10 IoU levels used in the evaluation. for each IoU threshold. There are 10 IoU levels used in the evaluation.
Example: Example:
```python ```python
@ -220,7 +220,7 @@ def run(
Args: Args:
data (str | dict): Path to a dataset yaml file or a dataset dictionary. data (str | dict): Path to a dataset yaml file or a dataset dictionary.
weights (str | list[str], optional): Path to the model weights file(s). Supports various formats: PyTorch, weights (str | list[str], optional): Path to the model weights file(s). Supports various formats including PyTorch,
TorchScript, ONNX, OpenVINO, TensorRT, CoreML, TensorFlow SavedModel, TensorFlow GraphDef, TensorFlow Lite, TorchScript, ONNX, OpenVINO, TensorRT, CoreML, TensorFlow SavedModel, TensorFlow GraphDef, TensorFlow Lite,
TensorFlow Edge TPU, and PaddlePaddle. TensorFlow Edge TPU, and PaddlePaddle.
batch_size (int, optional): Batch size for inference. Default is 32. batch_size (int, optional): Batch size for inference. Default is 32.
@ -473,15 +473,14 @@ def parse_opt():
Args: Args:
data (str): Path to the dataset YAML file, default is 'data/coco128.yaml'. data (str): Path to the dataset YAML file, default is 'data/coco128.yaml'.
weights (List[str]): List of paths to the model weight files, default is 'yolov5s.pt'. weights (list[str]): List of paths to the model weight files, default is 'yolov5s.pt'.
batch_size (int): Batch size for inference, default is 32. batch_size (int): Batch size for inference, default is 32.
imgsz (int): Inference image size in pixels, default is 640. imgsz (int): Inference image size in pixels, default is 640.
conf_thres (float): Confidence threshold for predictions, default is 0.001. conf_thres (float): Confidence threshold for predictions, default is 0.001.
iou_thres (float): IoU threshold for Non-Max Suppression (NMS), default is 0.6. iou_thres (float): IoU threshold for Non-Max Suppression (NMS), default is 0.6.
max_det (int): Maximum number of detections per image, default is 300. max_det (int): Maximum number of detections per image, default is 300.
task (str): Task type - options are 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'. task (str): Task type - options are 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'.
device (str): Device to run the model on, e.g., '0' or '0,1,2,3' or 'cpu'. Default is empty to let device (str): Device to run the model on, e.g., '0' or '0,1,2,3' or 'cpu'. Default is empty to let the system choose automatically.
the system choose automatically.
workers (int): Maximum number of dataloader workers per rank in DDP mode, default is 8. workers (int): Maximum number of dataloader workers per rank in DDP mode, default is 8.
single_cls (bool): If set, treats the dataset as a single-class dataset. Default is False. single_cls (bool): If set, treats the dataset as a single-class dataset. Default is False.
augment (bool): If set, performs augmented inference. Default is False. augment (bool): If set, performs augmented inference. Default is False.
@ -552,10 +551,10 @@ def main(opt):
Args: Args:
opt (argparse.Namespace): Parsed command-line options. opt (argparse.Namespace): Parsed command-line options.
- This includes values for parameters like 'data', 'weights', 'batch_size', 'imgsz', 'conf_thres', 'iou_thres', This includes values for parameters like 'data', 'weights', 'batch_size', 'imgsz', 'conf_thres', 'iou_thres',
'max_det', 'task', 'device', 'workers', 'single_cls', 'augment', 'verbose', 'save_txt', 'save_hybrid', 'max_det', 'task', 'device', 'workers', 'single_cls', 'augment', 'verbose', 'save_txt', 'save_hybrid',
'save_conf', 'save_json', 'project', 'name', 'exist_ok', 'half', and 'dnn', essential for configuring 'save_conf', 'save_json', 'project', 'name', 'exist_ok', 'half', and 'dnn', essential for configuring
the YOLOv5 tasks. the YOLOv5 tasks.
Returns: Returns:
None None