ultralytics 8.3.89__py3-none-any.whl → 8.3.91__py3-none-any.whl

ultralytics/nn/tasks.py

@@ -119,10 +119,10 @@ class BaseModel(torch.nn.Module):
 
         Args:
             x (torch.Tensor): The input tensor to the model.
-            profile (bool):  Print the computation time of each layer if True, defaults to False.
-            visualize (bool): Save the feature maps of the model if True, defaults to False.
-            augment (bool): Augment image during prediction, defaults to False.
-            embed (list, optional): A list of feature vectors/embeddings to return.
+            profile (bool): Print the computation time of each layer if True.
+            visualize (bool): Save the feature maps of the model if True.
+            augment (bool): Augment image during prediction.
+            embed (List, optional): A list of feature vectors/embeddings to return.
 
         Returns:
             (torch.Tensor): The last output of the model.
@@ -137,9 +137,9 @@ class BaseModel(torch.nn.Module):
 
         Args:
             x (torch.Tensor): The input tensor to the model.
-            profile (bool):  Print the computation time of each layer if True, defaults to False.
-            visualize (bool): Save the feature maps of the model if True, defaults to False.
-            embed (list, optional): A list of feature vectors/embeddings to return.
+            profile (bool): Print the computation time of each layer if True.
+            visualize (bool): Save the feature maps of the model if True.
+            embed (List, optional): A list of feature vectors/embeddings to return.
 
         Returns:
             (torch.Tensor): The last output of the model.
@@ -170,13 +170,12 @@ class BaseModel(torch.nn.Module):
 
     def _profile_one_layer(self, m, x, dt):
         """
-        Profile the computation time and FLOPs of a single layer of the model on a given input. Appends the results to
-        the provided list.
+        Profile the computation time and FLOPs of a single layer of the model on a given input.
 
         Args:
             m (torch.nn.Module): The layer to be profiled.
             x (torch.Tensor): The input data to the layer.
-            dt (list): A list to store the computation time of the layer.
+            dt (List): A list to store the computation time of the layer.
         """
         c = m == self.model[-1] and isinstance(x, list)  # is final layer list, copy input as inplace fix
         flops = thop.profile(m, inputs=[x.copy() if c else x], verbose=False)[0] / 1e9 * 2 if thop else 0  # GFLOPs
@@ -192,8 +191,8 @@ class BaseModel(torch.nn.Module):
 
     def fuse(self, verbose=True):
         """
-        Fuse the `Conv2d()` and `BatchNorm2d()` layers of the model into a single layer, in order to improve the
-        computation efficiency.
+        Fuse the `Conv2d()` and `BatchNorm2d()` layers of the model into a single layer for improved computation
+        efficiency.
 
         Returns:
             (torch.nn.Module): The fused model is returned.
@@ -225,7 +224,7 @@ class BaseModel(torch.nn.Module):
         Check if the model has less than a certain threshold of BatchNorm layers.
 
         Args:
-            thresh (int, optional): The threshold number of BatchNorm layers. Default is 10.
+            thresh (int, optional): The threshold number of BatchNorm layers.
 
         Returns:
             (bool): True if the number of BatchNorm layers in the model is less than the threshold, False otherwise.
@@ -235,21 +234,21 @@ class BaseModel(torch.nn.Module):
 
     def info(self, detailed=False, verbose=True, imgsz=640):
         """
-        Prints model information.
+        Print model information.
 
         Args:
-            detailed (bool): if True, prints out detailed information about the model. Defaults to False
-            verbose (bool): if True, prints out the model information. Defaults to False
-            imgsz (int): the size of the image that the model will be trained on. Defaults to 640
+            detailed (bool): If True, prints out detailed information about the model.
+            verbose (bool): If True, prints out the model information.
+            imgsz (int): The size of the image that the model will be trained on.
         """
         return model_info(self, detailed=detailed, verbose=verbose, imgsz=imgsz)
 
     def _apply(self, fn):
         """
-        Applies a function to all the tensors in the model that are not parameters or registered buffers.
+        Apply a function to all tensors in the model that are not parameters or registered buffers.
 
         Args:
