dgenerate-ultralytics-headless 8.3.134__py3-none-any.whl

ultralytics/utils/patches.py

@@ -0,0 +1,143 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""Monkey patches to update/extend functionality of existing functions."""
+
+import time
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+
+# OpenCV Multilanguage-friendly functions ------------------------------------------------------------------------------
+_imshow = cv2.imshow  # copy to avoid recursion errors
+
+
+def imread(filename: str, flags: int = cv2.IMREAD_COLOR):
+    """
+    Read an image from a file.
+
+    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.
+
+    Returns:
+        (np.ndarray): The read image.
+
+    Examples:
+        >>> img = imread("path/to/image.jpg")
+        >>> img = imread("path/to/image.jpg", cv2.IMREAD_GRAYSCALE)
+    """
+    file_bytes = np.fromfile(filename, np.uint8)
+    if filename.endswith((".tiff", ".tif")):
+        success, frames = cv2.imdecodemulti(file_bytes, cv2.IMREAD_UNCHANGED)
+        if success:
+            # 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
+
+
+def imwrite(filename: str, img: np.ndarray, params=None):
+    """
+    Write an image to a file.
+
+    Args:
+        filename (str): Path to the file to write.
+        img (np.ndarray): Image to write.
+        params (List[int], optional): Additional parameters for image encoding.
+
+    Returns:
+        (bool): True if the file was written successfully, False otherwise.
+
+    Examples:
+        >>> import numpy as np
+        >>> img = np.zeros((100, 100, 3), dtype=np.uint8)  # Create a black image
+        >>> success = imwrite("output.jpg", img)  # Write image to file
+        >>> print(success)
+        True
+    """
+    try:
+        cv2.imencode(Path(filename).suffix, img, params)[1].tofile(filename)
+        return True
+    except Exception:
+        return False
+
+
+def imshow(winname: str, mat: np.ndarray):
+    """
+    Display an image in the specified window.
+
+    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.
+
+    Args:
+        winname (str): Name of the window where the image will be displayed. If a window with this name already
+            exists, the image will be displayed in that window.
+        mat (np.ndarray): Image to be shown. Should be a valid numpy array representing an image.
+
+    Examples:
+        >>> import numpy as np
+        >>> img = np.zeros((300, 300, 3), dtype=np.uint8)  # Create a black image
+        >>> img[:100, :100] = [255, 0, 0]  # Add a blue square
+        >>> imshow("Example Window", img)  # Display the image
+    """
+    _imshow(winname.encode("unicode_escape").decode(), mat)
+
+
+# PyTorch functions ----------------------------------------------------------------------------------------------------
+_torch_load = torch.load  # copy to avoid recursion errors
+_torch_save = torch.save
+
+
+def torch_load(*args, **kwargs):
+    """
+    Load a PyTorch model with updated arguments to avoid warnings.
+
+    This function wraps torch.load and adds the 'weights_only' argument for PyTorch 1.13.0+ to prevent warnings.
+
+    Args:
+        *args (Any): Variable length argument list to pass to torch.load.
+        **kwargs (Any): Arbitrary keyword arguments to pass to torch.load.
+
+    Returns:
+        (Any): The loaded PyTorch object.
+
+    Notes:
+        For PyTorch versions 2.0 and above, this function automatically sets 'weights_only=False'
+        if the argument is not provided, to avoid deprecation warnings.
+    """
+    from ultralytics.utils.torch_utils import TORCH_1_13
+
+    if TORCH_1_13 and "weights_only" not in kwargs:
+        kwargs["weights_only"] = False
+
+    return _torch_load(*args, **kwargs)
+
+
+def torch_save(*args, **kwargs):
+    """
+    Save PyTorch objects with retry mechanism for robustness.
+
+    This function wraps torch.save with 3 retries and exponential backoff in case of save failures, which can occur
+    due to device flushing delays or antivirus scanning.
+
+    Args:
+        *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")
+    """
+    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
+            if i == 3:
+                raise e
+            time.sleep((2**i) / 2)  # exponential standoff: 0.5s, 1.0s, 2.0s