From c5ffbbf1c3f69c8e7acf2ec5106fb8edad33c3de Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 26 Aug 2024 15:46:32 +0200 Subject: [PATCH] Add https://www.reddit.com/r/Ultralytics/ badge (#13284) * Refactor code for speed and clarity * Update README.zh-CN.md Signed-off-by: Glenn Jocher --------- Signed-off-by: Glenn Jocher Co-authored-by: UltralyticsAssistant --- .github/workflows/merge-main-into-prs.yml | 4 +- README.md | 2 +- README.zh-CN.md | 1 + export.py | 2 + models/common.py | 78 +++++++++++++++-------- models/tf.py | 62 ++++++++++++------ models/yolo.py | 15 +++-- utils/__init__.py | 3 +- utils/activations.py | 8 +++ utils/augmentations.py | 12 ++-- utils/dataloaders.py | 11 +++- utils/general.py | 9 ++- utils/loggers/__init__.py | 3 +- utils/loss.py | 11 +++- utils/metrics.py | 3 +- utils/plots.py | 3 +- utils/segment/dataloaders.py | 2 + utils/segment/loss.py | 3 +- utils/segment/metrics.py | 2 + utils/torch_utils.py | 3 +- 20 files changed, 164 insertions(+), 73 deletions(-) diff --git a/.github/workflows/merge-main-into-prs.yml b/.github/workflows/merge-main-into-prs.yml index eae8e00..05a0363 100644 --- a/.github/workflows/merge-main-into-prs.yml +++ b/.github/workflows/merge-main-into-prs.yml @@ -34,11 +34,11 @@ jobs: g = Github(os.getenv('GITHUB_TOKEN')) repo = g.get_repo(os.getenv('GITHUB_REPOSITORY')) - + # Fetch the default branch name default_branch_name = repo.default_branch default_branch = repo.get_branch(default_branch_name) - + for pr in repo.get_pulls(state='open', sort='created'): try: # Get full names for repositories and branches diff --git a/README.md b/README.md index ac8b63a..8872959 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ YOLOv5 CI YOLOv5 Citation Docker Pulls - Discord Ultralytics Forums + Discord Ultralytics Forums Ultralytics Reddit
Run on Gradient Open In Colab diff --git a/README.zh-CN.md b/README.zh-CN.md index bb45872..530fed5 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -10,6 +10,7 @@ YOLOv5 CI YOLOv5 Citation Docker Pulls + Discord Ultralytics Forums Ultralytics Reddit
Run on Gradient Open In Colab diff --git a/export.py b/export.py index 5bc9238..2c697ec 100644 --- a/export.py +++ b/export.py @@ -91,6 +91,8 @@ MACOS = platform.system() == "Darwin" # macOS environment class iOSModel(torch.nn.Module): + """An iOS-compatible wrapper for YOLOv5 models that normalizes input images based on their dimensions.""" + def __init__(self, model, im): """ Initializes an iOS compatible model with normalization based on image dimensions. diff --git a/models/common.py b/models/common.py index 1e0ffdd..8ad53d5 100644 --- a/models/common.py +++ b/models/common.py @@ -71,7 +71,8 @@ def autopad(k, p=None, d=1): class Conv(nn.Module): - # Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation) + """Applies a convolution, batch normalization, and activation function to an input tensor in a neural network.""" + default_act = nn.SiLU() # default activation def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True): @@ -91,7 +92,8 @@ class Conv(nn.Module): class DWConv(Conv): - # Depth-wise convolution + """Implements a depth-wise convolution layer with optional activation for efficient spatial filtering.""" + def __init__(self, c1, c2, k=1, s=1, d=1, act=True): """Initializes a depth-wise convolution layer with optional activation; args: input channels (c1), output channels (c2), kernel size (k), stride (s), dilation (d), and activation flag (act). @@ -100,7 +102,8 @@ class DWConv(Conv): class DWConvTranspose2d(nn.ConvTranspose2d): - # Depth-wise transpose convolution + """A depth-wise transpose convolutional layer for upsampling in neural networks, particularly in YOLOv5 models.""" + def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0): """Initializes a depth-wise transpose convolutional layer for YOLOv5; args: input channels (c1), output channels (c2), kernel size (k), stride (s), input padding (p1), output padding (p2). @@ -109,7 +112,8 @@ class DWConvTranspose2d(nn.ConvTranspose2d): class TransformerLayer(nn.Module): - # Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance) + """Transformer layer with multihead attention and linear layers, optimized by removing LayerNorm.""" + def __init__(self, c, num_heads): """ Initializes a transformer layer, sans LayerNorm for performance, with multihead attention and linear layers. @@ -132,7 +136,8 @@ class TransformerLayer(nn.Module): class TransformerBlock(nn.Module): - # Vision Transformer https://arxiv.org/abs/2010.11929 + """A Transformer block for vision tasks with convolution, position embeddings, and Transformer layers.""" + def __init__(self, c1, c2, num_heads, num_layers): """Initializes a Transformer block for vision tasks, adapting dimensions if necessary and stacking specified layers. @@ -157,7 +162,8 @@ class TransformerBlock(nn.Module): class Bottleneck(nn.Module): - # Standard bottleneck + """A bottleneck layer with optional shortcut and group convolution for efficient feature extraction.""" + def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): """Initializes a standard bottleneck layer with optional shortcut and group convolution, supporting channel expansion. @@ -176,7 +182,8 @@ class Bottleneck(nn.Module): class BottleneckCSP(nn.Module): - # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks + """CSP bottleneck layer for feature extraction with cross-stage partial connections and optional shortcuts.""" + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): """Initializes CSP bottleneck with optional shortcuts; args: ch_in, ch_out, number of repeats, shortcut bool, groups, expansion. @@ -201,7 +208,8 @@ class BottleneckCSP(nn.Module): class CrossConv(nn.Module): - # Cross Convolution Downsample + """Implements a cross convolution layer with downsampling, expansion, and optional shortcut.""" + def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False): """ Initializes CrossConv with downsampling, expanding, and optionally shortcutting; `c1` input, `c2` output @@ -221,7 +229,8 @@ class CrossConv(nn.Module): class C3(nn.Module): - # CSP Bottleneck with 3 convolutions + """Implements a CSP Bottleneck module with three convolutions for enhanced feature extraction in neural networks.""" + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): """Initializes C3 module with options for channel count, bottleneck repetition, shortcut usage, group convolutions, and expansion. @@ -239,7 +248,8 @@ class C3(nn.Module): class C3x(C3): - # C3 module with cross-convolutions + """Extends the C3 module with cross-convolutions for enhanced feature extraction in neural networks.""" + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): """Initializes C3x module with cross-convolutions, extending C3 with customizable channel dimensions, groups, and expansion. @@ -250,7 +260,8 @@ class C3x(C3): class C3TR(C3): - # C3 module with TransformerBlock() + """C3 module with TransformerBlock for enhanced feature extraction in object detection models.""" + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): """Initializes C3 module with TransformerBlock for enhanced feature extraction, accepts channel sizes, shortcut config, group, and expansion. @@ -261,7 +272,8 @@ class C3TR(C3): class C3SPP(C3): - # C3 module with SPP() + """Extends the C3 module with an SPP layer for enhanced spatial feature extraction and customizable channels.""" + def __init__(self, c1, c2, k=(5, 9, 13), n=1, shortcut=True, g=1, e=0.5): """Initializes a C3 module with SPP layer for advanced spatial feature extraction, given channel sizes, kernel sizes, shortcut, group, and expansion ratio. @@ -272,7 +284,8 @@ class C3SPP(C3): class C3Ghost(C3): - # C3 module with GhostBottleneck() + """Implements a C3 module with Ghost Bottlenecks for efficient feature extraction in YOLOv5.""" + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): """Initializes YOLOv5's C3 module with Ghost Bottlenecks for efficient feature extraction.""" super().__init__(c1, c2, n, shortcut, g, e) @@ -281,7 +294,8 @@ class C3Ghost(C3): class SPP(nn.Module): - # Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729 + """Implements Spatial Pyramid Pooling (SPP) for feature extraction, ref: https://arxiv.org/abs/1406.4729.""" + def __init__(self, c1, c2, k=(5, 9, 13)): """Initializes SPP layer with Spatial Pyramid Pooling, ref: https://arxiv.org/abs/1406.4729, args: c1 (input channels), c2 (output channels), k (kernel sizes).""" super().__init__() @@ -301,7 +315,8 @@ class SPP(nn.Module): class SPPF(nn.Module): - # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher + """Implements a fast Spatial Pyramid Pooling (SPPF) layer for efficient feature extraction in YOLOv5 models.""" + def __init__(self, c1, c2, k=5): """ Initializes YOLOv5 SPPF layer with given channels and kernel size for YOLOv5 model, combining convolution and @@ -326,7 +341,8 @@ class SPPF(nn.Module): class Focus(nn.Module): - # Focus wh information into c-space + """Focuses spatial information into channel space using slicing and convolution for efficient feature extraction.""" + def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): """Initializes Focus module to concentrate width-height info into channel space with configurable convolution parameters. @@ -342,7 +358,8 @@ class Focus(nn.Module): class GhostConv(nn.Module): - # Ghost Convolution https://github.com/huawei-noah/ghostnet + """Implements Ghost Convolution for efficient feature extraction, see https://github.com/huawei-noah/ghostnet.""" + def __init__(self, c1, c2, k=1, s=1, g=1, act=True): """Initializes GhostConv with in/out channels, kernel size, stride, groups, and activation; halves out channels for efficiency. @@ -359,7 +376,8 @@ class GhostConv(nn.Module): class GhostBottleneck(nn.Module): - # Ghost Bottleneck https://github.com/huawei-noah/ghostnet + """Efficient bottleneck layer using Ghost Convolutions, see https://github.com/huawei-noah/ghostnet.""" + def __init__(self, c1, c2, k=3, s=1): """Initializes GhostBottleneck with ch_in `c1`, ch_out `c2`, kernel size `k`, stride `s`; see https://github.com/huawei-noah/ghostnet.""" super().__init__() @@ -379,7 +397,8 @@ class GhostBottleneck(nn.Module): class Contract(nn.Module): - # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40) + """Contracts spatial dimensions into channel dimensions for efficient processing in neural networks.""" + def __init__(self, gain=2): """Initializes a layer to contract spatial dimensions (width-height) into channels, e.g., input shape (1,64,80,80) to (1,256,40,40). @@ -399,7 +418,8 @@ class Contract(nn.Module): class Expand(nn.Module): - # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160) + """Expands spatial dimensions by redistributing channels, e.g., from (1,64,80,80) to (1,16,160,160).""" + def __init__(self, gain=2): """ Initializes the Expand module to increase spatial dimensions by redistributing channels, with an optional gain @@ -422,7 +442,8 @@ class Expand(nn.Module): class Concat(nn.Module): - # Concatenate a list of tensors along dimension + """Concatenates tensors along a specified dimension for efficient tensor manipulation in neural networks.""" + def __init__(self, dimension=1): """Initializes a Concat module to concatenate tensors along a specified dimension.""" super().__init__() @@ -436,7 +457,8 @@ class Concat(nn.Module): class DetectMultiBackend(nn.Module): - # YOLOv5 MultiBackend class for python inference on various backends + """YOLOv5 MultiBackend class for inference on various backends including PyTorch, ONNX, TensorRT, and more.""" + def __init__(self, weights="yolov5s.pt", device=torch.device("cpu"), dnn=False, data=None, fp16=False, fuse=True): """Initializes DetectMultiBackend with support for various inference backends, including PyTorch and ONNX.""" # PyTorch: weights = *.pt @@ -778,7 +800,8 @@ class DetectMultiBackend(nn.Module): class AutoShape(nn.Module): - # YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS + """AutoShape class for robust YOLOv5 inference with preprocessing, NMS, and support for various input formats.""" + conf = 0.25 # NMS confidence threshold iou = 0.45 # NMS IoU threshold agnostic = False # NMS class-agnostic @@ -889,7 +912,8 @@ class AutoShape(nn.Module): class Detections: - # YOLOv5 detections class for inference results + """Manages YOLOv5 detection results with methods for visualization, saving, cropping, and exporting detections.""" + def __init__(self, ims, pred, files, times=(0, 0, 0), names=None, shape=None): """Initializes the YOLOv5 Detections class with image info, predictions, filenames, timing and normalization.""" super().__init__() @@ -1047,7 +1071,8 @@ class Detections: class Proto(nn.Module): - # YOLOv5 mask Proto module for segmentation models + """YOLOv5 mask Proto module for segmentation models, performing convolutions and upsampling on input tensors.""" + def __init__(self, c1, c_=256, c2=32): """Initializes YOLOv5 Proto module for segmentation with input, proto, and mask channels configuration.""" super().__init__() @@ -1062,7 +1087,8 @@ class Proto(nn.Module): class Classify(nn.Module): - # YOLOv5 classification head, i.e. x(b,c1,20,20) to x(b,c2) + """YOLOv5 classification head with convolution, pooling, and dropout layers for channel transformation.""" + 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 diff --git a/models/tf.py b/models/tf.py index d79e486..59bb7e0 100644 --- a/models/tf.py +++ b/models/tf.py @@ -49,7 +49,8 @@ from utils.general import LOGGER, make_divisible, print_args class TFBN(keras.layers.Layer): - # TensorFlow BatchNormalization wrapper + """TensorFlow BatchNormalization wrapper for initializing with optional pretrained weights.""" + def __init__(self, w=None): """Initializes a TensorFlow BatchNormalization layer with optional pretrained weights.""" super().__init__() @@ -67,7 +68,8 @@ class TFBN(keras.layers.Layer): class TFPad(keras.layers.Layer): - # Pad inputs in spatial dimensions 1 and 2 + """Pads input tensors in spatial dimensions 1 and 2 with specified integer or tuple padding values.""" + def __init__(self, pad): """ Initializes a padding layer for spatial dimensions 1 and 2 with specified padding, supporting both int and tuple @@ -87,7 +89,8 @@ class TFPad(keras.layers.Layer): class TFConv(keras.layers.Layer): - # Standard convolution + """Implements a standard convolutional layer with optional batch normalization and activation for TensorFlow.""" + def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None): """ Initializes a standard convolution layer with optional batch normalization and activation; supports only @@ -118,7 +121,8 @@ class TFConv(keras.layers.Layer): class TFDWConv(keras.layers.Layer): - # Depthwise convolution + """Initializes a depthwise convolution layer with optional batch normalization and activation for TensorFlow.""" + def __init__(self, c1, c2, k=1, s=1, p=None, act=True, w=None): """ Initializes a depthwise convolution layer with optional batch normalization and activation for TensorFlow @@ -147,7 +151,8 @@ class TFDWConv(keras.layers.Layer): class TFDWConvTranspose2d(keras.layers.Layer): - # Depthwise ConvTranspose2d + """Implements a depthwise ConvTranspose2D layer for TensorFlow with specific settings.""" + def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0, w=None): """ Initializes depthwise ConvTranspose2D layer with specific channel, kernel, stride, and padding settings. @@ -179,7 +184,8 @@ class TFDWConvTranspose2d(keras.layers.Layer): class TFFocus(keras.layers.Layer): - # Focus wh information into c-space + """Focuses spatial information into channel space using pixel shuffling and convolution for TensorFlow models.""" + def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None): """ Initializes TFFocus layer to focus width and height information into channel space with custom convolution @@ -201,7 +207,8 @@ class TFFocus(keras.layers.Layer): class TFBottleneck(keras.layers.Layer): - # Standard bottleneck + """Implements a TensorFlow bottleneck layer with optional shortcut connections for efficient feature extraction.""" + def __init__(self, c1, c2, shortcut=True, g=1, e=0.5, w=None): """ Initializes a standard bottleneck layer for TensorFlow models, expanding and contracting channels with optional @@ -223,7 +230,8 @@ class TFBottleneck(keras.layers.Layer): class TFCrossConv(keras.layers.Layer): - # Cross Convolution + """Implements a cross convolutional layer with optional expansion, grouping, and shortcut for TensorFlow.""" + def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False, w=None): """Initializes cross convolution layer with optional expansion, grouping, and shortcut addition capabilities.""" super().__init__() @@ -238,7 +246,8 @@ class TFCrossConv(keras.layers.Layer): class TFConv2d(keras.layers.Layer): - # Substitution for PyTorch nn.Conv2D + """Implements a TensorFlow 2D convolution layer, mimicking PyTorch's nn.Conv2D for specified filters and stride.""" + def __init__(self, c1, c2, k, s=1, g=1, bias=True, w=None): """Initializes a TensorFlow 2D convolution layer, mimicking PyTorch's nn.Conv2D functionality for given filter sizes and stride. @@ -261,7 +270,8 @@ class TFConv2d(keras.layers.Layer): class TFBottleneckCSP(keras.layers.Layer): - # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks + """Implements a CSP bottleneck layer for TensorFlow models to enhance gradient flow and efficiency.""" + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None): """ Initializes CSP bottleneck layer with specified channel sizes, count, shortcut option, groups, and expansion @@ -289,7 +299,8 @@ class TFBottleneckCSP(keras.layers.Layer): class TFC3(keras.layers.Layer): - # CSP Bottleneck with 3 convolutions + """CSP bottleneck layer with 3 convolutions for TensorFlow, supporting optional shortcuts and group convolutions.""" + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None): """ Initializes CSP Bottleneck with 3 convolutions, supporting optional shortcuts and group convolutions. @@ -313,7 +324,8 @@ class TFC3(keras.layers.Layer): class TFC3x(keras.layers.Layer): - # 3 module with cross-convolutions + """A TensorFlow layer for enhanced feature extraction using cross-convolutions in object detection models.""" + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None): """ Initializes layer with cross-convolutions for enhanced feature extraction in object detection models. @@ -335,7 +347,8 @@ class TFC3x(keras.layers.Layer): class TFSPP(keras.layers.Layer): - # Spatial pyramid pooling layer used in YOLOv3-SPP + """Implements spatial pyramid pooling for YOLOv3-SPP with specific channels and kernel sizes.""" + def __init__(self, c1, c2, k=(5, 9, 13), w=None): """Initializes a YOLOv3-SPP layer with specific input/output channels and kernel sizes for pooling.""" super().__init__() @@ -351,7 +364,8 @@ class TFSPP(keras.layers.Layer): class TFSPPF(keras.layers.Layer): - # Spatial pyramid pooling-Fast layer + """Implements a fast spatial pyramid pooling layer for TensorFlow with optimized feature extraction.""" + def __init__(self, c1, c2, k=5, w=None): """Initializes a fast spatial pyramid pooling layer with customizable in/out channels, kernel size, and weights. @@ -373,7 +387,8 @@ class TFSPPF(keras.layers.Layer): class TFDetect(keras.layers.Layer): - # TF YOLOv5 Detect layer + """Implements YOLOv5 object detection layer in TensorFlow for predicting bounding boxes and class probabilities.""" + def __init__(self, nc=80, anchors=(), ch=(), imgsz=(640, 640), w=None): """Initializes YOLOv5 detection layer for TensorFlow with configurable classes, anchors, channels, and image size. @@ -427,7 +442,8 @@ class TFDetect(keras.layers.Layer): class TFSegment(TFDetect): - # YOLOv5 Segment head for segmentation models + """YOLOv5 segmentation head for TensorFlow, combining detection and segmentation.""" + def __init__(self, nc=80, anchors=(), nm=32, npr=256, ch=(), imgsz=(640, 640), w=None): """Initializes YOLOv5 Segment head with specified channel depths, anchors, and input size for segmentation models. @@ -450,6 +466,8 @@ class TFSegment(TFDetect): class TFProto(keras.layers.Layer): + """Implements convolutional and upsampling layers for feature extraction in YOLOv5 segmentation.""" + def __init__(self, c1, c_=256, c2=32, w=None): """Initializes TFProto layer with convolutional and upsampling layers for feature extraction and transformation. @@ -466,7 +484,8 @@ class TFProto(keras.layers.Layer): class TFUpsample(keras.layers.Layer): - # TF version of torch.nn.Upsample() + """Implements a TensorFlow upsampling layer with specified size, scale factor, and interpolation mode.""" + def __init__(self, size, scale_factor, mode, w=None): """ Initializes a TensorFlow upsampling layer with specified size, scale_factor, and mode, ensuring scale_factor is @@ -488,7 +507,8 @@ class TFUpsample(keras.layers.Layer): class TFConcat(keras.layers.Layer): - # TF version of torch.concat() + """Implements TensorFlow's version of torch.concat() for concatenating tensors along the last dimension.""" + def __init__(self, dimension=1, w=None): """Initializes a TensorFlow layer for NCHW to NHWC concatenation, requiring dimension=1.""" super().__init__() @@ -581,7 +601,8 @@ def parse_model(d, ch, model, imgsz): class TFModel: - # TF YOLOv5 model + """Implements YOLOv5 model in TensorFlow, supporting TensorFlow, Keras, and TFLite formats for object detection.""" + def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, model=None, imgsz=(640, 640)): """Initializes TF YOLOv5 model with specified configuration, channels, classes, model instance, and input size. @@ -653,7 +674,8 @@ class TFModel: class AgnosticNMS(keras.layers.Layer): - # TF Agnostic NMS + """Performs agnostic non-maximum suppression (NMS) on detected objects using IoU and confidence thresholds.""" + def call(self, input, topk_all, iou_thres, conf_thres): """Performs agnostic NMS on input tensors using given thresholds and top-K selection.""" return tf.map_fn( diff --git a/models/yolo.py b/models/yolo.py index d89c5da..c0dd946 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -70,7 +70,8 @@ except ImportError: class Detect(nn.Module): - # YOLOv5 Detect head for detection models + """YOLOv5 Detect head for processing input tensors and generating detection outputs in object detection models.""" + stride = None # strides computed during build dynamic = False # force grid reconstruction export = False # export mode @@ -127,7 +128,8 @@ class Detect(nn.Module): class Segment(Detect): - # YOLOv5 Segment head for segmentation models + """YOLOv5 Segment head for segmentation models, extending Detect with mask and prototype layers.""" + def __init__(self, nc=80, anchors=(), nm=32, npr=256, ch=(), inplace=True): """Initializes YOLOv5 Segment head with options for mask count, protos, and channel adjustments.""" super().__init__(nc, anchors, ch, inplace) @@ -214,7 +216,8 @@ class BaseModel(nn.Module): class DetectionModel(BaseModel): - # YOLOv5 detection model + """YOLOv5 detection model class for object detection tasks, supporting custom configurations and anchors.""" + def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, anchors=None): """Initializes YOLOv5 model with configuration file, input channels, number of classes, and custom anchors.""" super().__init__() @@ -332,14 +335,16 @@ Model = DetectionModel # retain YOLOv5 'Model' class for backwards compatibilit class SegmentationModel(DetectionModel): - # YOLOv5 segmentation model + """YOLOv5 segmentation model for object detection and segmentation tasks with configurable parameters.""" + def __init__(self, cfg="yolov5s-seg.yaml", ch=3, nc=None, anchors=None): """Initializes a YOLOv5 segmentation model with configurable params: cfg (str) for configuration, ch (int) for channels, nc (int) for num classes, anchors (list).""" super().__init__(cfg, ch, nc, anchors) class ClassificationModel(BaseModel): - # YOLOv5 classification model + """YOLOv5 classification model for image classification tasks, initialized with a config file or detection model.""" + def __init__(self, cfg=None, model=None, nc=1000, cutoff=10): """Initializes YOLOv5 model with config file `cfg`, input channels `ch`, number of classes `nc`, and `cuttoff` index. diff --git a/utils/__init__.py b/utils/__init__.py index c7ece49..185afd6 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -12,7 +12,8 @@ def emojis(str=""): class TryExcept(contextlib.ContextDecorator): - # YOLOv5 TryExcept class. Usage: @TryExcept() decorator or 'with TryExcept():' context manager + """A context manager and decorator for error handling that prints an optional message with emojis on exception.""" + def __init__(self, msg=""): """Initializes TryExcept with an optional message, used as a decorator or context manager for error handling.""" self.msg = msg diff --git a/utils/activations.py b/utils/activations.py index 47f0a99..44bec8c 100644 --- a/utils/activations.py +++ b/utils/activations.py @@ -7,6 +7,8 @@ import torch.nn.functional as F class SiLU(nn.Module): + """Applies the Sigmoid-weighted Linear Unit (SiLU) activation function, also known as Swish.""" + @staticmethod def forward(x): """ @@ -18,6 +20,8 @@ class SiLU(nn.Module): class Hardswish(nn.Module): + """Applies the Hardswish activation function, which is efficient for mobile and embedded devices.""" + @staticmethod def forward(x): """ @@ -38,7 +42,11 @@ class Mish(nn.Module): class MemoryEfficientMish(nn.Module): + """Efficiently applies the Mish activation function using custom autograd for reduced memory usage.""" + class F(torch.autograd.Function): + """Implements a custom autograd function for memory-efficient Mish activation.""" + @staticmethod def forward(ctx, x): """Applies the Mish activation function, a smooth ReLU alternative, to the input tensor `x`.""" diff --git a/utils/augmentations.py b/utils/augmentations.py index bdbe077..af4c405 100644 --- a/utils/augmentations.py +++ b/utils/augmentations.py @@ -18,7 +18,8 @@ IMAGENET_STD = 0.229, 0.224, 0.225 # RGB standard deviation class Albumentations: - # YOLOv5 Albumentations class (optional, only used if package is installed) + """Provides optional data augmentation for YOLOv5 using Albumentations library if installed.""" + def __init__(self, size=640): """Initializes Albumentations class for optional data augmentation in YOLOv5 with specified input size.""" self.transform = None @@ -378,7 +379,8 @@ def classify_transforms(size=224): class LetterBox: - # YOLOv5 LetterBox class for image preprocessing, i.e. T.Compose([LetterBox(size), ToTensor()]) + """Resizes and pads images to specified dimensions while maintaining aspect ratio for YOLOv5 preprocessing.""" + def __init__(self, size=(640, 640), auto=False, stride=32): """Initializes a LetterBox object for YOLOv5 image preprocessing with optional auto sizing and stride adjustment. @@ -405,7 +407,8 @@ class LetterBox: class CenterCrop: - # YOLOv5 CenterCrop class for image preprocessing, i.e. T.Compose([CenterCrop(size), ToTensor()]) + """Applies center crop to an image, resizing it to the specified size while maintaining aspect ratio.""" + def __init__(self, size=640): """Initializes CenterCrop for image preprocessing, accepting single int or tuple for size, defaults to 640.""" super().__init__() @@ -424,7 +427,8 @@ class CenterCrop: class ToTensor: - # YOLOv5 ToTensor class for image preprocessing, i.e. T.Compose([LetterBox(size), ToTensor()]) + """Converts BGR np.array image from HWC to RGB CHW format, normalizes to [0, 1], and supports FP16 if half=True.""" + def __init__(self, half=False): """Initializes ToTensor for YOLOv5 image preprocessing, with optional half precision (half=True for FP16).""" super().__init__() diff --git a/utils/dataloaders.py b/utils/dataloaders.py index 0593e52..61358eb 100644 --- a/utils/dataloaders.py +++ b/utils/dataloaders.py @@ -131,6 +131,8 @@ def seed_worker(worker_id): # Inherit from DistributedSampler and override iterator # https://github.com/pytorch/pytorch/blob/master/torch/utils/data/distributed.py class SmartDistributedSampler(distributed.DistributedSampler): + """A distributed sampler ensuring deterministic shuffling and balanced data distribution across GPUs.""" + def __iter__(self): """Yields indices for distributed data sampling, shuffled deterministically based on epoch and seed.""" g = torch.Generator() @@ -259,7 +261,8 @@ class _RepeatSampler: class LoadScreenshots: - # YOLOv5 screenshot dataloader, i.e. `python detect.py --source "screen 0 100 100 512 256"` + """Loads and processes screenshots for YOLOv5 detection from specified screen regions using mss.""" + def __init__(self, source, img_size=640, stride=32, auto=True, transforms=None): """ Initializes a screenshot dataloader for YOLOv5 with specified source region, image size, stride, auto, and @@ -428,7 +431,8 @@ class LoadImages: class LoadStreams: - # YOLOv5 streamloader, i.e. `python detect.py --source 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP streams` + """Loads and processes video streams for YOLOv5, supporting various sources including YouTube and IP cameras.""" + def __init__(self, sources="file.streams", img_size=640, stride=32, auto=True, transforms=None, vid_stride=1): """Initializes a stream loader for processing video streams with YOLOv5, supporting various sources including YouTube. @@ -531,7 +535,8 @@ def img2label_paths(img_paths): class LoadImagesAndLabels(Dataset): - # YOLOv5 train_loader/val_loader, loads images and labels for training and validation + """Loads images and their corresponding labels for training and validation in YOLOv5.""" + cache_version = 0.6 # dataset labels *.cache version rand_interp_methods = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4] diff --git a/utils/general.py b/utils/general.py index 57db68a..8c0b2fc 100644 --- a/utils/general.py +++ b/utils/general.py @@ -188,7 +188,8 @@ CONFIG_DIR = user_config_dir() # Ultralytics settings dir class Profile(contextlib.ContextDecorator): - # YOLOv5 Profile class. Usage: @Profile() decorator or 'with Profile():' context manager + """Context manager and decorator for profiling code execution time, with optional CUDA synchronization.""" + def __init__(self, t=0.0, device: torch.device = None): """Initializes a profiling context for YOLOv5 with optional timing threshold and device specification.""" self.t = t @@ -213,7 +214,8 @@ class Profile(contextlib.ContextDecorator): class Timeout(contextlib.ContextDecorator): - # YOLOv5 Timeout class. Usage: @Timeout(seconds) decorator or 'with Timeout(seconds):' context manager + """Enforces a timeout on code execution, raising TimeoutError if the specified duration is exceeded.""" + def __init__(self, seconds, *, timeout_msg="", suppress_timeout_errors=True): """Initializes a timeout context/decorator with defined seconds, optional message, and error suppression.""" self.seconds = int(seconds) @@ -239,7 +241,8 @@ class Timeout(contextlib.ContextDecorator): class WorkingDirectory(contextlib.ContextDecorator): - # Usage: @WorkingDirectory(dir) decorator or 'with WorkingDirectory(dir):' context manager + """Context manager/decorator to temporarily change the working directory within a 'with' statement or decorator.""" + def __init__(self, new_dir): """Initializes a context manager/decorator to temporarily change the working directory.""" self.dir = new_dir # new dir diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py index a37ae0f..92b4bcb 100644 --- a/utils/loggers/__init__.py +++ b/utils/loggers/__init__.py @@ -76,7 +76,8 @@ def _json_default(value): class Loggers: - # YOLOv5 Loggers class + """Initializes and manages various logging utilities for tracking YOLOv5 training and validation metrics.""" + def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, include=LOGGERS): """Initializes loggers for YOLOv5 training and validation metrics, paths, and options.""" self.save_dir = save_dir diff --git a/utils/loss.py b/utils/loss.py index 7dc2ad3..fd5912f 100644 --- a/utils/loss.py +++ b/utils/loss.py @@ -14,7 +14,8 @@ def smooth_BCE(eps=0.1): class BCEBlurWithLogitsLoss(nn.Module): - # BCEwithLogitLoss() with reduced missing label effects. + """Modified BCEWithLogitsLoss to reduce missing label effects in YOLOv5 training with optional alpha smoothing.""" + def __init__(self, alpha=0.05): """Initializes a modified BCEWithLogitsLoss with reduced missing label effects, taking optional alpha smoothing parameter. @@ -37,7 +38,8 @@ class BCEBlurWithLogitsLoss(nn.Module): class FocalLoss(nn.Module): - # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) + """Applies focal loss to address class imbalance by modifying BCEWithLogitsLoss with gamma and alpha parameters.""" + def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): """Initializes FocalLoss with specified loss function, gamma, and alpha values; modifies loss reduction to 'none'. @@ -71,7 +73,8 @@ class FocalLoss(nn.Module): class QFocalLoss(nn.Module): - # Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) + """Implements Quality Focal Loss to address class imbalance by modulating loss based on prediction confidence.""" + def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): """Initializes Quality Focal Loss with given loss function, gamma, alpha; modifies reduction to 'none'.""" super().__init__() @@ -101,6 +104,8 @@ class QFocalLoss(nn.Module): class ComputeLoss: + """Computes the total loss for YOLOv5 model predictions, including classification, box, and objectness losses.""" + sort_obj_iou = False # Compute losses diff --git a/utils/metrics.py b/utils/metrics.py index d7a5820..85c1a70 100644 --- a/utils/metrics.py +++ b/utils/metrics.py @@ -122,7 +122,8 @@ def compute_ap(recall, precision): class ConfusionMatrix: - # Updated version of https://github.com/kaanakan/object_detection_confusion_matrix + """Generates and visualizes a confusion matrix for evaluating object detection classification performance.""" + def __init__(self, nc, conf=0.25, iou_thres=0.45): """Initializes ConfusionMatrix with given number of classes, confidence, and IoU threshold.""" self.matrix = np.zeros((nc + 1, nc + 1)) diff --git a/utils/plots.py b/utils/plots.py index 89ed1bb..e899ea4 100644 --- a/utils/plots.py +++ b/utils/plots.py @@ -29,7 +29,8 @@ matplotlib.use("Agg") # for writing to files only class Colors: - # Ultralytics color palette https://ultralytics.com/ + """Provides an RGB color palette derived from Ultralytics color scheme for visualization tasks.""" + def __init__(self): """ Initializes the Colors class with a palette derived from Ultralytics color scheme, converting hex codes to RGB. diff --git a/utils/segment/dataloaders.py b/utils/segment/dataloaders.py index c2be5f0..5f5666c 100644 --- a/utils/segment/dataloaders.py +++ b/utils/segment/dataloaders.py @@ -83,6 +83,8 @@ def create_dataloader( class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels): # for training/testing + """Loads images, labels, and segmentation masks for training and testing YOLO models with augmentation support.""" + def __init__( self, path, diff --git a/utils/segment/loss.py b/utils/segment/loss.py index d4bc9d3..b3e76ae 100644 --- a/utils/segment/loss.py +++ b/utils/segment/loss.py @@ -12,7 +12,8 @@ from .general import crop_mask class ComputeLoss: - # Compute losses + """Computes the YOLOv5 model's loss components including classification, objectness, box, and mask losses.""" + def __init__(self, model, autobalance=False, overlap=False): """Initializes the compute loss function for YOLOv5 models with options for autobalancing and overlap handling. diff --git a/utils/segment/metrics.py b/utils/segment/metrics.py index 52ceb64..091b5b1 100644 --- a/utils/segment/metrics.py +++ b/utils/segment/metrics.py @@ -54,6 +54,8 @@ def ap_per_class_box_and_mask( class Metric: + """Computes performance metrics like precision, recall, F1 score, and average precision for model evaluation.""" + def __init__(self) -> None: """Initializes performance metric attributes for precision, recall, F1 score, average precision, and class indices. diff --git a/utils/torch_utils.py b/utils/torch_utils.py index af3d046..8bf6585 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -420,7 +420,8 @@ def smart_resume(ckpt, optimizer, ema=None, weights="yolov5s.pt", epochs=300, re class EarlyStopping: - # YOLOv5 simple early stopper + """Implements early stopping to halt training when no improvement is observed for a specified number of epochs.""" + def __init__(self, patience=30): """Initializes simple early stopping mechanism for YOLOv5, with adjustable patience for non-improving epochs.""" self.best_fitness = 0.0 # i.e. mAP