ultralytics 8.3.143__py3-none-any.whl → 8.3.145__py3-none-any.whl

ultralytics/utils/patches.py

@@ -3,6 +3,7 @@
 
 import time
 from pathlib import Path
+from typing import List, Optional
 
 import cv2
 import numpy as np
@@ -12,16 +13,16 @@ import torch
 _imshow = cv2.imshow  # copy to avoid recursion errors
 
 
-def imread(filename: str, flags: int = cv2.IMREAD_COLOR):
+def imread(filename: str, flags: int = cv2.IMREAD_COLOR) -> Optional[np.ndarray]:
     """
-    Read an image from a file.
+    Read an image from a file with multilanguage filename support.
 
     Args:
         filename (str): Path to the file to read.
-        flags (int): Flag that can take values of cv2.IMREAD_*. Controls how the image is read.
+        flags (int, optional): Flag that can take values of cv2.IMREAD_*. Controls how the image is read.
 
     Returns:
-        (np.ndarray): The read image.
+        (np.ndarray | None): The read image array, or None if reading fails.
 
     Examples:
         >>> img = imread("path/to/image.jpg")
@@ -31,17 +32,17 @@ def imread(filename: str, flags: int = cv2.IMREAD_COLOR):
     if filename.endswith((".tiff", ".tif")):
         success, frames = cv2.imdecodemulti(file_bytes, cv2.IMREAD_UNCHANGED)
         if success:
-            # handle RGB images in tif/tiff format
+            # Handle RGB images in tif/tiff format
             return frames[0] if len(frames) == 1 and frames[0].ndim == 3 else np.stack(frames, axis=2)
         return None
     else:
         im = cv2.imdecode(file_bytes, flags)
-        return im[..., None] if im.ndim == 2 else im  # always make sure there's 3 dimensions
+        return im[..., None] if im.ndim == 2 else im  # Always ensure 3 dimensions
 
 
-def imwrite(filename: str, img: np.ndarray, params=None):
+def imwrite(filename: str, img: np.ndarray, params: Optional[List[int]] = None) -> bool:
     """
-    Write an image to a file.
+    Write an image to a file with multilanguage filename support.
 
     Args:
         filename (str): Path to the file to write.
@@ -65,12 +66,12 @@ def imwrite(filename: str, img: np.ndarray, params=None):
         return False
 
 
-def imshow(winname: str, mat: np.ndarray):
+def imshow(winname: str, mat: np.ndarray) -> None:
     """
-    Display an image in the specified window.
+    Display an image in the specified window with multilanguage window name support.
 
-    This function is a wrapper around OpenCV's imshow function that displays an image in a named window. It is
-    particularly useful for visualizing images during development and debugging.
+    This function is a wrapper around OpenCV's imshow function that displays an image in a named window. It handles
+    multilanguage window names by encoding them properly for OpenCV compatibility.
 
     Args:
         winname (str): Name of the window where the image will be displayed. If a window with this name already
@@ -127,9 +128,6 @@ def torch_save(*args, **kwargs):
         *args (Any): Positional arguments to pass to torch.save.
         **kwargs (Any): Keyword arguments to pass to torch.save.
 
-    Returns:
-        (Any): Result of torch.save operation if successful, None otherwise.
-
     Examples:
         >>> model = torch.nn.Linear(10, 1)
         >>> torch_save(model.state_dict(), "model.pt")
@@ -137,7 +135,7 @@ def torch_save(*args, **kwargs):
     for i in range(4):  # 3 retries
         try:
             return _torch_save(*args, **kwargs)
-        except RuntimeError as e:  # unable to save, possibly waiting for device to flush or antivirus scan
+        except RuntimeError as e:  # Unable to save, possibly waiting for device to flush or antivirus scan
             if i == 3:
                 raise e
-            time.sleep((2**i) / 2)  # exponential standoff: 0.5s, 1.0s, 2.0s
+            time.sleep((2**i) / 2)  # Exponential backoff: 0.5s, 1.0s, 2.0s

ultralytics/utils/plotting.py

@@ -18,20 +18,21 @@ from ultralytics.utils.files import increment_path
 
 class Colors:
     """