-            fn (function): the function to apply to the model
+            fn (function): The function to apply to the model.
 
         Returns:
             (BaseModel): An updated BaseModel object.
@@ -264,11 +263,11 @@ class BaseModel(torch.nn.Module):
 
     def load(self, weights, verbose=True):
         """
-        Load the weights into the model.
+        Load weights into the model.
 
         Args:
             weights (dict | torch.nn.Module): The pre-trained weights to be loaded.
-            verbose (bool, optional): Whether to log the transfer progress. Defaults to True.
+            verbose (bool, optional): Whether to log the transfer progress.
         """
         model = weights["model"] if isinstance(weights, dict) else weights  # torchvision models are not dicts
         csd = model.float().state_dict()  # checkpoint state_dict as FP32
@@ -282,8 +281,8 @@ class BaseModel(torch.nn.Module):
         Compute loss.
 
         Args:
-            batch (dict): Batch to compute loss on
-            preds (torch.Tensor | List[torch.Tensor]): Predictions.
+            batch (dict): Batch to compute loss on.
+            preds (torch.Tensor | List[torch.Tensor], optional): Predictions.
         """
         if getattr(self, "criterion", None) is None:
             self.criterion = self.init_criterion()
@@ -300,7 +299,15 @@ class DetectionModel(BaseModel):
     """YOLO detection model."""
 
     def __init__(self, cfg="yolo11n.yaml", ch=3, nc=None, verbose=True):  # model, input channels, number of classes
-        """Initialize the YOLO detection model with the given config and parameters."""
+        """
+        Initialize the YOLO detection model with the given config and parameters.
+
+        Args:
+            cfg (str | dict): Model configuration file path or dictionary.
+            ch (int): Number of input channels.
+            nc (int, optional): Number of classes.
+            verbose (bool): Whether to display model information.
+        """
         super().__init__()
         self.yaml = cfg if isinstance(cfg, dict) else yaml_model_load(cfg)  # cfg dict
         if self.yaml["backbone"][0][2] == "Silence":
@@ -327,7 +334,7 @@ class DetectionModel(BaseModel):
             m.inplace = self.inplace
 
             def _forward(x):
-                """Performs a forward pass through the model, handling different Detect subclass types accordingly."""
+                """Perform a forward pass through the model, handling different Detect subclass types accordingly."""
                 if self.end2end:
                     return self.forward(x)["one2many"]
                 return self.forward(x)[0] if isinstance(m, (Segment, Pose, OBB)) else self.forward(x)
@@ -345,7 +352,15 @@ class DetectionModel(BaseModel):
             LOGGER.info("")
 
     def _predict_augment(self, x):
-        """Perform augmentations on input image x and return augmented inference and train outputs."""
+        """
+        Perform augmentations on input image x and return augmented inference and train outputs.
+
+        Args:
+            x (torch.Tensor): Input image tensor.
+
+        Returns:
+            (torch.Tensor): Augmented inference output.
+        """
         if getattr(self, "end2end", False) or self.__class__.__name__ != "DetectionModel":
             LOGGER.warning("WARNING ⚠️ Model does not support 'augment=True', reverting to single-scale prediction.")
             return self._predict_once(x)
@@ -363,7 +378,19 @@ class DetectionModel(BaseModel):
 
     @staticmethod
     def _descale_pred(p, flips, scale, img_size, dim=1):
-        """De-scale predictions following augmented inference (inverse operation)."""
+        """
+        De-scale predictions following augmented inference (inverse operation).
+
+        Args:
+            p (torch.Tensor): Predictions tensor.
+            flips (int): Flip type (0=none, 2=ud, 3=lr).
+            scale (float): Scale factor.
+            img_size (tuple): Original image size (height, width).
+            dim (int): Dimension to split at.
+
+        Returns:
+            (torch.Tensor): De-scaled predictions.
+        """
         p[:, :4] /= scale  # de-scale
         x, y, wh, cls = p.split((1, 1, 2, p.shape[dim] - 4), dim)
         if flips == 2:
@@ -373,7 +400,15 @@ class DetectionModel(BaseModel):
         return torch.cat((x, y, wh, cls), dim)
 
     def _clip_augmented(self, y):
-        """Clip YOLO augmented inference tails."""
+        """
+        Clip YOLO augmented inference tails.
+
+        Args:
+            y (List[torch.Tensor]): List of detection tensors.
+
+        Returns:
+            (List[torch.Tensor]): Clipped detection tensors.
+        """
         nl = self.model[-1].nl  # number of detection layers (P3-P5)
         g = sum(4**x for x in range(nl))  # grid points
         e = 1  # exclude layer count
@@ -392,7 +427,15 @@ class OBBModel(DetectionModel):
     """YOLO Oriented Bounding Box (OBB) model."""
 
     def __init__(self, cfg="yolo11n-obb.yaml", ch=3, nc=None, verbose=True):
-        """Initialize YOLO OBB model with given config and parameters."""
+        """
+        Initialize YOLO OBB model with given config and parameters.
+
+        Args:
+            cfg (str | dict): Model configuration file path or dictionary.
+            ch (int): Number of input channels.
+            nc (int, optional): Number of classes.
+            verbose (bool): Whether to display model information.
+        """
         super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
 
     def init_criterion(self):
@@ -404,7 +447,15 @@ class SegmentationModel(DetectionModel):
     """YOLO segmentation model."""
 
     def __init__(self, cfg="yolo11n-seg.yaml", ch=3, nc=None, verbose=True):
-        """Initialize YOLOv8 segmentation model with given config and parameters."""
+        """
+        Initialize YOLOv8 segmentation model with given config and parameters.
+
+        Args:
+            cfg (str | dict): Model configuration file path or dictionary.
+            ch (int): Number of input channels.
+            nc (int, optional): Number of classes.
+            verbose (bool): Whether to display model information.
+        """
         super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
 
     def init_criterion(self):
@@ -416,7 +467,16 @@ class PoseModel(DetectionModel):
     """YOLO pose model."""
 
     def __init__(self, cfg="yolo11n-pose.yaml", ch=3, nc=None, data_kpt_shape=(None, None), verbose=True):
-        """Initialize YOLOv8 Pose model."""
+        """
+        Initialize YOLOv8 Pose model.
+
+        Args:
+            cfg (str | dict): Model configuration file path or dictionary.
+            ch (int): Number of input channels.
+            nc (int, optional): Number of classes.
+            data_kpt_shape (tuple): Shape of keypoints data.
+            verbose (bool): Whether to display model information.
+        """
         if not isinstance(cfg, dict):
             cfg = yaml_model_load(cfg)  # load model YAML
         if any(data_kpt_shape) and list(data_kpt_shape) != list(cfg["kpt_shape"]):
@@ -433,12 +493,28 @@ class ClassificationModel(BaseModel):
     """YOLO classification model."""
 
     def __init__(self, cfg="yolo11n-cls.yaml", ch=3, nc=None, verbose=True):
-        """Init ClassificationModel with YAML, channels, number of classes, verbose flag."""
+        """
+        Initialize ClassificationModel with YAML, channels, number of classes, verbose flag.
+
+        Args:
+            cfg (str | dict): Model configuration file path or dictionary.
+            ch (int): Number of input channels.
+            nc (int, optional): Number of classes.
+            verbose (bool): Whether to display model information.
+        """
         super().__init__()
         self._from_yaml(cfg, ch, nc, verbose)
 
     def _from_yaml(self, cfg, ch, nc, verbose):
-        """Set YOLOv8 model configurations and define the model architecture."""
+        """
+        Set YOLOv8 model configurations and define the model architecture.
+
+        Args:
+            cfg (str | dict): Model configuration file path or dictionary.
+            ch (int): Number of input channels.
+            nc (int, optional): Number of classes.
+            verbose (bool): Whether to display model information.
+        """
         self.yaml = cfg if isinstance(cfg, dict) else yaml_model_load(cfg)  # cfg dict
 
         # Define model
@@ -455,7 +531,13 @@ class ClassificationModel(BaseModel):
 
     @staticmethod
     def reshape_outputs(model, nc):
-        """Update a TorchVision classification model to class count 'n' if required."""
+        """
+        Update a TorchVision classification model to class count 'n' if required.
+
+        Args:
+            model (torch.nn.Module): Model to update.
+            nc (int): New number of classes.
+        """
         name, m = list((model.model if hasattr(model, "model") else model).named_children())[-1]  # last module
         if isinstance(m, Classify):  # YOLO Classify() head
             if m.linear.out_features != nc:
@@ -500,10 +582,10 @@ class RTDETRDetectionModel(DetectionModel):
         Initialize the RTDETRDetectionModel.
 
         Args:
-            cfg (str): Configuration file name or path.
+            cfg (str | dict): Configuration file name or path.
             ch (int): Number of input channels.
-            nc (int, optional): Number of classes. Defaults to None.
-            verbose (bool, optional): Print additional information during initialization. Defaults to True.
+            nc (int, optional): Number of classes.
+            verbose (bool): Print additional information during initialization.
         """
         super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
 
