mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge branch 'master' into Maj-Pierre-1

pull/11869/head
Ultralytics Assistant 2024-07-16 01:21:04 +02:00 committed by GitHub
parent 0854767d8c 12be49963c
commit 1c5d4e67e8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 216 additions and 213 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
benchmarks.py
View File

@ -78,10 +78,11 @@ def run(
None. Logs information about the benchmark results, including the format, size, mAP50-95, and inference time.
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.
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.
Examples:
Example:
```python
$ python benchmarks.py --weights yolov5s.pt --img 640
```
@ -190,7 +191,6 @@ def test(
$ 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
"""
@ -263,7 +263,7 @@ def main(opt):
Args:
opt (argparse.Namespace): Parsed command-line arguments including options for weights, image size, batch size, data
configuration, device, and other flags for inference settings.
configuration, device, and other flags for inference settings.
Returns:
None: This function does not return any value. It leverages side-effects such as logging and running benchmarks.

31
detect.py
View File

@ -102,15 +102,15 @@ def run(
Args:
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.
Default is 'data/images'.
source (str | Path): Input source, which can be a file, directory, URL, glob pattern, screen capture, or webcam
index. Default is 'data/images'.
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).
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.
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
uses the best available device.
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.
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_csv (bool): If True, save results in a CSV file. Default is False.
@ -137,15 +137,15 @@ def run(
None
Examples:
```python
from ultralytics import run
```python
from ultralytics import run
# Run inference on an image
run(source='data/images/example.jpg', weights='yolov5s.pt', device='0')
# Run inference on an image
run(source='data/images/example.jpg', weights='yolov5s.pt', device='0')
# Run inference on a video with specific confidence threshold
run(source='data/videos/example.mp4', weights='yolov5s.pt', conf_thres=0.4, device='0')
```
# Run inference on a video with specific confidence threshold
run(source='data/videos/example.mp4', weights='yolov5s.pt', conf_thres=0.4, device='0')
```
"""
source = str(source)
save_img = not nosave and not source.endswith(".txt") # save inference images
@ -316,7 +316,7 @@ def run(
def parse_opt():
"""
Parses command-line arguments for YOLOv5 detection, setting inference options and model configurations.
Parse command-line arguments for YOLOv5 detection, allowing custom inference options and model configurations.
Args:
--weights (str | list[str], optional): Model path or Triton URL. Defaults to ROOT / 'yolov5s.pt'.
@ -346,7 +346,8 @@ def parse_opt():
--hide-conf (bool, optional): Flag to hide confidences in the output. Defaults to False.
--half (bool, optional): Flag to use FP16 half-precision inference. Defaults to False.
--dnn (bool, optional): Flag to use OpenCV DNN for ONNX inference. Defaults to False.
--vid-stride (int, optional): Video frame-rate stride, determining the number of frames to skip in between consecutive frames. Defaults to 1.
--vid-stride (int, optional): Video frame-rate stride, determining the number of frames to skip in between
consecutive frames. Defaults to 1.
Returns:
argparse.Namespace: Parsed command-line arguments as an argparse.Namespace object.
@ -403,8 +404,8 @@ def main(opt):
None
Note:
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:
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:
https://github.com/ultralytics/ultralytics
Example usage:

186
export.py
View File

@ -97,15 +97,15 @@ class iOSModel(torch.nn.Module):
Args:
model (torch.nn.Module): The PyTorch model to be adapted for iOS compatibility.
im (torch.Tensor): An input tensor representing a batch of images with shape (batch, channel, height, width).
im (torch.Tensor): An input tensor representing a batch of images with shape (B, C, H, W).
Returns:
None: This method does not return any value.
Notes:
This initializer configures normalization based on the input image dimensions, which is critical for ensuring the
model's compatibility and proper functionality on iOS devices. The normalization step involves dividing by the image
width if the image is square; otherwise, additional conditions might apply (trimmed for brevity).
This initializer configures normalization based on the input image dimensions, which is critical for
ensuring the model's compatibility and proper functionality on iOS devices. The normalization step
involves dividing by the image width if the image is square; otherwise, additional conditions might apply.
"""
super().__init__()
b, c, h, w = im.shape # batch, channel, height, width
@ -120,14 +120,15 @@ class iOSModel(torch.nn.Module):
def forward(self, x):
"""
Runs a forward pass on the input tensor, returning class confidences and normalized coordinates.
Run a forward pass on the input tensor, returning class confidences and normalized coordinates.
Args:
x (torch.Tensor): Input tensor containing the image data.
x (torch.Tensor): Input tensor containing the image data with shape (batch, channels, height, width).
Returns:
torch.Tensor: Concatenated tensor with normalized coordinates (xywh), confidence scores (conf),
and class probabilities (cls).
and class probabilities (cls), having shape (N, 4 + 1 + C), where N is the number of predictions,
and C is the number of classes.
Examples:
```python
@ -144,9 +145,9 @@ def export_formats():
Returns a DataFrame of supported YOLOv5 model export formats and their properties.
Returns:
pd.DataFrame: A pandas DataFrame containing supported export formats and their properties. The DataFrame includes
columns for format name, CLI argument suffix, file extension or directory name, and boolean flags indicating if the
export format supports training and detection.
pandas.DataFrame: A DataFrame containing supported export formats and their properties. The DataFrame
includes columns for format name, CLI argument suffix, file extension or directory name, and boolean flags
indicating if the export format supports training and detection.
Examples:
```python
@ -156,7 +157,7 @@ def export_formats():
Notes:
The DataFrame contains the following columns:
- Format: The name of the model format (e.g., PyTorch, TorchScript, ONNX, etc.)
- Format: The name of the model format (e.g., PyTorch, TorchScript, ONNX, etc.).
- Include Argument: The argument to use with the export script to include this format.
- File Suffix: File extension or directory name associated with the format.
- Supports Training: Whether the format supports training.
@ -181,7 +182,7 @@ def export_formats():
def try_export(inner_func):
"""
Logs success or failure, execution time, and file size for YOLOv5 model export functions wrapped with @try_export.
Log success or failure, execution time, and file size for YOLOv5 model export functions wrapped with @try_export.
Args:
inner_func (Callable): The model export function to be wrapped by the decorator.
@ -202,7 +203,8 @@ def try_export(inner_func):
```
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).
"""
inner_args = get_default_args(inner_func)
@ -224,7 +226,7 @@ def try_export(inner_func):
@try_export
def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")):
"""
Exports a YOLOv5 model to the TorchScript format.
Export a YOLOv5 model to the TorchScript format.
Args:
model (torch.nn.Module): The YOLOv5 model to be exported.
@ -235,8 +237,8 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:"
Returns:
(str | None, torch.jit.ScriptModule | None): A tuple containing the file path of the exported model
(as a string) and the TorchScript model (as a torch.jit.ScriptModule). If the export fails, both elements
of the tuple will be None.
(as a string) and the TorchScript model (as a torch.jit.ScriptModule). If the export fails, both elements
of the tuple will be None.
Notes:
- This function uses tracing to create the TorchScript model.
@ -303,6 +305,26 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX
```
pip install onnx onnx-simplifier onnxruntime onnxruntime-gpu
```
Example:
```python
from pathlib import Path
import torch
from models.experimental import attempt_load
from utils.torch_utils import select_device
# Load model
weights = 'yolov5s.pt'
device = select_device('')
model = attempt_load(weights, map_location=device)
# Example input tensor
im = torch.zeros(1, 3, 640, 640).to(device)
# Export model
file_path = Path('yolov5s.onnx')
export_onnx(model, im, file_path, opset=12, dynamic=True, simplify=True)
```
"""
check_requirements("onnx>=1.12.0")
import onnx
@ -361,10 +383,10 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX
@try_export
def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")):
"""
Exports a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization.
Export a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization.
Args:
file (Path): The path to the output file where the OpenVINO model will be saved.
file (Path): Path to the output file where the OpenVINO model will be saved.
metadata (dict): Dictionary including model metadata such as names and strides.
half (bool): If True, export the model with FP16 precision.
int8 (bool): If True, export the model with INT8 quantization.
@ -372,8 +394,8 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:
prefix (str): Prefix string for logging purposes (default is "OpenVINO:").
Returns:
(str, openvino.runtime.Model | None): Returns the OpenVINO model file path and openvino.runtime.Model object if
export is successful; otherwise, returns None.
(str, openvino.runtime.Model | None): The OpenVINO model file path and openvino.runtime.Model object if export is
successful; otherwise, None.
Notes:
- Requires `openvino-dev` package version 2023.0 or higher. Install with:
@ -451,7 +473,7 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:
@try_export
def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
"""
Export a YOLOv5 model to PaddlePaddle format using X2Paddle, saving the converted model and metadata.
Export a YOLOv5 PyTorch model to PaddlePaddle format using X2Paddle, saving the converted model and metadata.
Args:
model (torch.nn.Module): The YOLOv5 model to be exported.
@ -554,12 +576,12 @@ def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")):
@try_export
def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr("TensorRT:")):
"""
Exports a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0.
Export a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0.
Args:
model (torch.nn.Module): YOLOv5 model to be exported.
im (torch.Tensor): Input tensor of shape (B, C, H, W).
file (Path): Path to save the exported model.
file (pathlib.Path): Path to save the exported model.
half (bool): Set to True to export with FP16 precision.
dynamic (bool): Set to True to enable dynamic input shapes.
simplify (bool): Set to True to simplify the model during export.
@ -568,7 +590,7 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose
prefix (str): Log message prefix.
Returns:
(Path, None): Tuple containing the path to the exported model and None.
(pathlib.Path, None): Tuple containing the path to the exported model and None.
Raises:
AssertionError: If executed on CPU instead of GPU.
@ -666,7 +688,7 @@ def export_saved_model(
prefix=colorstr("TensorFlow SavedModel:"),
):
"""
Exports a YOLOv5 model to the TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression
Export a YOLOv5 model to the TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression
(NMS).
Args:
@ -791,7 +813,7 @@ def export_tflite(
):
# YOLOv5 TensorFlow Lite export
"""
Exports a YOLOv5 model to TensorFlow Lite format with optional INT8 quantization and NMS support.
Export a YOLOv5 model to TensorFlow Lite format with optional INT8 quantization and NMS support.
Args:
keras_model (tf.keras.Model): The Keras model to be exported.
@ -805,7 +827,8 @@ def export_tflite(
prefix (str): Prefix for log messages.
Returns:
(str | None, tflite.Model | None): The file path of the exported TFLite model and the TFLite model instance, or None if export failed.
(str | None, tflite.Model | None): The file path of the exported TFLite model and the TFLite model instance, or None
if the export failed.
Example:
```python
@ -863,8 +886,6 @@ def export_tflite(
@try_export
def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
"""
Clear and concise summary line describing the function's purpose:
Exports a YOLOv5 model to Edge TPU compatible TFLite format; requires Linux and Edge TPU compiler.
Args:
@ -927,28 +948,30 @@ def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
@try_export
def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
"""
Exports a YOLOv5 model to TensorFlow.js format, optionally with uint8 quantization.
Convert a YOLOv5 model to TensorFlow.js format with optional uint8 quantization.
Args:
file (Path): Path to the input model file.
int8 (bool): If True, applies uint8 quantization.
prefix (str): Prefix for logging information (default: colorstr("TensorFlow.js:")).
file (Path): Path to the YOLOv5 model file to be converted, typically having a ".pt" or ".onnx" extension.
int8 (bool): If True, applies uint8 quantization during the conversion process.
prefix (str): Optional prefix for logging messages, default is 'TensorFlow.js:' with color formatting.
Returns:
(str, None): The output directory path as a string and None.
(str, None): Tuple containing the output directory path as a string and None.
Notes:
This function requires `tensorflowjs` to be installed. You can install it using:
```shell
pip install tensorflowjs
```
- This function requires the `tensorflowjs` package. Install it using:
```shell
pip install tensorflowjs
```
- The converted TensorFlow.js model will be saved in a directory with the "_web_model" suffix appended to the original file name.
- The conversion involves running shell commands that invoke the TensorFlow.js converter tool.
Example usage:
Example:
```python
export_tfjs(Path('yolov5s.onnx'), int8=False)
from pathlib import Path
file = Path('yolov5.onnx')
export_tfjs(file, int8=False)
```
The TensorFlow.js converted model is saved in the directory specified by `file` with "_web_model" suffix.
"""
check_requirements("tensorflowjs")
import tensorflowjs as tfjs
@ -987,11 +1010,13 @@ def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
def add_tflite_metadata(file, metadata, num_outputs):
"""
Adds metadata to a TensorFlow Lite (TFLite) model file, supporting multiple outputs, based on TensorFlow guidelines.
Adds metadata to a TensorFlow Lite (TFLite) model file, supporting multiple outputs according to TensorFlow
guidelines.
Args:
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 the TFLite metadata schema.
Common keys include "name", "description", "version", "author", and "license".
num_outputs (int): Number of output tensors the model has, used to configure the metadata properly.
Returns:
@ -1047,20 +1072,20 @@ def add_tflite_metadata(file, metadata, num_outputs):
def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")):
"""
Converts a PyTorch YOLOv5 model to CoreML format with Non-Maximum Suppression (NMS), handling different input/output
shapes and saving the model.
Convert a PyTorch YOLOv5 model to CoreML format with Non-Maximum Suppression (NMS), handling different input/output
shapes, and saving the model.
Args:
model (torch.nn.Module): The YOLOv5 PyTorch model.
im (torch.Tensor): Input tensor example with shape [N, C, H, W], where N is the batch size, C is the number of
channels, H is the height, and W is the width.
model (torch.nn.Module): The YOLOv5 PyTorch model to be converted.
im (torch.Tensor): Example input tensor with shape (N, C, H, W), where N is the batch size, C is the number of channels,
H is the height, and W is the width.
file (Path): Path to save the converted CoreML model.
names (dict[int, str]): Dictionary mapping class indices to class names.
y (torch.Tensor): Output tensor from the PyTorch model's forward pass.
prefix (str): Custom prefix for logging messages.
Returns:
Path: Path to the saved CoreML model (.mlmodel).
(Path): Path to the saved CoreML model (.mlmodel).
Raises:
AssertionError: If the number of class names does not match the number of classes in the model.
@ -1075,18 +1100,14 @@ def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:
from pathlib import Path
import torch
# Load YOLOv5 model and an example input tensor
model = torch.load("yolov5s.pt")
im = torch.zeros(1, 3, 640, 640) # Example input tensor
model = torch.load('yolov5s.pt') # Load YOLOv5 model
im = torch.zeros((1, 3, 640, 640)) # Example input tensor
# Define class names
names = {0: "person", 1: "bicycle", 2: "car", ...}
names = {0: "person", 1: "bicycle", 2: "car", ...} # Define class names
# Perform forward pass to get model output
y = model(im)
y = model(im) # Perform forward pass to get model output
# Convert to CoreML
output_file = Path("yolov5s.mlmodel")
output_file = Path('yolov5s.mlmodel') # Convert to CoreML
pipeline_coreml(model, im, output_file, names, y)
```
"""
@ -1246,8 +1267,6 @@ def run(
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.
Args:
@ -1424,14 +1443,14 @@ def run(
def parse_opt(known=False):
"""
Parses command-line arguments for YOLOv5 model export configurations.
Parse command-line options for YOLOv5 model export configurations.
Args:
known (bool): If True, `argparse.ArgumentParser.parse_known_args` is used to parse command-line arguments; otherwise,
`argparse.ArgumentParser.parse_args` is used. Defaults to False.
known (bool): If True, uses `argparse.ArgumentParser.parse_known_args`; otherwise, uses `argparse.ArgumentParser.parse_args`.
Default is False.
Returns:
argparse.Namespace: An object containing parsed command-line arguments.
argparse.Namespace: Object containing parsed command-line arguments.
Example:
```python
@ -1475,44 +1494,7 @@ def parse_opt(known=False):
def main(opt):
"""
Exports the YOLOv5 model to specified formats, including ONNX, TensorRT, CoreML, and TensorFlow.
Args:
opt (argparse.Namespace): Parsed command-line arguments containing the export configurations.
- data (str): Path to the dataset YAML configuration file (e.g., 'data/coco128.yaml').
- weights (list[str] | str): Paths to the pretrained model weights file(s) (e.g., 'yolov5s.pt').
- imgsz (list[int]): Image size as a list [height, width].
- batch_size (int): Batch size for exporting the model.
- device (str): Device to run the export on, such as '0' for GPU, or 'cpu' for CPU.
- half (bool): Flag to export the model with FP16 half-precision.
- inplace (bool): Set the YOLOv5 Detect() module inplace mode to True.
- keras (bool): Flag to use Keras for TensorFlow SavedModel export.
- optimize (bool): Optimize TorchScript model for mobile deployment.
- int8 (bool): Apply INT8 quantization for CoreML or TensorFlow models.
- per_tensor (bool): Apply per-tensor quantization for TensorFlow models.
- dynamic (bool): Enable dynamic axes for ONNX, TensorFlow, or TensorRT exports.
- simplify (bool): Simplify ONNX model during export.
- opset (int): ONNX opset version.
- verbose (bool): Enable verbose logging for TensorRT export.
- workspace (int): TensorRT workspace size in GB.
- nms (bool): Add non-maximum suppression (NMS) to the TensorFlow model.
- agnostic_nms (bool): Add class-agnostic NMS to the TensorFlow model.
- topk_per_class (int): Top-K boxes per class to keep 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 NMS.
- conf_thres (float): Confidence threshold for NMS.
- include (list[str]): List of formats to include in export (e.g., ['torchscript', 'onnx']).
Returns:
list[str]: List of paths to the exported model files.
Example:
```python
opt = parse_opt()
main(opt)
```
"""
"""Run(**vars(opt)) # Execute the run function with parsed options."""
for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
run(**vars(opt))

102
hubconf.py
View File

@ -24,10 +24,11 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
classes (int, optional): Number of classes the model is expected to detect. Defaults to 80.
autoshape (bool, optional): If True, applies the YOLOv5 .autoshape() wrapper for various input formats. Defaults to True.
verbose (bool, optional): If True, prints detailed information during the model creation/loading process. Defaults to True.
device (str | torch.device | None, optional): Device to use for model parameters (e.g., 'cpu', 'cuda'). If None, selects the best available device. Defaults to None.
device (str | torch.device | None, optional): Device to use for model parameters (e.g., 'cpu', 'cuda'). If None, selects
the best available device. Defaults to None.
Returns:
DetectMultiBackend | AutoShape: The loaded YOLOv5 model, potentially wrapped with AutoShape if specified.
(DetectMultiBackend | AutoShape): The loaded YOLOv5 model, potentially wrapped with AutoShape if specified.
Examples:
```python
@ -108,8 +109,8 @@ def custom(path="path/to/model.pt", autoshape=True, _verbose=True, device=None):
Args:
path (str): Path to the custom model file (e.g., 'path/to/model.pt').
autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model if True, enabling compatibility with various input types
(default is True).
autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model if True, enabling compatibility with various input
types (default is True).
_verbose (bool): If True, prints all informational messages to the screen; otherwise, operates silently
(default is True).
device (str | torch.device | None): Device to load the model on, e.g., 'cpu', 'cuda', torch.device('cuda:0'), etc.
@ -174,32 +175,42 @@ def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
"""
Creates YOLOv5-small model with options for pretraining, input channels, class count, autoshaping, verbosity, and
device.
Create a YOLOv5-small (yolov5s) model with options for pretraining, input channels, class count, autoshaping,
verbosity, and device configuration.
Args:
pretrained (bool, optional): Flag to load pretrained weights into the model. Defaults to True.
channels (int, optional): Number of input channels. Defaults to 3.
classes (int, optional): Number of model classes. Defaults to 80.
autoshape (bool, optional): Whether to apply YOLOv5 .autoshape() wrapper to the model for preprocessed inputs.
autoshape (bool, optional): Whether to wrap the model with YOLOv5's .autoshape() for handling various input formats.
Defaults to True.
_verbose (bool, optional): Flag to print detailed information on model loading. Defaults to True.
device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda'. If None,
auto-select the best available device. Defaults to None.
_verbose (bool, optional): Flag to print detailed information regarding model loading. Defaults to True.
device (str | torch.device | None, optional): Device to use for model computation, can be 'cpu', 'cuda', or
torch.device instances. If None, automatically selects the best available device. Defaults to None.
Returns:
torch.nn.Module: The YOLOv5-small model loaded with specified configurations and optionally pretrained weights.
torch.nn.Module: The YOLOv5-small model configured and loaded according to the specified parameters.
Usage:
Example:
```python
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # official model
model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s') # from branch
model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt') # custom/local model
model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local') # local repo
# Load the official YOLOv5-small model with pretrained weights
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
# Load the YOLOv5-small model from a specific branch
model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s')
# Load a custom YOLOv5-small model from a local checkpoint
model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt')
# Load a local YOLOv5-small model specifying source as local repository
model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local')
```
For more information, visit https://pytorch.org/hub/ultralytics_yolov5.
Notes:
For more details on model loading and customization, visit
the [YOLOv5 PyTorch Hub Documentation](https://pytorch.org/hub/ultralytics_yolov5).
"""
return _create("yolov5s", pretrained, channels, classes, autoshape, _verbose, device)
@ -213,22 +224,26 @@ def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
pretrained (bool, optional): Whether to load pretrained weights into the model. Default is True.
channels (int, optional): Number of input channels. Default is 3.
classes (int, optional): Number of model classes. Default is 80.
autoshape (bool, optional): Apply YOLOv5 .autoshape() wrapper to the model for handling various input formats. Default is True.
autoshape (bool, optional): Apply YOLOv5 .autoshape() wrapper to the model for handling various input formats.
Default is True.
_verbose (bool, optional): Whether to print detailed information to the screen. Default is True.
device (str | torch.device | None, optional): Device specification to use for model parameters (e.g., 'cpu', 'cuda'). Default is None.
device (str | torch.device | None, optional): Device specification to use for model parameters (e.g., 'cpu', 'cuda').
Default is None.
Returns:
torch.nn.Module: The instantiated YOLOv5-medium model.
Usage Example:
```python
import torch
```python
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5m') # Load YOLOv5-medium from Ultralytics repository
model = torch.hub.load('ultralytics/yolov5:master', 'yolov5m') # Load from the master branch
model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5m.pt') # Load a custom/local YOLOv5-medium model
model = torch.hub.load('.', 'custom', 'yolov5m.pt', source='local') # Load from a local repository
```
model = torch.hub.load('ultralytics/yolov5', 'yolov5m') # Load YOLOv5-medium from Ultralytics repository
model = torch.hub.load('ultralytics/yolov5:master', 'yolov5m') # Load from the master branch
model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5m.pt') # Load a custom/local YOLOv5-medium model
model = torch.hub.load('.', 'custom', 'yolov5m.pt', source='local') # Load from a local repository
```
For more information, visit https://pytorch.org/hub/ultralytics_yolov5.
"""
return _create("yolov5m", pretrained, channels, classes, autoshape, _verbose, device)
@ -251,7 +266,7 @@ def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
YOLOv5 model (torch.nn.Module): The YOLOv5-large model instantiated with specified configurations and possibly
pretrained weights.
Example:
Examples:
```python
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5l')
@ -266,21 +281,22 @@ def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
"""
Instantiates the YOLOv5-xlarge model with customizable pretraining, channel count, class count, autoshaping,
verbosity, and device.
Perform object detection using the YOLOv5-xlarge model with options for pretraining, input channels, class count,
autoshaping, verbosity, and device specification.
Args:
pretrained (bool): If True, loads pretrained weights into the model. Defaults to True.
channels (int): Number of input channels. Defaults to 3.
classes (int): Number of model classes. Defaults to 80.
autoshape (bool): If True, applies YOLOv5 .autoshape() wrapper to the model for easier image handling. Defaults to
channels (int): Number of input channels for the model. Defaults to 3.
classes (int): Number of model classes for object detection. Defaults to 80.
autoshape (bool): If True, applies the YOLOv5 .autoshape() wrapper for handling different input formats. Defaults to
True.
_verbose (bool): If True, prints detailed information to the screen. Defaults to True.
device (str | torch.device | None): Device for model parameters, e.g., 'cpu', 'cuda:0', or a torch.device object.
_verbose (bool): If True, prints detailed information during model loading. Defaults to True.
device (str | torch.device | None): Device specification for computing the model, e.g., 'cpu', 'cuda:0', torch.device('cuda').
Defaults to None.
Returns:
torch.nn.Module: The instantiated YOLOv5-xlarge model.
torch.nn.Module: The YOLOv5-xlarge model loaded with the specified parameters, optionally with pretrained weights and
autoshaping applied.
Example:
```python
@ -288,7 +304,7 @@ def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=Tr
model = torch.hub.load('ultralytics/yolov5', 'yolov5x')
```
For more details and usage, refer to the official YOLOv5 PyTorch Hub models documentation:
For additional details, refer to the official YOLOv5 PyTorch Hub models documentation:
https://pytorch.org/hub/ultralytics_yolov5
"""
return _create("yolov5x", pretrained, channels, classes, autoshape, _verbose, device)
@ -308,7 +324,7 @@ def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
Default is None.
Returns:
torch.nn.Module: YOLOv5 model loaded with the specified configurations.
torch.nn.Module: YOLOv5-nano-P6 model loaded with the specified configurations.
Example:
```python
@ -362,7 +378,7 @@ def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
"""
Creates YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity, and
Create YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity, and
device.
Args:
@ -389,15 +405,15 @@ def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
```
Notes:
- The model can be loaded with pre-trained weights for better performance on specific tasks.
- The autoshape feature simplifies input handling by allowing various popular data formats.
- The model can be loaded with pre-trained weights for better performance on specific tasks.
- The autoshape feature simplifies input handling by allowing various popular data formats.
"""
return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device)
def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
"""
Instantiates the YOLOv5-large-P6 model with customizable pretraining, channel and class counts, autoshaping,
Instantiate the YOLOv5-large-P6 model with options for pretraining, channel and class counts, autoshaping,
verbosity, and device selection.
Args:
@ -407,12 +423,12 @@ def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=T
autoshape (bool, optional): If True, apply YOLOv5 .autoshape() wrapper to the model for input flexibility. 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 torch.device.
If None, automatically selects the best available device. Default is None.
If None, automatically selects the best available device. Default is None.
Returns:
torch.nn.Module: The instantiated YOLOv5-large-P6 model.
Usage:
Example:
```python
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5l6') # official model

35
train.py
View File

@ -102,7 +102,7 @@ GIT_INFO = check_git_info()
def train(hyp, opt, device, callbacks):
"""
Trains a YOLOv5 model on a custom dataset using specified hyperparameters, options, and device, managing datasets,
Train a YOLOv5 model on a custom dataset using specified hyperparameters, options, and device, managing datasets,
model architecture, loss computation, and optimizer steps.
Args:
@ -543,13 +543,13 @@ def train(hyp, opt, device, callbacks):
def parse_opt(known=False):
"""
Parses command-line arguments for YOLOv5 training, validation, and testing.
Parse command-line arguments for YOLOv5 training, validation, and testing.
Args:
known (bool, optional): If True, parses known arguments, ignoring the unknown. Defaults to False.
Returns:
argparse.Namespace: Parsed command-line arguments.
(argparse.Namespace): Parsed command-line arguments containing options for YOLOv5 execution.
Example:
```python
@ -559,9 +559,9 @@ def parse_opt(known=False):
```
Links:
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
- 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
"""
parser = argparse.ArgumentParser()
parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")
@ -618,7 +618,7 @@ def parse_opt(known=False):
def main(opt, callbacks=Callbacks()):
"""
Runs training or hyperparameter evolution with specified options and optional callbacks.
Runs the main entry point for training or hyperparameter evolution with specified options and optional callbacks.
Args:
opt (argparse.Namespace): The command-line arguments parsed for YOLOv5 training and evolution.
@ -629,7 +629,7 @@ def main(opt, callbacks=Callbacks()):
None
Note:
For detailed usage, visit:
For detailed usage, refer to:
https://github.com/ultralytics/yolov5/tree/master/models
"""
if RANK in {-1, 0}:
@ -888,23 +888,27 @@ def main(opt, callbacks=Callbacks()):
def generate_individual(input_ranges, individual_length):
"""
Generate a random individual with gene values within specified input ranges.
Generate an individual with random hyperparameters within specified ranges.
Args:
input_ranges (list[tuple[float, float]]): List of tuples where each tuple contains the lower and upper bounds
for the corresponding gene.
individual_length (int): The number of genes in the individual.
for the corresponding gene (hyperparameter).
individual_length (int): The number of genes (hyperparameters) in the individual.
Returns:
list[float]: A list representing a generated individual with random gene values within the specified ranges.
Examples:
Example:
```python
input_ranges = [(0.01, 0.1), (0.1, 1.0), (0.9, 2.0)]
individual_length = 3
individual = generate_individual(input_ranges, individual_length)
print(individual) # Output: [0.035, 0.678, 1.456] (example output)
```
Note:
The individual returned will have a length equal to `individual_length`, with each gene value being a floating-point
number within its specified range in `input_ranges`.
"""
individual = []
for i in range(individual_length):
@ -915,7 +919,7 @@ def generate_individual(input_ranges, individual_length):
def run(**kwargs):
"""
Executes YOLOv5 training with given options, allowing optional overrides through keyword arguments.
Execute YOLOv5 training with specified options, allowing optional overrides through keyword arguments.
Args:
weights (str, optional): Path to initial weights. Defaults to ROOT / 'yolov5s.pt'.
@ -931,8 +935,9 @@ def run(**kwargs):
noval (bool, optional): Only validate at the final epoch. Defaults to False.
noautoanchor (bool, optional): Disable AutoAnchor. Defaults to False.
noplots (bool, optional): Do not save plot files. Defaults to False.
evolve (int, optional): Evolve hyperparameters for a specified number of generations. Use 300 if provided without a value.
evolve_population (str, optional): Directory for loading population during evolution. Defaults to ROOT / 'data/hyps'.
evolve (int, optional): Evolve hyperparameters for a specified number of generations. Use 300 if provided without a
value.
evolve_population (str, optional): Directory for loading population during evolution. Defaults to ROOT / 'data/ hyps'.
resume_evolve (str, optional): Resume hyperparameter evolution from the last generation. Defaults to None.
bucket (str, optional): gsutil bucket for saving checkpoints. Defaults to an empty string.
cache (str, optional): Cache image data in 'ram' or 'disk'. Defaults to None.

65
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.
Args:
predn (torch.Tensor): Predicted bounding boxes and associated confidence scores and classes in xyxy format,
tensor of shape (N, 6) where N is the number of detections.
predn (torch.Tensor): Predicted bounding boxes and associated confidence scores and classes in xyxy format, 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.
shape (tuple): Shape of the original image as (height, width).
file (str | Path): File path where the result will be saved.
@ -219,7 +219,7 @@ def run(
Evaluates a YOLOv5 model on a dataset and logs performance metrics.
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 including PyTorch,
TorchScript, ONNX, OpenVINO, TensorRT, CoreML, TensorFlow SavedModel, TensorFlow GraphDef, TensorFlow Lite,
TensorFlow Edge TPU, and PaddlePaddle.
@ -469,34 +469,34 @@ def run(
def parse_opt():
"""
Parses command-line options for YOLOv5 model inference configuration.
Parse command-line options for configuring YOLOv5 model inference.
Args:
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'.
batch_size (int): Batch size for inference, default is 32.
imgsz (int): Inference image size in pixels, default is 640.
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.
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'.
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.
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.
augment (bool): If set, performs augmented inference. Default is False.
verbose (bool): If set, reports mAP by class. Default is False.
save_txt (bool): If set, saves results to *.txt files. Default is False.
save_hybrid (bool): If set, saves label+prediction hybrid results to *.txt files. Default is False.
save_conf (bool): If set, saves confidences in --save-txt labels. Default is False.
save_json (bool): If set, saves results to a COCO-JSON file. Default is False.
project (str): Project directory to save results to. Default is 'runs/val'.
name (str): Name of the directory to save results to. Default is 'exp'.
exist_ok (bool): If set, existing directory will not be incremented. Default is False.
half (bool): If set, uses FP16 half-precision inference. Default is False.
dnn (bool): If set, uses OpenCV DNN for ONNX inference. Default is False.
data (str, optional): Path to the dataset YAML file. Default is 'data/coco128.yaml'.
weights (list[str], optional): List of paths to model weight files. Default is 'yolov5s.pt'.
batch_size (int, optional): Batch size for inference. Default is 32.
imgsz (int, optional): Inference image size in pixels. Default is 640.
conf_thres (float, optional): Confidence threshold for predictions. Default is 0.001.
iou_thres (float, optional): IoU threshold for Non-Max Suppression (NMS). Default is 0.6.
max_det (int, optional): Maximum number of detections per image. Default is 300.
task (str, optional): Task type - options are 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'.
device (str, optional): 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.
workers (int, optional): Maximum number of dataloader workers per rank in DDP mode. Default is 8.
single_cls (bool, optional): If set, treats the dataset as a single-class dataset. Default is False.
augment (bool, optional): If set, performs augmented inference. Default is False.
verbose (bool, optional): If set, reports mAP by class. Default is False.
save_txt (bool, optional): If set, saves results to *.txt files. Default is False.
save_hybrid (bool, optional): If set, saves label+prediction hybrid results to *.txt files. Default is False.
save_conf (bool, optional): If set, saves confidences in --save-txt labels. Default is False.
save_json (bool, optional): If set, saves results to a COCO-JSON file. Default is False.
project (str, optional): Project directory to save results to. Default is 'runs/val'.
name (str, optional): Name of the directory to save results to. Default is 'exp'.
exist_ok (bool, optional): If set, existing directory will not be incremented. Default is False.
half (bool, optional): If set, uses FP16 half-precision inference. Default is False.
dnn (bool, optional): If set, uses OpenCV DNN for ONNX inference. Default is False.
Returns:
argparse.Namespace: Parsed command-line options
argparse.Namespace: Parsed command-line options.
Notes:
- The '--data' parameter is checked to ensure it ends with 'coco.yaml' if '--save-json' is set.
@ -508,7 +508,7 @@ def parse_opt():
```python
$ python val.py --weights yolov5s.pt --data coco128.yaml --img 640
```
Different model formats could be used instead of yolov5s.pt:
Different model formats could be used instead of `yolov5s.pt`:
```python
$ python val.py --weights yolov5s.pt yolov5s.torchscript yolov5s.onnx yolov5s_openvino_model yolov5s.engine
```
@ -551,17 +551,16 @@ def main(opt):
Args:
opt (argparse.Namespace): Parsed command-line options.
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',
'save_conf', 'save_json', 'project', 'name', 'exist_ok', 'half', and 'dnn', essential for configuring
the YOLOv5 tasks.
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', 'save_conf', 'save_json', 'project', 'name', 'exist_ok', 'half', and 'dnn', essential
for configuring the YOLOv5 tasks.
Returns:
None
Examples:
To validate a trained YOLOv5 model on the COCO dataset with a specific weights file, use:
```python
$ python val.py --weights yolov5s.pt --data coco128.yaml --img 640
```