-    Ultralytics color palette https://docs.ultralytics.com/reference/utils/plotting/#ultralytics.utils.plotting.Colors.
+    Ultralytics color palette for visualization and plotting.
 
     This class provides methods to work with the Ultralytics color palette, including converting hex color codes to
-    RGB values.
+    RGB values and accessing predefined color schemes for object detection and pose estimation.
 
     Attributes:
-        palette (List[Tuple]): List of RGB color values.
+        palette (List[tuple]): List of RGB color tuples for general use.
         n (int): The number of colors in the palette.
         pose_palette (np.ndarray): A specific color palette array for pose estimation with dtype np.uint8.
 
     Examples:
         >>> from ultralytics.utils.plotting import Colors
         >>> colors = Colors()
-        >>> colors(5, True)  # ff6fdd or (255, 111, 221)
+        >>> colors(5, True)  # Returns BGR format: (221, 111, 255)
+        >>> colors(5, False)  # Returns RGB format: (255, 111, 221)
 
     ## Ultralytics Color Palette
 
@@ -85,7 +86,8 @@ class Colors:
 
     !!! note "Ultralytics Brand Colors"
 
-        For Ultralytics brand colors see [https://www.ultralytics.com/brand](https://www.ultralytics.com/brand). Please use the official Ultralytics colors for all marketing materials.
+        For Ultralytics brand colors see [https://www.ultralytics.com/brand](https://www.ultralytics.com/brand).
+        Please use the official Ultralytics colors for all marketing materials.
     """
 
     def __init__(self):
@@ -140,13 +142,22 @@ class Colors:
             dtype=np.uint8,
         )
 
-    def __call__(self, i, bgr=False):
-        """Convert hex color codes to RGB values."""
+    def __call__(self, i: int, bgr: bool = False) -> tuple:
+        """
+        Convert hex color codes to RGB values.
+
+        Args:
+            i (int): Color index.
+            bgr (bool, optional): Whether to return BGR format instead of RGB.
+
+        Returns:
+            (tuple): RGB or BGR color tuple.
+        """
         c = self.palette[int(i) % self.n]
         return (c[2], c[1], c[0]) if bgr else c
 
     @staticmethod
-    def hex2rgb(h):
+    def hex2rgb(h: str) -> tuple:
         """Convert hex color codes to RGB values (i.e. default PIL order)."""
         return tuple(int(h[1 + i : 1 + i + 2], 16) for i in (0, 2, 4))
 
@@ -159,9 +170,9 @@ class Annotator:
     Ultralytics Annotator for train/val mosaics and JPGs and predictions annotations.
 
     Attributes:
-        im (Image.Image or np.ndarray): The image to annotate.
+        im (Image.Image | np.ndarray): The image to annotate.
         pil (bool): Whether to use PIL or cv2 for drawing annotations.
-        font (ImageFont.truetype or ImageFont.load_default): Font used for text annotations.
+        font (ImageFont.truetype | ImageFont.load_default): Font used for text annotations.
         lw (float): Line width for drawing.
         skeleton (List[List[int]]): Skeleton structure for keypoints.
         limb_color (List[int]): Color palette for limbs.
@@ -173,9 +184,18 @@ class Annotator:
         >>> from ultralytics.utils.plotting import Annotator
         >>> im0 = cv2.imread("test.png")
         >>> annotator = Annotator(im0, line_width=10)
+        >>> annotator.box_label([10, 10, 100, 100], "person", (255, 0, 0))
     """
 
-    def __init__(self, im, line_width=None, font_size=None, font="Arial.ttf", pil=False, example="abc"):
+    def __init__(
+        self,
+        im,
+        line_width: Optional[int] = None,
+        font_size: Optional[int] = None,
+        font: str = "Arial.ttf",
+        pil: bool = False,
+        example: str = "abc",
+    ):
         """Initialize the Annotator class with image and line width along with color palette for keypoints and limbs."""
         non_ascii = not is_ascii(example)  # non-latin labels, i.e. asian, arabic, cyrillic
         input_is_pil = isinstance(im, Image.Image)
@@ -254,7 +274,7 @@ class Annotator:
             (104, 31, 17),
         }
 
-    def get_txt_color(self, color=(128, 128, 128), txt_color=(255, 255, 255)):
+    def get_txt_color(self, color: tuple = (128, 128, 128), txt_color: tuple = (255, 255, 255)) -> tuple:
         """
         Assign text color based on background color.
 
@@ -278,7 +298,7 @@ class Annotator:
         else:
             return txt_color
 
-    def box_label(self, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255)):
+    def box_label(self, box, label: str = "", color: tuple = (128, 128, 128), txt_color: tuple = (255, 255, 255)):
         """
         Draw a bounding box on an image with a given label.
 
@@ -340,7 +360,7 @@ class Annotator:
                     lineType=cv2.LINE_AA,
                 )
 