@@ -519,7 +601,7 @@ class RTDETRDetectionModel(DetectionModel):
 
         Args:
             batch (dict): Dictionary containing image and label data.
-            preds (torch.Tensor, optional): Precomputed model predictions. Defaults to None.
+            preds (torch.Tensor, optional): Precomputed model predictions.
 
         Returns:
             (tuple): A tuple containing the total loss and main three losses in a tensor.
@@ -564,11 +646,11 @@ class RTDETRDetectionModel(DetectionModel):
 
         Args:
             x (torch.Tensor): The input tensor.
-            profile (bool, optional): If True, profile the computation time for each layer. Defaults to False.
-            visualize (bool, optional): If True, save feature maps for visualization. Defaults to False.
-            batch (dict, optional): Ground truth data for evaluation. Defaults to None.
-            augment (bool, optional): If True, perform data augmentation during inference. Defaults to False.
-            embed (list, optional): A list of feature vectors/embeddings to return.
+            profile (bool): If True, profile the computation time for each layer.
+            visualize (bool): If True, save feature maps for visualization.
+            batch (dict, optional): Ground truth data for evaluation.
+            augment (bool): If True, perform data augmentation during inference.
+            embed (List, optional): A list of feature vectors/embeddings to return.
 
         Returns:
             (torch.Tensor): Model's output tensor.
