HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Fixed YOLO heads docstrings (#16822)

Browse source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Francesco Mattioli 2024-10-11 13:41:40 +02:00 • committed by GitHub
parent 75b3b695be
commit 8e92930a60
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 9 additions and 75 deletions
Download patch file Download diff file Expand all files Collapse all files

18
ultralytics/nn/modules/head.py
View file

@ -19,7 +19,7 @@ __all__ = "Detect", "Segment", "Pose", "Classify", "OBB", "RTDETRDecoder", "v10D
class Detect(nn.Module):
"""YOLOv8 Detect head for detection models."""
"""YOLO Detect head for detection models."""
dynamic = False # force grid reconstruction
export = False # export mode
@ -30,7 +30,7 @@ class Detect(nn.Module):
strides = torch.empty(0) # init
def __init__(self, nc=80, ch=()):
"""Initializes the YOLOv8 detection layer with specified number of classes and channels."""
"""Initializes the YOLO detection layer with specified number of classes and channels."""
super().__init__()
self.nc = nc # number of classes
self.nl = len(ch) # number of detection layers
@ -162,7 +162,7 @@ class Detect(nn.Module):
class Segment(Detect):
"""YOLOv8 Segment head for segmentation models."""
"""YOLO Segment head for segmentation models."""
def __init__(self, nc=80, nm=32, npr=256, ch=()):
"""Initialize the YOLO model attributes such as the number of masks, prototypes, and the convolution layers."""
@ -187,7 +187,7 @@ class Segment(Detect):
class OBB(Detect):
"""YOLOv8 OBB detection head for detection with rotation models."""
"""YOLO OBB detection head for detection with rotation models."""
def __init__(self, nc=80, ne=1, ch=()):
"""Initialize OBB with number of classes `nc` and layer channels `ch`."""
@ -217,7 +217,7 @@ class OBB(Detect):
class Pose(Detect):
"""YOLOv8 Pose head for keypoints models."""
"""YOLO Pose head for keypoints models."""
def __init__(self, nc=80, kpt_shape=(17, 3), ch=()):
"""Initialize YOLO network with default parameters and Convolutional Layers."""
@ -257,10 +257,10 @@ class Pose(Detect):
class Classify(nn.Module):
"""YOLOv8 classification head, i.e. x(b,c1,20,20) to x(b,c2)."""
"""YOLO classification head, i.e. x(b,c1,20,20) to x(b,c2)."""
def __init__(self, c1, c2, k=1, s=1, p=None, g=1):
"""Initializes YOLOv8 classification head to transform input tensor from (b,c1,20,20) to (b,c2) shape."""
"""Initializes YOLO classification head to transform input tensor from (b,c1,20,20) to (b,c2) shape."""
super().__init__()
c_ = 1280 # efficientnet_b0 size
self.conv = Conv(c1, c_, k, s, p, g)
@ -277,10 +277,10 @@ class Classify(nn.Module):
class WorldDetect(Detect):
"""Head for integrating YOLOv8 detection models with semantic understanding from text embeddings."""
"""Head for integrating YOLO detection models with semantic understanding from text embeddings."""
def __init__(self, nc=80, embed=512, with_bn=False, ch=()):
"""Initialize YOLOv8 detection layer with nc classes and layer channels ch."""
"""Initialize YOLO detection layer with nc classes and layer channels ch."""
super().__init__(nc, ch)
c3 = max(ch[0], min(self.nc, 100))
self.cv3 = nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, embed, 1)) for x in ch)