diff --git a/benchmarks.py b/benchmarks.py index 4ca7122551e..c849eed6f07 100644 --- a/benchmarks.py +++ b/benchmarks.py @@ -60,6 +60,7 @@ def run( pt_only=False, # test PyTorch only hard_fail=False, # throw error on benchmark failure ): + """Run YOLOv5 benchmarks on multiple export formats and log results for model performance evaluation.""" y, t = [], time.time() device = select_device(device) model_type = type(attempt_load(weights, fuse=False)) # DetectionModel, SegmentationModel, etc. @@ -124,6 +125,7 @@ def test( pt_only=False, # test PyTorch only hard_fail=False, # throw error on benchmark failure ): + """Run YOLOv5 export tests for all supported formats and log the results, including inference speed and mAP.""" y, t = [], time.time() device = select_device(device) for i, (name, f, suffix, gpu) in export.export_formats().iterrows(): # index, (name, file, suffix, gpu-capable) diff --git a/classify/predict.py b/classify/predict.py index 4dc3735f3a7..33140e9b56c 100644 --- a/classify/predict.py +++ b/classify/predict.py @@ -84,6 +84,7 @@ def run( dnn=False, # use OpenCV DNN for ONNX inference vid_stride=1, # video frame-rate stride ): + """Conducts YOLOv5 classification inference on diverse input sources and saves results.""" source = str(source) save_img = not nosave and not source.endswith(".txt") # save inference images is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) diff --git a/classify/val.py b/classify/val.py index 23dbe7bfa42..8ce48f0645b 100644 --- a/classify/val.py +++ b/classify/val.py @@ -68,6 +68,7 @@ def run( criterion=None, pbar=None, ): + """Validates a YOLOv5 classification model on a dataset, computing metrics like top1 and top5 accuracy.""" # Initialize/load model and set device training = model is not None if training: # called by train.py diff --git a/detect.py b/detect.py index 38380185f6a..474a4570f05 100644 --- a/detect.py +++ b/detect.py @@ -98,6 +98,7 @@ def run( dnn=False, # use OpenCV DNN for ONNX inference vid_stride=1, # video frame-rate stride ): + """Runs YOLOv5 detection inference on various sources like images, videos, directories, streams, etc.""" source = str(source) save_img = not nosave and not source.endswith(".txt") # save inference images is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) diff --git a/export.py b/export.py index 7d79e8bf6a4..77c972ba27f 100644 --- a/export.py +++ b/export.py @@ -407,6 +407,9 @@ def export_saved_model( keras=False, prefix=colorstr("TensorFlow SavedModel:"), ): + """Exports a YOLOv5 model to TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression + (NMS). + """ # YOLOv5 TensorFlow SavedModel export try: import tensorflow as tf @@ -477,6 +480,7 @@ def export_tflite( keras_model, im, file, int8, per_tensor, data, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:") ): # YOLOv5 TensorFlow Lite export + """Exports YOLOv5 model to TensorFlow Lite format with optional FP16, INT8, and NMS support.""" import tensorflow as tf LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") @@ -784,6 +788,7 @@ def run( iou_thres=0.45, # TF.js NMS: IoU threshold conf_thres=0.25, # TF.js NMS: confidence threshold ): + """Exports YOLOv5 model to specified formats including ONNX, TensorRT, CoreML, and TensorFlow; see https://github.com/ultralytics/yolov5.""" t = time.time() include = [x.lower() for x in include] # to lowercase fmts = tuple(export_formats()["Argument"][1:]) # --include arguments diff --git a/models/common.py b/models/common.py index 781f999445d..049dfc0b9e0 100644 --- a/models/common.py +++ b/models/common.py @@ -1066,6 +1066,9 @@ class Classify(nn.Module): def __init__( self, c1, c2, k=1, s=1, p=None, g=1, dropout_p=0.0 ): # ch_in, ch_out, kernel, stride, padding, groups, dropout probability + """Initializes YOLOv5 classification head with convolution, pooling, and dropout layers for input to output + channel transformation. + """ super().__init__() c_ = 1280 # efficientnet_b0 size self.conv = Conv(c1, c_, k, s, autopad(k, p), g) diff --git a/models/tf.py b/models/tf.py index c65938c4b31..9884ec3db35 100644 --- a/models/tf.py +++ b/models/tf.py @@ -612,6 +612,7 @@ def predict( iou_thres=0.45, conf_thres=0.25, ): + """Runs inference on input data, with an option for TensorFlow NMS.""" y = [] # outputs x = inputs for m in self.model.layers: @@ -730,6 +731,7 @@ def run( dynamic=False, # dynamic batch size ): # PyTorch model + """Exports YOLOv5 model from PyTorch to TensorFlow and Keras formats, performing inference for validation.""" im = torch.zeros((batch_size, 3, *imgsz)) # BCHW image model = attempt_load(weights, device=torch.device("cpu"), inplace=True, fuse=False) _ = model(im) # inference diff --git a/segment/predict.py b/segment/predict.py index 109a68415b0..0bccaaaae9f 100644 --- a/segment/predict.py +++ b/segment/predict.py @@ -97,6 +97,7 @@ def run( vid_stride=1, # video frame-rate stride retina_masks=False, ): + """Run YOLOv5 segmentation inference on diverse sources including images, videos, directories, and streams.""" source = str(source) save_img = not nosave and not source.endswith(".txt") # save inference images is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) diff --git a/segment/val.py b/segment/val.py index b5e9f7557ec..ab8a66a90c3 100644 --- a/segment/val.py +++ b/segment/val.py @@ -184,6 +184,9 @@ def run( compute_loss=None, callbacks=Callbacks(), ): + """Validates a YOLOv5 segmentation model on specified dataset, producing metrics, plots, and optional JSON + output. + """ if save_json: check_requirements("pycocotools>=2.0.6") process = process_mask_native # more accurate diff --git a/utils/augmentations.py b/utils/augmentations.py index 1840d47d46c..4a6e441d7c4 100644 --- a/utils/augmentations.py +++ b/utils/augmentations.py @@ -157,6 +157,7 @@ def random_perspective( # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10)) # targets = [cls, xyxy] + """Applies random perspective transformation to an image, modifying the image and corresponding labels.""" height = im.shape[0] + border[0] * 2 # shape(h,w,c) width = im.shape[1] + border[1] * 2 @@ -336,6 +337,9 @@ def classify_albumentations( auto_aug=False, ): # YOLOv5 classification Albumentations (optional, only used if package is installed) + """Sets up and returns Albumentations transforms for YOLOv5 classification tasks depending on augmentation + settings. + """ prefix = colorstr("albumentations: ") try: import albumentations as A diff --git a/utils/dataloaders.py b/utils/dataloaders.py index dacb0e0b33d..21308f0cedb 100644 --- a/utils/dataloaders.py +++ b/utils/dataloaders.py @@ -174,6 +174,7 @@ def create_dataloader( shuffle=False, seed=0, ): + """Creates and returns a configured DataLoader instance for loading and processing image datasets.""" if rect and shuffle: LOGGER.warning("WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False") shuffle = False @@ -552,6 +553,7 @@ def __init__( rank=-1, seed=0, ): + """Initializes the YOLOv5 dataset loader, handling images and their labels, caching, and preprocessing.""" self.img_size = img_size self.augment = augment self.hyp = hyp @@ -1351,6 +1353,7 @@ def create_classification_dataloader( path, imgsz=224, batch_size=16, augment=True, cache=False, rank=-1, workers=8, shuffle=True ): # Returns Dataloader object to be used with YOLOv5 Classifier + """Creates a DataLoader for image classification, supporting caching, augmentation, and distributed training.""" with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP dataset = ClassificationDataset(root=path, imgsz=imgsz, augment=augment, cache=cache) batch_size = min(batch_size, len(dataset)) diff --git a/utils/segment/augmentations.py b/utils/segment/augmentations.py index 5773b56f4d7..d7dd8aec669 100644 --- a/utils/segment/augmentations.py +++ b/utils/segment/augmentations.py @@ -30,6 +30,7 @@ def random_perspective( # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10)) # targets = [cls, xyxy] + """Applies random perspective, rotation, scale, shear, and translation augmentations to an image and targets.""" height = im.shape[0] + border[0] * 2 # shape(h,w,c) width = im.shape[1] + border[1] * 2 diff --git a/utils/segment/dataloaders.py b/utils/segment/dataloaders.py index 0804818deca..c2be5f0dfe9 100644 --- a/utils/segment/dataloaders.py +++ b/utils/segment/dataloaders.py @@ -39,6 +39,7 @@ def create_dataloader( overlap_mask=False, seed=0, ): + """Creates a dataloader for training, validating, or testing YOLO models with various dataset options.""" if rect and shuffle: LOGGER.warning("WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False") shuffle = False @@ -102,6 +103,7 @@ def __init__( rank=-1, seed=0, ): + """Initializes the dataset with image, label, and mask loading capabilities for training/testing.""" super().__init__( path, img_size, diff --git a/val.py b/val.py index 221226b4a45..c1e8a6aa309 100644 --- a/val.py +++ b/val.py @@ -148,6 +148,7 @@ def run( callbacks=Callbacks(), compute_loss=None, ): + """Evaluates model on a dataset and logs performance metrics, results are saved to specific directories.""" # Initialize/load model and set device training = model is not None if training: # called by train.py