@@ -596,13 +678,28 @@ class WorldModel(DetectionModel):
     """YOLOv8 World Model."""
 
     def __init__(self, cfg="yolov8s-world.yaml", ch=3, nc=None, verbose=True):
-        """Initialize YOLOv8 world model with given config and parameters."""
+        """
+        Initialize YOLOv8 world model with given config and parameters.
+
+        Args:
+            cfg (str | dict): Model configuration file path or dictionary.
+            ch (int): Number of input channels.
+            nc (int, optional): Number of classes.
+            verbose (bool): Whether to display model information.
+        """
         self.txt_feats = torch.randn(1, nc or 80, 512)  # features placeholder
         self.clip_model = None  # CLIP model placeholder
         super().__init__(cfg=cfg, ch=ch, nc=nc, verbose=verbose)
 
     def set_classes(self, text, batch=80, cache_clip_model=True):
-        """Set classes in advance so that model could do offline-inference without clip model."""
+        """
+        Set classes in advance so that model could do offline-inference without clip model.
+
+        Args:
+            text (List[str]): List of class names.
+            batch (int): Batch size for processing text tokens.
+            cache_clip_model (bool): Whether to cache the CLIP model.
+        """
         try:
             import clip
         except ImportError:
@@ -628,11 +725,11 @@ class WorldModel(DetectionModel):
 
         Args:
             x (torch.Tensor): The input tensor.
