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

* Refactor code for speed and clarity

* Auto-format by https://ultralytics.com/actions

* Update export.py

* Auto-format by https://ultralytics.com/actions

* Update dataloaders.py

---------

Co-authored-by: UltralyticsAssistant <web@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 YOLOv3 benchmarks on multiple export formats and validate performance metrics."""
     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 YOLOv3 export tests for various formats and log the results, including export success status."""
     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
 ):
+    """Performs YOLOv3 classification inference on various input sources and saves or displays 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,
 ):
+    """Evaluate a YOLOv3 classification model on the specified dataset, providing accuracy metrics."""
     # Initialize/load model and set device
     training = model is not None
     if training:  # called by train.py

detect.py

@@ -95,6 +95,7 @@ def run(
     dnn=False,  # use OpenCV DNN for ONNX inference
     vid_stride=1,  # video frame-rate stride
 ):
+    """Performs YOLOv3 detection on various input sources including images, videos, streams, and YouTube URLs."""
     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

@@ -415,6 +415,7 @@ def export_saved_model(
     keras=False,
     prefix=colorstr("TensorFlow SavedModel:"),
 ):
+    """Exports YOLOv3 model to TensorFlow SavedModel format; includes NMS and configuration options."""
     # YOLOv3 TensorFlow SavedModel export
     try:
         import tensorflow as tf
@@ -781,6 +782,7 @@ def run(
     iou_thres=0.45,  # TF.js NMS: IoU threshold
     conf_thres=0.25,  # TF.js NMS: confidence threshold
 ):
+    """Exports a PyTorch model to specified formats like ONNX, CoreML, TensorRT."""
     t = time.time()
     include = [x.lower() for x in include]  # to lowercase
     fmts = tuple(export_formats()["Argument"][1:])  # --include arguments

models/common.py

@@ -1025,6 +1025,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 YOLOv3 classification head with convolution, pooling and dropout layers for feature extraction
+        and classification.
+        """
         super().__init__()
         c_ = 1280  # efficientnet_b0 size
         self.conv = Conv(c1, c_, k, s, autopad(k, p), g)

models/tf.py

@@ -565,6 +565,7 @@ def predict(
         iou_thres=0.45,
         conf_thres=0.25,
     ):
+        """Performs inference on input data using a YOLOv3 model, including optional TensorFlow NMS."""
         y = []  # outputs
         x = inputs
         for m in self.model.layers:
@@ -683,6 +684,7 @@ def run(
     dynamic=False,  # dynamic batch size
 ):
     # PyTorch model
+    """Exports and summarizes both PyTorch and TensorFlow models for YOLOv5-based object detection."""
     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,
 ):
+    """Performs YOLOv3 segmentation inference on various sources such as images, videos, 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

@@ -178,6 +178,7 @@ def run(
     compute_loss=None,
     callbacks=Callbacks(),
 ):
+    """Validates a trained YOLOv3 segmentation model using a specified dataset and evaluation metrics."""
     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 a random perspective transformation to an image and its bounding boxes for data augmentation."""
     height = im.shape[0] + border[0] * 2  # shape(h,w,c)
     width = im.shape[1] + border[1] * 2
 
@@ -321,6 +322,7 @@ def classify_albumentations(
     auto_aug=False,
 ):
     # YOLOv3 classification Albumentations (optional, only used if package is installed)
+    """Generates an Albumentations transform pipeline for image classification with optional augmentations."""
     prefix = colorstr("albumentations: ")
     try:
         import albumentations as A

utils/dataloaders.py

@@ -142,6 +142,7 @@ def create_dataloader(
     shuffle=False,
     seed=0,
 ):
+    """Creates a DataLoader for training, with options for augmentation, caching, and parallelization."""
     if rect and shuffle:
         LOGGER.warning("WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False")
         shuffle = False
@@ -511,6 +512,7 @@ def __init__(
         min_items=0,
         prefix="",
     ):
+        """Initializes a dataset with images and labels for YOLOv3 training and validation."""
         self.img_size = img_size
         self.augment = augment
         self.hyp = hyp
@@ -1302,6 +1304,9 @@ 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 YOLOv3 Classifier
+    """Creates a DataLoader for image classification tasks with options for augmentation, caching, 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

@@ -27,6 +27,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 augmentation including rotation, translation, scale, and shear transformations."""
     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 images and labels with optional augmentations and distributed sampling."""
     if rect and shuffle:
         LOGGER.warning("WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False")
         shuffle = False
@@ -99,6 +100,7 @@ def __init__(
         downsample_ratio=1,
         overlap=False,
     ):
+        """Initializes image, label, and mask loading for training/testing with optional augmentations."""
         super().__init__(
             path,
             img_size,

val.py

@@ -145,6 +145,7 @@ def run(
     callbacks=Callbacks(),
     compute_loss=None,
 ):
+    """Validates a trained YOLO model on a dataset and saves detection results in specified formats."""
     # Initialize/load model and set device
     training = model is not None
     if training:  # called by train.py