Skip to content

Commit

Permalink
Ultralytics Code Refactor https://ultralytics.com/actions (#2236)
Browse files Browse the repository at this point in the history
* 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 <[email protected]>
  • Loading branch information
glenn-jocher and UltralyticsAssistant authored Jun 30, 2024
1 parent b93ce58 commit 8ed1e65
Show file tree
Hide file tree
Showing 14 changed files with 25 additions and 0 deletions.
2 changes: 2 additions & 0 deletions benchmarks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
1 change: 1 addition & 0 deletions classify/predict.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
1 change: 1 addition & 0 deletions classify/val.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions detect.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions export.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions models/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions models/tf.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions segment/predict.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
1 change: 1 addition & 0 deletions segment/val.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions utils/augmentations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
5 changes: 5 additions & 0 deletions utils/dataloaders.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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))
Expand Down
1 change: 1 addition & 0 deletions utils/segment/augmentations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 2 additions & 0 deletions utils/segment/dataloaders.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions val.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 8ed1e65

Please sign in to comment.