-            profile (bool, optional): If True, profile the computation time for each layer. Defaults to False.
-            visualize (bool, optional): If True, save feature maps for visualization. Defaults to False.
-            txt_feats (torch.Tensor): The text features, use it if it's given. Defaults to None.
-            augment (bool, optional): If True, perform data augmentation during inference. Defaults to False.
-            embed (list, optional): A list of feature vectors/embeddings to return.
+            profile (bool): If True, profile the computation time for each layer.
+            visualize (bool): If True, save feature maps for visualization.
+            txt_feats (torch.Tensor, optional): The text features, use it if it's given.
+            augment (bool): If True, perform data augmentation during inference.
+            embed (List, optional): A list of feature vectors/embeddings to return.
 
         Returns:
             (torch.Tensor): Model's output tensor.
@@ -671,7 +768,7 @@ class WorldModel(DetectionModel):
 
         Args:
             batch (dict): Batch to compute loss on.
-            preds (torch.Tensor | List[torch.Tensor]): Predictions.
+            preds (torch.Tensor | List[torch.Tensor], optional): Predictions.
         """
         if not hasattr(self, "criterion"):
             self.criterion = self.init_criterion()
@@ -689,7 +786,18 @@ class Ensemble(torch.nn.ModuleList):
         super().__init__()
 
     def forward(self, x, augment=False, profile=False, visualize=False):
-        """Function generates the YOLO network's final layer."""
+        """
+        Generate the YOLO network's final layer.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+            augment (bool): Whether to augment the input.
+            profile (bool): Whether to profile the model.
+            visualize (bool): Whether to visualize the features.
+
+        Returns:
+            (tuple): Tuple containing the concatenated predictions and None.
+        """
         y = [module(x, augment, profile, visualize)[0] for module in self]
         # y = torch.stack(y).max(0)[0]  # max ensemble
         # y = torch.stack(y).mean(0)  # mean ensemble
@@ -765,7 +873,16 @@ class SafeUnpickler(pickle.Unpickler):
     """Custom Unpickler that replaces unknown classes with SafeClass."""
 
     def find_class(self, module, name):
