Ultralytics Code Refactor https://ultralytics.com/actions (#13153)

* Refactor code for speed and clarity

* Update dataloaders.py

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

---------

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

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)

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)

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

detect.py

@@ -97,6 +97,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)

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

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)

models/tf.py

@@ -612,6 +612,7 @@ class TFModel:
         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

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)

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

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

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 @@ class LoadImagesAndLabels(Dataset):
         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))

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
 

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 @@ class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels):  # for training/testing
         rank=-1,
         seed=0,
     ):
+        """Initializes the dataset with image, label, and mask loading capabilities for training/testing."""
         super().__init__(
             path,
             img_size,

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