-    def masks(self, masks, colors, im_gpu, alpha=0.5, retina_masks=False):
+    def masks(self, masks, colors, im_gpu, alpha: float = 0.5, retina_masks: bool = False):
         """
         Plot masks on image.
 
@@ -376,7 +396,15 @@ class Annotator:
             # Convert im back to PIL and update draw
             self.fromarray(self.im)
 
-    def kpts(self, kpts, shape=(640, 640), radius=None, kpt_line=True, conf_thres=0.25, kpt_color=None):
+    def kpts(
+        self,
+        kpts,
+        shape: tuple = (640, 640),
+        radius: Optional[int] = None,
+        kpt_line: bool = True,
+        conf_thres: float = 0.25,
+        kpt_color: Optional[tuple] = None,
+    ):
         """
         Plot keypoints on the image.
 
@@ -436,11 +464,11 @@ class Annotator:
             # Convert im back to PIL and update draw
             self.fromarray(self.im)
 
-    def rectangle(self, xy, fill=None, outline=None, width=1):
+    def rectangle(self, xy, fill=None, outline=None, width: int = 1):
         """Add rectangle to image (PIL-only)."""
         self.draw.rectangle(xy, fill, outline, width)
 
-    def text(self, xy, text, txt_color=(255, 255, 255), anchor="top", box_color=()):
+    def text(self, xy, text: str, txt_color: tuple = (255, 255, 255), anchor: str = "top", box_color: tuple = ()):
         """
         Add text to an image using PIL or cv2.
 
@@ -480,7 +508,7 @@ class Annotator:
         """Return annotated image as array."""
         return np.asarray(self.im)
 
-    def show(self, title=None):
+    def show(self, title: Optional[str] = None):
         """Show the annotated image."""
         im = Image.fromarray(np.asarray(self.im)[..., ::-1])  # Convert numpy array to PIL Image with RGB to BGR
         if IS_COLAB or IS_KAGGLE:  # can not use IS_JUPYTER as will run for all ipython environments
@@ -491,12 +519,12 @@ class Annotator:
         else:
             im.show(title=title)
 
-    def save(self, filename="image.jpg"):
+    def save(self, filename: str = "image.jpg"):
         """Save the annotated image to 'filename'."""
         cv2.imwrite(filename, np.asarray(self.im))
 
     @staticmethod
-    def get_bbox_dimension(bbox=None):
+    def get_bbox_dimension(bbox: Optional[tuple] = None):
         """
         Calculate the dimensions and area of a bounding box.
 
@@ -592,7 +620,16 @@ def plot_labels(boxes, cls, names=(), save_dir=Path(""), on_plot=None):
         on_plot(fname)
 
 
