Spaces:
Running
Running
| # 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, optional): Flag that can take values of cv2.IMREAD_*. Defaults to cv2.IMREAD_COLOR. | |
| Returns: | |
| (np.ndarray): The read image. | |
| """ | |
| return cv2.imdecode(np.fromfile(filename, np.uint8), flags) | |
| 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 of ints, optional): Additional parameters. See OpenCV documentation. | |
| Returns: | |
| (bool): True if the file was written, False otherwise. | |
| """ | |
| try: | |
| cv2.imencode(Path(filename).suffix, img, params)[1].tofile(filename) | |
| return True | |
| except Exception: | |
| return False | |
| def imshow(winname: str, mat: np.ndarray): | |
| """ | |
| Displays an image in the specified window. | |
| Args: | |
| winname (str): Name of the window. | |
| mat (np.ndarray): Image to be shown. | |
| """ | |
| _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. | |
| Note: | |
| 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): | |
| """ | |
| Optionally use dill to serialize lambda functions where pickle does not, adding robustness with 3 retries and | |
| exponential standoff in case of save failure. | |
| Args: | |
| *args (tuple): Positional arguments to pass to torch.save. | |
| **kwargs (Any): Keyword arguments to pass to torch.save. | |
| """ | |
| 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 | |