-        """Attempt to find a class, returning SafeClass if not among safe modules."""
+        """
+        Attempt to find a class, returning SafeClass if not among safe modules.
+
+        Args:
+            module (str): Module name.
+            name (str): Class name.
+
+        Returns:
+            (type): Found class or SafeClass.
+        """
         safe_modules = (
             "torch",
             "collections",
@@ -791,13 +908,13 @@ def torch_safe_load(weight, safe_only=False):
         weight (str): The file path of the PyTorch model.
         safe_only (bool): If True, replace unknown classes with SafeClass during loading.
 
+    Returns:
+        ckpt (dict): The loaded model checkpoint.
+        file (str): The loaded filename.
+
     Examples:
         >>> from ultralytics.nn.tasks import torch_safe_load
         >>> ckpt, file = torch_safe_load("path/to/best.pt", safe_only=True)
-
-    Returns:
-        ckpt (dict): The loaded model checkpoint.
-        file (str): The loaded filename
     """
     from ultralytics.utils.downloads import attempt_download_asset
 
@@ -858,7 +975,18 @@ def torch_safe_load(weight, safe_only=False):
 
 
 def attempt_load_weights(weights, device=None, inplace=True, fuse=False):
-    """Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a."""
+    """
+    Load an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a.
+
+    Args:
+        weights (str | List[str]): Model weights path(s).
+        device (torch.device, optional): Device to load model to.
+        inplace (bool): Whether to do inplace operations.
+        fuse (bool): Whether to fuse model.
+
+    Returns:
+        (torch.nn.Module): Loaded model.
+    """
     ensemble = Ensemble()
     for w in weights if isinstance(weights, list) else [weights]:
         ckpt, w = torch_safe_load(w)  # load ckpt
@@ -896,7 +1024,18 @@ def attempt_load_weights(weights, device=None, inplace=True, fuse=False):
 
 
 def attempt_load_one_weight(weight, device=None, inplace=True, fuse=False):
-    """Loads a single model weights."""
+    """
+    Load a single model weights.
+
+    Args:
+        weight (str): Model weight path.
+        device (torch.device, optional): Device to load model to.
+        inplace (bool): Whether to do inplace operations.
+        fuse (bool): Whether to fuse model.
+
+    Returns:
+        (tuple): Tuple containing the model and checkpoint.
+    """
     ckpt, weight = torch_safe_load(weight)  # load ckpt
     args = {**DEFAULT_CFG_DICT, **(ckpt.get("train_args", {}))}  # combine model and default args, preferring model args
     model = (ckpt.get("ema") or ckpt["model"]).to(device).float()  # FP32 model
@@ -922,7 +1061,17 @@ def attempt_load_one_weight(weight, device=None, inplace=True, fuse=False):
 
 
 def parse_model(d, ch, verbose=True):  # model_dict, input_channels(3)
-    """Parse a YOLO model.yaml dictionary into a PyTorch model."""
+    """
+    Parse a YOLO model.yaml dictionary into a PyTorch model.
+
+    Args:
+        d (dict): Model dictionary.
+        ch (int): Input channels.
+        verbose (bool): Whether to print model details.
+
+    Returns:
+        (tuple): Tuple containing the PyTorch model and sorted list of output layers.
+    """
     import ast
 
     # Args
@@ -1086,7 +1235,15 @@ def parse_model(d, ch, verbose=True):  # model_dict, input_channels(3)
 
 
 def yaml_model_load(path):
-    """Load a YOLOv8 model from a YAML file."""
+    """
+    Load a YOLOv8 model from a YAML file.
+
+    Args:
+        path (str | Path): Path to the YAML file.
+
+    Returns:
+        (dict): Model dictionary.
+    """
     path = Path(path)
     if path.stem in (f"yolov{d}{x}6" for x in "nsmlx" for d in (5, 8)):
         new_stem = re.sub(r"(\d+)([nslmx])6(.+)?$", r"\1\2-p6\3", path.stem)
@@ -1103,15 +1260,13 @@ def yaml_model_load(path):
 
 def guess_model_scale(model_path):
     """
-    Takes a path to a YOLO model's YAML file as input and extracts the size character of the model's scale. The function
-    uses regular expression matching to find the pattern of the model scale in the YAML file name, which is denoted by
-    n, s, m, l, or x. The function returns the size character of the model scale as a string.
+    Extract the size character n, s, m, l, or x of the model's scale from the model path.
 
     Args:
         model_path (str | Path): The path to the YOLO model's YAML file.
 
     Returns:
-        (str): The size character of the model's scale, which can be n, s, m, l, or x.
+        (str): The size character of the model's scale (n, s, m, l, or x).
     """
     try:
         return re.search(r"yolo[v]?\d+([nslmx])", Path(model_path).stem).group(1)  # returns n, s, m, l, or x
@@ -1127,10 +1282,7 @@ def guess_model_task(model):
         model (torch.nn.Module | dict): PyTorch model or model configuration in YAML format.
 
     Returns:
-        (str): Task of the model ('detect', 'segment', 'classify', 'pose').
-
-    Raises:
-        SyntaxError: If the task of the model could not be determined.
+        (str): Task of the model ('detect', 'segment', 'classify', 'pose', 'obb').
     """
 
     def cfg2task(cfg):

ultralytics/solutions/ai_gym.py

@@ -33,7 +33,7 @@ class AIGym(BaseSolution):
 
     def __init__(self, **kwargs):
         """
-        Initializes AIGym for workout monitoring using pose estimation and predefined angles.
+        Initialize AIGym for workout monitoring using pose estimation and predefined angles.
 
         Args:
             **kwargs (Any): Keyword arguments passed to the parent class constructor.
@@ -53,7 +53,7 @@ class AIGym(BaseSolution):
 
     def process(self, im0):
         """
-        Monitors workouts using Ultralytics YOLO Pose Model.
+        Monitor workouts using Ultralytics YOLO Pose Model.
 
         This function processes an input image to track and analyze human poses for workout monitoring. It uses
         the YOLO Pose model to detect keypoints, estimate angles, and count repetitions based on predefined

ultralytics/solutions/analytics.py

@@ -37,8 +37,8 @@ class Analytics(BaseSolution):
         color_mapping (Dict[str, str]): Dictionary mapping class labels to colors for consistent visualization.
 
     Methods:
-        process: Processes image data and updates the chart.
-        update_graph: Updates the chart with new data points.
+        process: Process image data and update the chart.
+        update_graph: Update the chart with new data points.
 
     Examples:
         >>> analytics = Analytics(analytics_type="line")
@@ -87,7 +87,7 @@ class Analytics(BaseSolution):
 
     def process(self, im0, frame_number):
         """
-        Processes image data and runs object tracking to update analytics charts.
+        Process image data and run object tracking to update analytics charts.
 
         Args:
             im0 (np.ndarray): Input image for processing.
@@ -127,7 +127,7 @@ class Analytics(BaseSolution):
 
     def update_graph(self, frame_number, count_dict=None, plot="line"):
         """
-        Updates the graph with new data for single or multiple classes.
+        Update the graph with new data for single or multiple classes.
 
         Args:
             frame_number (int): The current frame number.

ultralytics/solutions/heatmap.py

@@ -21,8 +21,8 @@ class Heatmap(ObjectCounter):
         annotator (SolutionAnnotator): Object for drawing annotations on the image.
 
     Methods:
-        heatmap_effect: Calculates and updates the heatmap effect for a given bounding box.
-        process: Generates and applies the heatmap effect to each frame.
+        heatmap_effect: Calculate and update the heatmap effect for a given bounding box.
+        process: Generate and apply the heatmap effect to each frame.
 
     Examples:
         >>> from ultralytics.solutions import Heatmap
@@ -33,7 +33,7 @@ class Heatmap(ObjectCounter):
 
     def __init__(self, **kwargs):
         """
-        Initializes the Heatmap class for real-time video stream heatmap generation based on object tracks.
+        Initialize the Heatmap class for real-time video stream heatmap generation based on object tracks.
 
         Args:
             **kwargs (Any): Keyword arguments passed to the parent ObjectCounter class.
@@ -50,7 +50,7 @@ class Heatmap(ObjectCounter):
 
     def heatmap_effect(self, box):
         """
-        Efficiently calculates heatmap area and effect location for applying colormap.
+        Efficiently calculate heatmap area and effect location for applying colormap.
 
         Args:
             box (List[float]): Bounding box coordinates [x0, y0, x1, y1].

ultralytics/solutions/instance_segmentation.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
 from ultralytics.utils.plotting import colors
@@ -13,9 +13,15 @@ class InstanceSegmentation(BaseSolution):
 
     Attributes:
         model (str): The segmentation model to use for inference.
+        line_width (int): Width of the bounding box and text lines.
+        names (Dict[int, str]): Dictionary mapping class indices to class names.
+        clss (List[int]): List of detected class indices.
+        track_ids (List[int]): List of track IDs for detected instances.
+        masks (List[numpy.ndarray]): List of segmentation masks for detected instances.
 
     Methods:
-        process: Processes the input image to perform instance segmentation and annotate results.
+        process: Process the input image to perform instance segmentation and annotate results.
+        extract_tracks: Extract tracks including bounding boxes, classes, and masks from model predictions.
 
     Examples:
         >>> segmenter = InstanceSegmentation()
@@ -26,7 +32,7 @@ class InstanceSegmentation(BaseSolution):
 
     def __init__(self, **kwargs):
         """
-        Initializes the InstanceSegmentation class for detecting and annotating segmented instances.
+        Initialize the InstanceSegmentation class for detecting and annotating segmented instances.
 
         Args:
             **kwargs (Any): Keyword arguments passed to the BaseSolution parent class.
@@ -37,7 +43,7 @@ class InstanceSegmentation(BaseSolution):
 
     def process(self, im0):
         """
-        Performs instance segmentation on the input image and annotates the results.
+        Perform instance segmentation on the input image and annotate the results.
 
         Args:
             im0 (numpy.ndarray): The input image for segmentation.