-def save_one_box(xyxy, im, file=Path("im.jpg"), gain=1.02, pad=10, square=False, BGR=False, save=True):
+def save_one_box(
+    xyxy,
+    im,
+    file: Path = Path("im.jpg"),
+    gain: float = 1.02,
+    pad: int = 10,
+    square: bool = False,
+    BGR: bool = False,
+    save: bool = True,
+):
     """
     Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop.
 
@@ -808,7 +845,14 @@ def plot_images(
 
 
 @plt_settings()
-def plot_results(file="path/to/results.csv", dir="", segment=False, pose=False, classify=False, on_plot=None):
+def plot_results(
+    file: str = "path/to/results.csv",
+    dir: str = "",
+    segment: bool = False,
+    pose: bool = False,
+    classify: bool = False,
+    on_plot: Optional[Callable] = None,
+):
     """
     Plot training results from a results CSV file. The function supports various types of data including segmentation,
     pose estimation, and classification. Plots are saved as 'results.png' in the directory where the CSV is located.
@@ -868,7 +912,7 @@ def plot_results(file="path/to/results.csv", dir="", segment=False, pose=False,
         on_plot(fname)
 
 
-def plt_color_scatter(v, f, bins=20, cmap="viridis", alpha=0.8, edgecolors="none"):
+def plt_color_scatter(v, f, bins: int = 20, cmap: str = "viridis", alpha: float = 0.8, edgecolors: str = "none"):
     """
     Plot a scatter plot with points colored based on a 2D histogram.
 
@@ -901,7 +945,7 @@ def plt_color_scatter(v, f, bins=20, cmap="viridis", alpha=0.8, edgecolors="none
     plt.scatter(v, f, c=colors, cmap=cmap, alpha=alpha, edgecolors=edgecolors)
 
 
-def plot_tune_results(csv_file="tune_results.csv"):
+def plot_tune_results(csv_file: str = "tune_results.csv"):
     """
     Plot the evolution results stored in a 'tune_results.csv' file. The function generates a scatter plot for each key
     in the CSV, color-coded based on fitness scores. The best-performing configurations are highlighted on the plots.
@@ -957,7 +1001,7 @@ def plot_tune_results(csv_file="tune_results.csv"):
     _save_one_file(csv_file.with_name("tune_fitness.png"))
 
 
-def output_to_target(output, max_det=300):
+def output_to_target(output, max_det: int = 300):
     """Convert model output to target format [batch_id, class_id, x, y, w, h, conf] for plotting."""
     targets = []
     for i, o in enumerate(output):
@@ -968,7 +1012,7 @@ def output_to_target(output, max_det=300):
     return targets[:, 0], targets[:, 1], targets[:, 2:-1], targets[:, -1]
 
 
-def output_to_rotated_target(output, max_det=300):
+def output_to_rotated_target(output, max_det: int = 300):
     """Convert model output to target format [batch_id, class_id, x, y, w, h, conf] for plotting."""
     targets = []
     for i, o in enumerate(output):
@@ -979,7 +1023,7 @@ def output_to_rotated_target(output, max_det=300):
     return targets[:, 0], targets[:, 1], targets[:, 2:-1], targets[:, -1]
 
 
-def feature_visualization(x, module_type, stage, n=32, save_dir=Path("runs/detect/exp")):
+def feature_visualization(x, module_type: str, stage: int, n: int = 32, save_dir: Path = Path("runs/detect/exp")):
     """
     Visualize feature maps of a given model module during inference.
 

ultralytics/utils/tal.py

@@ -26,8 +26,17 @@ class TaskAlignedAssigner(nn.Module):
         eps (float): A small value to prevent division by zero.
     """
 
-    def __init__(self, topk=13, num_classes=80, alpha=1.0, beta=6.0, eps=1e-9):
-        """Initialize a TaskAlignedAssigner object with customizable hyperparameters."""
+    def __init__(self, topk: int = 13, num_classes: int = 80, alpha: float = 1.0, beta: float = 6.0, eps: float = 1e-9):
+        """
+        Initialize a TaskAlignedAssigner object with customizable hyperparameters.
+
+        Args:
+            topk (int, optional): The number of top candidates to consider.
+            num_classes (int, optional): The number of object classes.
+            alpha (float, optional): The alpha parameter for the classification component of the task-aligned metric.
+            beta (float, optional): The beta parameter for the localization component of the task-aligned metric.
+            eps (float, optional): A small value to prevent division by zero.
+        """
         super().__init__()
         self.topk = topk
         self.num_classes = num_classes
@@ -196,12 +205,11 @@ class TaskAlignedAssigner(nn.Module):
         Select the top-k candidates based on the given metrics.
 
         Args:
-            metrics (torch.Tensor): A tensor of shape (b, max_num_obj, h*w), where b is the batch size,
-                              max_num_obj is the maximum number of objects, and h*w represents the
-                              total number of anchor points.
-            topk_mask (torch.Tensor): An optional boolean tensor of shape (b, max_num_obj, topk), where
-                                topk is the number of top candidates to consider. If not provided,
-                                the top-k values are automatically computed based on the given metrics.
+            metrics (torch.Tensor): A tensor of shape (b, max_num_obj, h*w), where b is the batch size, max_num_obj is
+                the maximum number of objects, and h*w represents the total number of anchor points.
+            topk_mask (torch.Tensor, optional): An optional boolean tensor of shape (b, max_num_obj, topk), where
+                topk is the number of top candidates to consider. If not provided, the top-k values are automatically
+                computed based on the given metrics.
 
         Returns:
             (torch.Tensor): A tensor of shape (b, max_num_obj, h*w) containing the selected top-k candidates.
@@ -239,11 +247,9 @@ class TaskAlignedAssigner(nn.Module):
                               (foreground) anchor points.
 
         Returns:
-            target_labels (torch.Tensor): Shape (b, h*w), containing the target labels for positive anchor points.
-            target_bboxes (torch.Tensor): Shape (b, h*w, 4), containing the target bounding boxes for positive
-                                          anchor points.
-            target_scores (torch.Tensor): Shape (b, h*w, num_classes), containing the target scores for positive
-                                          anchor points.
+            target_labels (torch.Tensor): Target labels for positive anchor points with shape (b, h*w).
+            target_bboxes (torch.Tensor): Target bounding boxes for positive anchor points with shape (b, h*w, 4).
+            target_scores (torch.Tensor): Target scores for positive anchor points with shape (b, h*w, num_classes).
         """
         # Assigned target labels, (b, 1)
         batch_ind = torch.arange(end=self.bs, dtype=torch.int64, device=gt_labels.device)[..., None]
@@ -277,7 +283,7 @@ class TaskAlignedAssigner(nn.Module):
         Args:
             xy_centers (torch.Tensor): Anchor center coordinates, shape (h*w, 2).
             gt_bboxes (torch.Tensor): Ground truth bounding boxes, shape (b, n_boxes, 4).
-            eps (float, optional): Small value for numerical stability. Defaults to 1e-9.
+            eps (float, optional): Small value for numerical stability.
 
         Returns:
             (torch.Tensor): Boolean mask of positive anchors, shape (b, n_boxes, h*w).
@@ -399,7 +405,7 @@ def dist2rbox(pred_dist, pred_angle, anchor_points, dim=-1):
         pred_dist (torch.Tensor): Predicted rotated distance with shape (bs, h*w, 4).
         pred_angle (torch.Tensor): Predicted angle with shape (bs, h*w, 1).
         anchor_points (torch.Tensor): Anchor points with shape (h*w, 2).
-        dim (int, optional): Dimension along which to split. Defaults to -1.
+        dim (int, optional): Dimension along which to split.
 
     Returns:
         (torch.Tensor): Predicted rotated bounding boxes with shape (bs, h*w, 4).